AIGC基础，从VAE到DDPM原理、代码详解

Python                  
class DownConvLayer(tf.keras.layers.Layer):                  
    def __init__(self, dim):                  
        super(DownConvLayer, self).__init__()                  
        self.conv = tf.keras.layers.Conv2D(dim, 3, activation=tf.keras.layers.ReLU(), use_bias=False, padding='same')                  
        self.pool = tf.keras.layers.MaxPool2D(2)                  
                 
    def call(self, x, training=False, **kwargs):                  
        x = self.conv(x)                  
        x = self.pool(x)                  
        return x                  
                 
                 
class UpConvLayer(tf.keras.layers.Layer):                  
    def __init__(self, dim):                  
        super(UpConvLayer, self).__init__()                  
        self.conv = tf.keras.layers.Conv2D(dim, 3, activation=tf.keras.layers.ReLU(), use_bias=False, padding='same')                  
        # 通过UpSampling2D上采样                  
        self.pool = tf.keras.layers.UpSampling2D(2)                  
                 
    def call(self, x, training=False, **kwargs):                  
        x = self.conv(x)                  
        x = self.pool(x)                  
        return x                  
                 
# 示例代码都是通过非常简单的卷积操作实现编码器和解码器                  
class Encoder(tf.keras.layers.Layer):                  
    def __init__(self, dim, layer_num=3):                  
        super(Encoder, self).__init__()                  
        self.convs = [DownConvLayer(dim) for _ in range(layer_num)]                  
                 
    def call(self, x, training=False, **kwargs):                  
        for conv in self.convs:                  
            x = conv(x, training)                  
        return x                  
                 
                 
class Decoder(tf.keras.layers.Layer):                  
    def __init__(self, dim, layer_num=3):                  
        super(Decoder, self).__init__()                  
        self.convs = [UpConvLayer(dim) for _ in range(layer_num)]                  
        self.final_conv = tf.keras.layers.Conv2D(1, 3, strides=1)                  
                 
    def call(self, x, training=False, **kwargs):                  
        for conv in self.convs:                  
            x = conv(x, training)                  
        # 将图像转成和输入图像shape一致                  
        reconstruct = self.final_conv(x)                  
        return reconstruct                  
                 
                 
class AutoEncoderModel(tf.keras.Model):                  
    def __init__(self):                  
        super(AutoEncoderModel, self).__init__()                  
        self.encoder = Encoder(64, layer_num=3)                  
        self.decoder = Decoder(64, layer_num=3)                  
                 
    def call(self, inputs, training=None, mask=None):                  
        image = inputs[0]                  
        # 得到图像的特征表示                  
        latent = self.encoder(image, training)                  
        # 通过特征重建图像                  
        reconstruct_img = self.decoder(latent, training)                  
        return reconstruct_img                  
                 
    @tf.function                  
    def train_step(self, data):                  
        img = data["image"]                  
        with tf.GradientTape() as tape:                  
            reconstruct_img = self((img,), True)                  
        trainable_vars = self.trainable_variables                  
        # 利用l2 loss 来判断重建图片和原始图像的一致性                  
        l2_loss = (reconstruct_img - img) ** 2                  
        l2_loss = tf.reduce_mean(tf.reduce_sum(                  
            l2_loss, axis=(1, 2, 3)                  
        ))                  
        gradients = tape.gradient(l2_loss, trainable_vars)                  
        self.optimizer.apply_gradients(zip(gradients, trainable_vars))                  
        return {"l2_loss": l2_loss}