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

网络结构图
该章节介绍VITGAN对抗生成网络中，Mapping NetWork 部分的代码实现。

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

Mapping NetWork 简介

Generator生成器
设计基于 ViT 架构的生成器是一项不平凡的任务。一个挑战是将 ViT 从预测一组类标签转换为在空间区域上生成像素。在介绍我们的模型之前，我们讨论了两个似是而非的基线模型。两种模型交换 ViT 的输入和输出以从嵌入生成像素，特别是从 MLP 的高斯噪声向量 z 派生的潜在向量 w，即 w = MLP(z)（在图 2 中称为映射网络 [26]）。两个基线生成器的输入序列不同。该设计的灵感来自于反转 ViT，其中 w 用于替换公式 4 中的分类嵌入 h0L。
该网络部分属于Generator生成器的开始部分，输入的Latent z实际上是噪声。这部分用于将噪声转换成w，实际上可以看作是可学习的一个多维图片输入特征。

代码实现

模型是由8层全连接+LeakyReLU组成，该结构出自论文：
Tero Karras, Samuli Laine, and Timo Aila. A style-based generator architecture for generative adversarial networks. In CVPR, 2019.

import numpy as np
import tensorflow as tf


class MappingNetwork(tf.keras.layers.Layer):
    def __init__(self, d_model, num_layers, lrmul=0.01):
        super().__init__()
        self.d_model = d_model
        self.num_layers = num_layers
        self.lrmul = lrmul
        self.ln = tf.keras.layers.LayerNormalization(
            axis=-1,
            epsilon=1e-6,
            center=False,
            scale=False,
        )
        self.dn_layers = []
        for _ in range(num_layers):
            self.dn_layers.append(tf.keras.layers.Dense(d_model, use_bias=True, kernel_initializer='he_uniform', kernel_regularizer='l2'))
            self.dn_layers.append(tf.keras.layers.LeakyReLU(lrmul))
        self.net = tf.keras.Sequential(self.dn_layers)

    def call(self, x, training):
        x = self.ln(x, training=training)
        x = self.net(x)
        return x

    def get_config(self):
        config = super(MappingNetwork, self).get_config()
        config.update({
            'd_model': self.d_model,
            'num_layers': self.num_layers,
            'lrmul': self.lrmul,
        })
        return config

if __name__ == "__main__":
    layer = MappingNetwork(256, 8)
    x = tf.random.uniform([2,1,256], dtype=tf.float32)
    o1 = layer(x)
    tf.print('o1:', tf.shape(o1))