生成模型——自回归模型详解与PixelCNN构建

自回归模型（Autoregressive models）

深度神经网络生成算法主要分为三类：

1. 生成对抗网络(Generative Adversarial Network, GAN)
2. 可变自动编码器(Variational Autoencoder, VAE)
3. 自回归模型(Autoregressive models)

VAE已经在《变分自编码器（VAE）原理与实现（tensorflow2.x）》中进行了介绍。GAN的详细信息参考《深度卷积生成对抗网络（DCGAN）原理与实现（采用Tensorflow2.x）》。将在这里介绍鲜为人知的自回归模型，尽管自回归在图像生成中并不常见，但自回归仍然是研究的活跃领域，DeepMind的WaveNet使用自回归来生成逼真的音频。在本文中，将介绍自回归模型并构建PixelCNN模型。

简介

Autoregressive中的“Auto”意味着自我(self)，而机器学习术语的回归(regress)意味着预测新的值。将它们放在一起，自回归意味着我们使用模型基于模型的过去数据点来预测新数据点。
设图像的概率分布是$p(x)$是像素的联合概率分布$p(x_1,x_2,\dots x_n)$，由于高维数很难建模。在这里，我们假设一个像素的值仅取决于它之前的像素的值。换句话说，当前像素仅以其前一像素为条件，即$p(x_i)=p(x_i|x_{i-1})p(x_{i-1})$，我们就可以将联合概率近似为条件概率的乘积：
$$p(x)=p(x_n,x_{n-1},\dots ,x_2,x_1)$$
$$p(x)=p(x_n|x_{n-1})...p(x_3|x_2)p(x_2|x_1)p(x_1)$$
举一个具体的例子，假设在图像的中心附近包含一个红色的苹果，并且该苹果被绿叶包围，在此情况下，假设仅存在两种可能的颜色：红色和绿色。 $x_1$是左上像素，所以$p(x_1)$表示左上像素是绿色还是红色的概率。如果$x_1$为绿色，则其右边$p(x_2)$的像素也可能也为绿色，因为它可能会有更多的叶子。但是，尽管可能性较小，但它也可能是红色的。
继续进行计算，我们最终将得到红色像素。从那个像素开始，接下来的几个像素也很可能也是红色的，这比必须同时考虑所有像素要简单得多。

PixelRNN

PixelRNN由DeepMind于2016年提出。正如名称RNN（Recurrent Neural Network, 递归神经网络）所暗示的那样，该模型使用一种称为长短期记忆（LSTM）的RNN来学习图像的分布。它在LSTM中的一个步骤中一次读取图像的一行，并使用一维卷积层对其进行处理，然后将激活信息馈送到后续层中以预测该行的像素。
由于LSTM运行缓慢，因此需要花费很长时间来训练和生成样本。因此，我们不会对其进行过多的研究，而将注意力转移到同一论文中提出的一种变体——PixelCNN。

使用TensorFlow 2构建PixelCNN模型

PixelCNN仅由卷积层组成，使其比PixelRNN快得多。在这里，我们将为使用MNIST数据集训练一个简单的PixelCNN模型。

输入和标签

MNIST由28 x 28 x 1灰度数字手写数字组成。它只有一个通道：

在本实验中，通过将图像转换为二进制数据来简化问题：0代表黑色，1代表白色：

def binarize(image, label):
    image = tf.cast(image, tf.float32)
    image = tf.math.round(image/255.)
    return image, tf.cast(image, tf.int32)

该函数需要两个输入——图像和标签。该函数的前两行将图像转换为二进制float32格式，即0.0或1.0。并且，我们将二进制图像转换为整数并返回，以遵循使用整数作为标签的惯例而已。返回的数据，将作为网络训练的输入和标签，都是28 x 28 x 1的二进制MNIST图像，它们仅在数据类型上有所不同。

掩膜

与PixelRNN逐行读取不同，PixelCNN在图像中从左到右，从上到下滑动卷积核。当执行卷积以预测当前像素时，传统的卷积核能够看到当前输入像素及其周围的像素，其中包括当前像素之后的像素信息，这与在简介部分的条件概率假设相悖。
为了避免这种情况，我们需要确保CNN在预测输出像素$x_i$时不会看到输入像素$x_i$。
这是通过使用掩膜卷积来实现的，其中在执行卷积之前将掩膜应用于卷积核权重。下图显示了一个7 x 7卷积核的掩膜，其中从中心开始的权重为0。这会阻止CNN看到它正在预测的像素（卷积核的中心）以及所有之后的像素。这称为A型掩膜，仅应用于输入层。

由于中心像素在第一层中被遮挡，因此我们不再需要在后面的层中隐藏中心要素。实际上，我们需要将卷积核中心设置为1，以使其能够读取先前层的特征，这称为B型掩膜。

实现自定义层

现在，我们将为掩膜卷积创建一个自定义层。我们可以使用从基类tf.keras.layers.Layer继承的子类在TensorFlow2.x中创建自定义层，以便将能够像使用其他Keras层一样使用它。以下是自定义层类的基本结构：

class MaskedConv2D(tf.keras.layers.Layer):
    def __init__(self):
        ...       
    def build(self, input_shape):
        ...
    def call(self, inputs):
        ...
        return output

build()将输入张量的形状作为参数，我们将使用此信息来确保创建正确形状的变量。构建图层时，此函数仅运行一次。我们可以通过声明不可训练的变量或常量来创建掩码，以使TensorFlow知道它不需要梯度来反向传播：

	def build(self, input_shape):
	        self.w = self.add_weight(shape=[self.kernel,
	                                        self.kernel,
	                                        input_shape[-1],
	                                        self.filters],
	                                initializer='glorot_normal',
	                                trainable=True)
	        self.b = self.add_weight(shape=(self.filters,),
	                                initializer='zeros',
	                                trainable=True)
	        mask = np.ones(self.kernel**2, dtype=np.float32)
	        center = len(mask)//2
	        mask[center+1:] = 0
	        if self.mask_type == 'A':
	            mask[center] = 0
	        mask = mask.reshape((self.kernel, self.kernel, 1, 1))
	        self.mask = tf.constant(mask, dtype='float32')

call()用来执行前向传递。在掩膜卷积层中，在使用低级tf.nn API执行卷积之前，我们将权重乘以掩码后将下半部分的值设为零：

	def call(self, inputs):
	        masked_w = tf.math.multiply(self.w, self.mask)
	        output=tf.nn.conv2d(inputs, masked_w, 1, "SAME") + self.b
	        return output

网络架构

PixelCNN架构非常简单。在使用A型掩膜的第一个7 x 7 conv2d图层之后，有几层带有B型掩膜的残差块：

Model: "PixelCnn"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
input_2 (InputLayer)         [(None, 28, 28, 1)]       0         
_________________________________________________________________
masked_conv2d_22 (MaskedConv (None, 28, 28, 128)       6400      
_________________________________________________________________
residual_block_7 (ResidualBl (None, 28, 28, 128)       53504     
_________________________________________________________________
residual_block_8 (ResidualBl (None, 28, 28, 128)       53504     
_________________________________________________________________
residual_block_9 (ResidualBl (None, 28, 28, 128)       53504     
_________________________________________________________________
residual_block_10 (ResidualB (None, 28, 28, 128)       53504     
_________________________________________________________________
residual_block_11 (ResidualB (None, 28, 28, 128)       53504     
_________________________________________________________________
residual_block_12 (ResidualB (None, 28, 28, 128)       53504     
_________________________________________________________________
residual_block_13 (ResidualB (None, 28, 28, 128)       53504     
_________________________________________________________________
conv2d_2 (Conv2D)            (None, 28, 28, 64)        8256      
_________________________________________________________________
conv2d_3 (Conv2D)            (None, 28, 28, 1)         65        
=================================================================
Total params: 389,249
Trainable params: 389,249
Non-trainable params: 0

下图说明了PixelCNN中使用的残差块架构：

交叉熵损失

交叉熵损失也称为对数损失，它衡量模型的性能，其中输出的概率在0到1之间。以下是二进制交叉熵损失的方程，其中只有两个类，标签y可以是0或1，$p(x)$是模型的预测：
$$BCE=-\frac{1}{N}\sum _{i=1}^N(y_{i}logp(x)+(1-y_i)log(1-p(x)))$$
在PixelCNN中，单个图像像素用作标签。在二值化MNIST中，我们要预测输出像素是0还是1，这使其成为使用交叉熵作为损失函数的分类问题。
最后，编译和训练神经网络，我们对损失和度量均使用二进制交叉熵，并使用RMSprop作为优化器。有许多不同的优化器可供使用，它们的主要区别在于它们根据过去的统计信息调整学习率的方式。没有一种最佳的优化器可以在所有情况下使用，因此建议尝试使用不同的优化器。
编译和训练pixelcnn模型：

pixelcnn = SimplePixelCnn()
pixelcnn.compile(
    loss = tf.keras.losses.BinaryCrossentropy(),
    optimizer=tf.keras.optimizers.RMSprop(learning_rate=0.001),
    metrics=[ tf.keras.metrics.BinaryCrossentropy()])
pixelcnn.fit(ds_train, epochs = 10, validation_data=ds_test)

接下来，我们将根据先前的模型生成一个新图像。

采样生成图片

训练后，我们可以通过以下步骤使用该模型生成新图像：

1. 创建一个具有与输入图像相同形状的空张量，并用1填充。将其馈入网络并获得$p(x1)$，即第一个像素的概率。
2. 从$p(x_1)$进行采样，并将采样值分配给输入张量中的像素$x_1$。
3. 再次将输入提供给网络，并对下一个像素执行步骤2。
4. 重复步骤2和3，直到生成$x_N$。

自回归模型的一个主要缺点是它生成速度慢，因为需要逐像素生成，而无法并行化。以下图像是我们的PixelCNN模型经过100个训练周期后生成的。它们看起来还不太像正确的数字，但我们现在可以凭空生成新图像。可以通过训练更长的模型并进行一些超参数调整来生成更好的数字。

完整代码

import tensorflow as tf
from tensorflow.keras import layers
from tensorflow.keras.activations import relu
from tensorflow.keras.models import Sequential

import tensorflow_datasets as tfds

import numpy as np
import matplotlib.pyplot as plt

import warnings
warnings.filterwarnings('ignore')
print(tf.__version__)

(ds_train, ds_test), ds_info = tfds.load(
    'mnist',
    split=['test', 'test'],
    shuffle_files=True,
    as_supervised=True,
    with_info=True)
fig = tfds.show_examples(ds_train, ds_info)

def binarize(image, label):
    image = tf.cast(image, tf.float32)
    image = tf.math.round(image/255.)
    return image, tf.cast(image, tf.int32)
    
ds_train = ds_train.map(binarize)
ds_train = ds_train.cache() # put dataset into memory
ds_train = ds_train.shuffle(ds_info.splits['train'].num_examples)
ds_train = ds_train.batch(128)
ds_test = ds_test.map(binarize).batch(128).prefetch(64)

class MaskedConv2D(layers.Layer):
    def __init__(self, mask_type, kernel=5, filters=1):
        super(MaskedConv2D, self).__init__()
        self.kernel = kernel
        self.filters = filters
        self.mask_type = mask_type
    
    def build(self, input_shape):
        
        self.w = self.add_weight(shape=[
            self.kernel, 
            self.kernel, 
            input_shape[-1],
            self.filters],
            initializer='glorot_normal',
            trainable=True)
        
        self.b = self.add_weight(shape=(self.filters,),
            initializer='zeros',
            trainable=True)
        
        # Create Mask
        mask = np.ones(self.kernel ** 2, dtype=np.float32)
        center = len(mask) // 2
        mask[center+1:] = 0
        if self.mask_type == 'A':
            mask[center] = 0
        
        mask = mask.reshape((self.kernel, self.kernel, 1, 1))

        self.mask = tf.constant(mask, dtype='float32')
    
    def call(self, inputs):
        # mask the convolution
        masked_w = tf.math.multiply(self.w, self.mask)

        # preform conv2d using low level API
        output = tf.nn.conv2d(inputs, masked_w, 1, 'SAME') + self.b

        return tf.nn.relu(output)

class ResidualBlock(layers.Layer):
    def __init__(self, h=32):
        super(ResidualBlock, self).__init__()

        self.forward = Sequential([
            MaskedConv2D('B', kernel=1, filters=h),
            MaskedConv2D('B', kernel=3, filters=h),
            MaskedConv2D('B', kernel=1, filters=2*h),
        ])
    
    def call(self, inputs):
        x = self.forward(inputs)
        return x + inputs

def SimplePixelCnn(
    hidden_features=64,
    output_features=64,
    resblocks_num=7):
    
    inputs = layers.Input(shape=[28, 28, 1])
    x = inputs
    
    x = MaskedConv2D('A', kernel=7, filters=2*hidden_features)(x)

    for _ in range(resblocks_num):
        x = ResidualBlock(hidden_features)(x)
    
    x = layers.Conv2D(output_features, (1,1), padding='same', activation='relu')(x)
    x = layers.Conv2D(1, (1,1), padding='same', activation='sigmoid')(x)

    return tf.keras.Model(inputs=inputs, outputs=x, name='PixelCnn')

pixel_cnn = SimplePixelCnn()
pixel_cnn.summary()

pixel_cnn.compile(
    loss=tf.keras.losses.BinaryCrossentropy(),
    optimizer=tf.keras.optimizers.RMSprop(learning_rate=0.001),
    metrics=[tf.keras.losses.BinaryCrossentropy()]
)

grid_row = 5
grid_col = 5
batch = grid_row * grid_col
h = w = 28
images = np.ones((batch,h,w,1), dtype=np.float32)

for row in range(h):
    for col in range(w):
        prob = pixel_cnn.predict(images)[:, row,col,0]
        pixel_samples = tf.random.categorical(
            tf.math.log(np.stack([1-prob,prob],1)),1
        )
        #print(pixel_samples.shape)
        images[:,row,col,0] = tf.reshape(pixel_samples, [batch])

# Display
f, axarr = plt.subplots(grid_row, grid_col, figsize=(grid_col*1.1,grid_row))

i = 0
for row in range(grid_row):
    for col in range(grid_col):
        axarr[row,col].imshow(images[i,:,:,0], cmap='gray')
        axarr[row,col].axis('off')
        i += 1
f.tight_layout(0.1, h_pad=0.2, w_pad=0.1)        
plt.show()