学习笔记：简易UNet网络部署

1.导包

import torch
from torch import nn

2.开始搭建

class BasicUNet(nn.Module):
    """一个十分简单的UNet网络部署"""
    # in_channels=1 表示输入图像的通道数。通常，灰度图像有一个通道，而彩色图像有三个通道（红、绿、蓝）。
    # 这个参数指定了网络模型接受的输入图像的通道数。
    # 在这个示例中，设置为1，表示网络期望接收单通道的灰度图像作为输入。
    # out_channels=1 表示网络模型的输出通道数。这通常用于指定网络的输出维度。
    # 在这个示例中，设置为1，表示网络的输出将是一个单通道的特征图。这意味着网络的目标是生成一个单通道的输出
    # 需要了解的术语：
    #         上下采样、最大池化、大步幅卷积层、卷积核、向前传播、激活函数、全连接层、残差连接。
    #         通道数、神经元、跳跃连接。
    def __init__(self, in_channels=1, out_channels=1):
        super().__init__()
        self.down_layers = torch.nn.ModuleList([
            nn.Conv2d(in_channels, 32, kernel_size=5, padding=2),
            nn.Conv2d(32, 64, kernel_size=5, padding=2),
            nn.Conv2d(64, 64, kernel_size=5, padding=2),
        ])
        self.up_layers = torch.nn.ModuleList([
            nn.Conv2d(64, 64, kernel_size=5, padding=2),
            nn.Conv2d(64, 32, kernel_size=5, padding=2),
            nn.Conv2d(32, out_channels, kernel_size=5, padding=2),
        ])
        # self.act = nn.n()
        self.act = nn.ReLU()
        # 定义一个激活函数
        self.downscale = nn.MaxPool2d(2)
        # 定义一个下采样最大池化层
        self.upscale = nn.Upsample(scale_factor=2)
        # 上采样图像，将输入数据的尺寸沿两个维度分别放大两倍。

3.定义向前传播的过程

    def forward(self, _x):
        h = []
        for i, l in enumerate(self.down_layers):
            # enumerate 是一个 Python 内置函数，
            # 用于迭代遍历可迭代对象（如列表、元组、字符串等）的元素，并在迭代过程中返回元素的索引和值。
            # 它会遍历 self.down_layers 中的每个元素。
            # 在每次迭代中，i 将被赋值为当前元素的索引，l 将被赋值为当前元素本身（即神经网络层）
            # self.down_layers 是一个 PyTorch 模型中的 ModuleList，
            # 它包含了一系列的神经网络层（在这个代码中是卷积层）。这些层按顺序存储在列表中。
            _x = self.act(l(_x))                               # 通过运算层与激活函数
            # 将输入 _x 通过卷积层 l 进行卷积运算，然后应用激活函数 self.act，
            # 最后将结果赋值给 _x，以便在网络的下一层继续处理。
            # 这是神经网络中典型的前向传播操作，用于逐层处理输入数据，提取特征并引入非线性。
            if i < 2:                                                # 选择第一、第二层
                h.append(_x)                                 # 排列供残差连接使用的数据 分别是32通道时的数据和64通道时的数据
                _x = self.downscale(_x)                 # 进行下采样以适配下一层的输入 降低图片的分辨率

        for i, l in enumerate(self.up_layers):
            if i > 0:                                                # 选择第二第三层
                _x = self.upscale(_x)                      # 上采样 提升图像的分辨率
                _x += h.pop()                                # 得到之前排列好的供残差连接使用的数据
                # Python 中列表（list）的方法，用于移除并返回列表的最后一个元素
            _x = self.act(l(_x))                              # 通过运算层与激活函数 通道数可能会发生改变
        return _x


# 在UNet网络的第三层下采样和第一层上采样中，通道数保持不变并且图片分辨率没有改变，
# 这样做的主要目的是为了实现跳跃连接（skip connection）或残差连接（residual connection），
# 以便更好地传递信息和梯度。

4.验证结果与预期是否相符

net = BasicUNet()
x = torch.rand(8, 1, 28, 28)
# net(x).shape
print(net(x).shape)
print(sum([p.numel() for p in net.parameters()]))
# 最终，输出的结果与预期那样，与输入的形状相同。
# torch.Size([8, 1, 28, 28])
# 我们所构建的UNet网络30多万个参数
# 309057

5.完整示例

import torch
# import torchvision
from torch import nn
# from torch.nn import functional as F
# from torch.utils.data import DataLoader
# from diffusers import DDPMScheduler, UNet2DModel
# from matplotlib import pyplot as plt


class BasicUNet(nn.Module):
    """一个十分简单的UNet网络部署"""
    # in_channels=1 表示输入图像的通道数。通常，灰度图像有一个通道，而彩色图像有三个通道（红、绿、蓝）。
    # 这个参数指定了网络模型接受的输入图像的通道数。
    # 在这个示例中，设置为1，表示网络期望接收单通道的灰度图像作为输入。
    # out_channels=1 表示网络模型的输出通道数。这通常用于指定网络的输出维度。
    # 在这个示例中，设置为1，表示网络的输出将是一个单通道的特征图。这意味着网络的目标是生成一个单通道的输出
    # 需要了解的术语：
    #         上下采样、最大池化、大步幅卷积层、卷积核、向前传播、激活函数、全连接层、残差连接。
    #         通道数、神经元、跳跃连接。
    def __init__(self, in_channels=1, out_channels=1):
        super().__init__()
        self.down_layers = torch.nn.ModuleList([
            nn.Conv2d(in_channels, 32, kernel_size=5, padding=2),
            nn.Conv2d(32, 64, kernel_size=5, padding=2),
            nn.Conv2d(64, 64, kernel_size=5, padding=2),
        ])
        self.up_layers = torch.nn.ModuleList([
            nn.Conv2d(64, 64, kernel_size=5, padding=2),
            nn.Conv2d(64, 32, kernel_size=5, padding=2),
            nn.Conv2d(32, out_channels, kernel_size=5, padding=2),
        ])
        # self.act = nn.n()
        self.act = nn.ReLU()
        # 定义一个激活函数
        self.downscale = nn.MaxPool2d(2)
        # 定义一个下采样最大池化层
        self.upscale = nn.Upsample(scale_factor=2)
        # 上采样图像，将输入数据的尺寸沿两个维度分别放大两倍。

    def forward(self, _x):
        h = []
        for i, l in enumerate(self.down_layers):
            # enumerate 是一个 Python 内置函数，
            # 用于迭代遍历可迭代对象（如列表、元组、字符串等）的元素，并在迭代过程中返回元素的索引和值。
            # 它会遍历 self.down_layers 中的每个元素。
            # 在每次迭代中，i 将被赋值为当前元素的索引，l 将被赋值为当前元素本身（即神经网络层）
            # self.down_layers 是一个 PyTorch 模型中的 ModuleList，
            # 它包含了一系列的神经网络层（在这个代码中是卷积层）。这些层按顺序存储在列表中。
            _x = self.act(l(_x))                               # 通过运算层与激活函数
            # 将输入 _x 通过卷积层 l 进行卷积运算，然后应用激活函数 self.act，
            # 最后将结果赋值给 _x，以便在网络的下一层继续处理。
            # 这是神经网络中典型的前向传播操作，用于逐层处理输入数据，提取特征并引入非线性。
            if i < 2:                                                # 选择第一、第二层
                h.append(_x)                                 # 排列供残差连接使用的数据 分别是32通道时的数据和64通道时的数据
                _x = self.downscale(_x)                 # 进行下采样以适配下一层的输入 降低图片的分辨率

        for i, l in enumerate(self.up_layers):
            if i > 0:                                                # 选择第二第三层
                _x = self.upscale(_x)                      # 上采样 提升图像的分辨率
                _x += h.pop()                                # 得到之前排列好的供残差连接使用的数据
                # Python 中列表（list）的方法，用于移除并返回列表的最后一个元素
            _x = self.act(l(_x))                              # 通过运算层与激活函数 通道数可能会发生改变
        return _x


# 在UNet网络的第三层下采样和第一层上采样中，通道数保持不变并且图片分辨率没有改变，
# 这样做的主要目的是为了实现跳跃连接（skip connection）或残差连接（residual connection），
# 以便更好地传递信息和梯度。
net = BasicUNet()
x = torch.rand(8, 1, 28, 28)
# net(x).shape
print(net(x).shape)
print(sum([p.numel() for p in net.parameters()]))
# 最终，输出的结果与预期那样，与输入的形状相同。
# torch.Size([8, 1, 28, 28])
# 我们所构建的UNet网络30多万个参数
# 309057