【人工智能笔记】第四十节：TF2实现VITGAN对抗生成网络，Siren网络 实现

该章节介绍VITGAN对抗生成网络中，Siren网络 部分的代码实现。

目录（文章发布后会补上链接）：

1. 网络结构简介
2. Mapping NetWork 实现
3. PositionalEmbedding 实现
4. MLP 实现
5. MSA多头注意力 实现
6. SLN自调制 实现
7. CoordinatesPositionalEmbedding 实现
8. ModulatedLinear 实现
9. Siren 实现
10. Generator生成器 实现
11. PatchEmbedding 实现
12. ISN 实现
13. Discriminator鉴别器 实现
14. VITGAN 实现

Siren网络 简介

Siren网络 是论文中资料46里提到的网络结构。

代码实现

原Siren网络实现：

import tensorflow as tf
import numpy as np


class SineLayer(tf.keras.layers.Layer):
    def __init__(
        self,
        emb_dim,
        use_bias=True,
        is_first=False,
        omega_0=30,
        demodulation=True,
        outermost_linear=False, 
    ):
        super().__init__()
        self.emb_dim = emb_dim
        self.use_bias = use_bias
        self.is_first = is_first
        self.omega_0 = omega_0
        self.demodulation = demodulation
        self.outermost_linear=outermost_linear
        kernel_initializer = None
        if is_first:
            kernel_initializer = tf.random_uniform_initializer(
                -1/emb_dim,
                1/emb_dim
            )
        else:
            kernel_initializer = tf.random_uniform_initializer(
                -np.sqrt(6 / emb_dim) / omega_0,
                np.sqrt(6 / emb_dim) / omega_0
            )
        self.linear = tf.keras.layers.Dense(
            units=emb_dim,
            use_bias=use_bias,
            kernel_initializer=kernel_initializer,
        )

    def call(self, x):
        x = self.linear(x)
        if self.outermost_linear:
            return x
        else:
            return tf.math.sin(self.omega_0 * x)


class Siren(tf.keras.layers.Layer):
    def __init__(
        self,
        hidden_dim,
        hidden_layers,
        out_dim,
        first_omega_0=30.0,
        hidden_omega_0=30.0,
        demodulation=True,
        outermost_linear=False,
        ):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.hidden_layers = hidden_layers
        self.out_dim = out_dim
        self.first_omega_0=first_omega_0
        self.hidden_omega_0=hidden_omega_0
        self.demodulation=demodulation
        self.outermost_linear=outermost_linear
        self.net_layers = []
        self.net_layers.append(
            SineLayer(
                hidden_dim, 
                is_first=True,
                omega_0=first_omega_0,
                demodulation=demodulation,
            )
        )

        for _ in range(hidden_layers):
            self.net_layers.append(
                SineLayer(
                    hidden_dim, 
                    is_first=False,
                    omega_0=hidden_omega_0,
                    demodulation=demodulation,
                )
            )

        self.net_layers.append(
            SineLayer(
                out_dim,
                is_first=False,
                omega_0=hidden_omega_0,
                demodulation=demodulation,
                outermost_linear=outermost_linear,
            )
        )

    def call(self, x):
        for layer in self.net_layers:
            x = layer(x)
        return x



if __name__ == "__main__":
    emb_dim = 768 # 16*16*3
    layer = Siren(
        hidden_dim=emb_dim,
        hidden_layers=3,
        out_dim=3,
        first_omega_0=30.0,
        hidden_omega_0=30.0,
        demodulation=True,
        outermost_linear=False, 
    )
    # e_fou = tf.random.uniform([2*196,256,768], dtype=tf.float32)
    y = tf.random.uniform([2,196,768], dtype=tf.float32)
    o1 = layer(y)
    tf.print('o1:', tf.shape(o1))

原Siren网络 + ModulatedLinear 的实现：

import tensorflow as tf
import numpy as np
import sys

sys.path.append('')

from models.modulated_linear import ModulatedLinear


class SineLayer(tf.keras.layers.Layer):
    def __init__(
        self,
        hidden_dim,
        output_dim,
        use_bias=False,
        is_first=False,
        omega_0=30,
        demodulation=True,
        outermost_linear=False, 
    ):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.output_dim = output_dim
        self.use_bias = use_bias
        self.is_first = is_first
        self.omega_0 = omega_0
        self.demodulation = demodulation
        self.outermost_linear=outermost_linear
        kernel_initializer = None
        if is_first:
            kernel_initializer = tf.random_uniform_initializer(
                -1/output_dim,
                1/output_dim
            )
        else:
            kernel_initializer = tf.random_uniform_initializer(
                -np.sqrt(6 / output_dim) / omega_0,
                np.sqrt(6 / output_dim) / omega_0
            )
        self.linear = ModulatedLinear(
            hidden_dim=hidden_dim,
            output_dim=output_dim,
            demodulation=demodulation,
            use_bias=use_bias,
            kernel_initializer=kernel_initializer,
        )

    def call(self, inputs):
        assert isinstance(inputs, tuple) and len(inputs) == 2
        x, style = inputs
        x = self.linear((x, style))
        if self.outermost_linear:
            return x
        else:
            return tf.math.sin(self.omega_0 * x)


class Siren(tf.keras.layers.Layer):
    def __init__(
        self,
        hidden_dim,
        hidden_layers,
        out_dim,
        first_omega_0=30.0,
        hidden_omega_0=30.0,
        demodulation=True,
        outermost_linear=False,
        ):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.hidden_layers = hidden_layers
        self.out_dim = out_dim
        self.first_omega_0=first_omega_0
        self.hidden_omega_0=hidden_omega_0
        self.demodulation=demodulation
        self.outermost_linear=outermost_linear
        self.net_layers = []
        self.net_layers.append(
            SineLayer(
                hidden_dim=hidden_dim,
                output_dim=hidden_dim,
                is_first=True,
                omega_0=first_omega_0,
                demodulation=demodulation,
            )
        )

        for _ in range(hidden_layers):
            self.net_layers.append(
                SineLayer(
                    hidden_dim=hidden_dim,
                    output_dim=hidden_dim,
                    is_first=False,
                    omega_0=hidden_omega_0,
                    demodulation=demodulation,
                )
            )

        self.net_layers.append(
            SineLayer(
                hidden_dim=hidden_dim,
                output_dim=out_dim,
                is_first=False,
                omega_0=hidden_omega_0,
                demodulation=demodulation,
                outermost_linear=outermost_linear,
            )
        )

    def call(self, inputs):
        assert isinstance(inputs, tuple) and len(inputs) == 2
        x, style = inputs
        for layer in self.net_layers:
            x = layer((x, style))
        return x



if __name__ == "__main__":
    emb_dim = 768 # 16*16*3
    layer = Siren(
        hidden_dim=emb_dim,
        hidden_layers=3,
        out_dim=3,
        first_omega_0=30.0,
        hidden_omega_0=30.0,
        demodulation=True,
        outermost_linear=False, 
    )
    e_fou = tf.random.uniform([2*196,256,768], dtype=tf.float32)
    y = tf.random.uniform([2,196,768], dtype=tf.float32)
    o1 = layer((e_fou, y))
    tf.print('o1:', tf.shape(o1))

参考资料：

• 论文地址：VITGAN: Training GANs with Vision Transformers
• 资料26：Implicit Neural Representations with Periodic
  Activation Functions
• 源码：https://github.com/tfwcn/VITGAN-tf2