手撕Diffusion系列 - 第六期 - Unet网络

最新推荐文章于 2025-02-18 08:46:48 发布

小菜鸟博士

最新推荐文章于 2025-02-18 08:46:48 发布

阅读量1.6k

点赞数 43

分类专栏：手撕Diffusion 文章标签：网络算法人工智能深度学习学习笔记

本文链接：https://blog.youkuaiyun.com/m0_62030579/article/details/145311319

版权

手撕Diffusion系列 - 第六期 - Unet网络

DDPM 原理图

DDPM包括两个过程：前向过程（forward process）和反向过程（reverse process），其中前向过程又称为扩散过程（diffusion process），如下图所示。无论是前向过程还是反向过程都是一个参数化的马尔可夫链（Markov chain），其中反向过程可以用来生成图片。

在这里插入图片描述

DDPM 整体大概流程

图中，由高斯随机噪声 $x_T$ 生成原始图片 $x_0$ 为反向过程，反之为前向过程（噪音扩散）。

Unet网络结构图

Unet的网络架构如下图所示：

在这里插入图片描述

Unet 网络结构图

通过上述的图像，我们可以发现对于Unet来说，其编码器包含5个卷积，4个最大池化层，解码器包含4个上采样，4个反卷积层。我们分别对每个卷积层，池化层，上采样层和反卷积层进行分析。
- Part1 卷积层：对于一个卷积层来说，我们的操作是先进行第一个二维卷积扩大卷积通道数(如下图从1通道到64通道)，然后Relu激活函数，然后再第二次卷积，通道数不变。(注意：1. 图像可能图像的大小也进行了变化，但是我们这里不进行变化，2. 为了融入时间t元素，我们会在第一次卷积之后在图像上面添加对应的时间t的embdding编码)
- 卷积层的构造示例图
- Part2 下采样层：下采样层采用的是最大池化，也就是在一个池化核大小的图像元素内，取其最大值。我们这里的下采样都是对其图像进行边长减小二分之一，比如 (512，512) -> (256，256) ，所以我们的核的大小取(2,2)，步长取2，不padding。通道数不变。
- max下采样示例图
- Part3 上采样层：对于上采样，我们和下采样反着来，下采样图像边长变小一半，上采样图像边长就放大一倍数。同时，由于解码器过程中，有残差操作，并且这个操作是直接对通道进行合并，因此通道数会比编码时候多一倍，为了达成平衡，解码器的上采样同时改变通道数，把通道数减小(如下图：(看这里就把通道减小了一半，从1024->512))。
- 上采样示例图
- Part4 反采样层：对于反卷积层，也没啥特别的，和卷积层反着来就行，通道增加，只不过由于有残差，导致和卷积层的通道减小有延迟，不是完全对称的(如图中所示，这里左边的卷积是把64->128，而右边的反卷积层确是把256->128，所以说非完全对称)。
- 反卷积层示例图

Uet 网络代码

Part1 首先在conv_block.py设置一个卷积的类(用于正反卷积)

第一部分：引入相关库函数

# 该模块主要是为了实现Unet里面的横向卷积过程，每个上下采样后，都需要进行两次卷积，主要是改变图像通道数，和容易batch_t_emding信息

'''
# Part1 引入相关的库函数
'''
import torch
from dataset import *
from torch import nn

第二部分：设置卷积类

'''
# Part2 将卷积设置为一个类函数 
'''


class Conv_Block(nn.Module):
    def __init__(self, in_channel, out_channel, time_emding_size):
        super().__init__(