pytorch笔记——CycleGAN代码（2）：models.py

cyclegan训练模型
务必保证datasets.py能正确读取数据。

目录

pytorch笔记——CycleGAN代码（1）：datasets.py
pytorch笔记——CycleGAN代码（2）：models.py
pytorch笔记——CycleGAN代码（3）：utils.py
pytorch笔记——CycleGAN代码（4）：train.py

视频学习：
【【目前B站最全pytorch系列项目实战】70个练手项目合集，七天练完，练完即可就业！（从入门到精通，小白也能学会）-哔哩哔哩】 https://b23.tv/OG2TPEH

环境
python 3.9
tensorflow 2.11.0
torch 1.3.1
cuda 0.0.1
Windows
cpu

代码

所有可能出现的问题及解决方法都在代码注释行里。

from torch import nn
import torch.nn.functional as F

class resBlock(nn.Module):

    def __init__(self, in_channel):
        super(resBlock, self).__init__()
        conv_block = [
            # 对输入的feature大小敏感，避免图像的损失，故使用pad方法
            nn.ReflectionPad2d(1),
            nn.Conv2d(in_channels=in_channel, out_channels=in_channel, kernel_size=3),
            nn.InstanceNorm2d(in_channel),
            nn.ReLU(inplace=True),
            nn.ReflectionPad2d(1),
            nn.Conv2d(in_channels=in_channel, out_channels=in_channel, kernel_size=3),
            nn.InstanceNorm2d(in_channel)
        ]
        self.conv_block = nn.Sequential(*conv_block)

    def forward(self, x):
        return x + self.conv_block(x)


class Generator(nn.Module):
    def __init__(self):
        super(Generator, self).__init__()
        net = [
            nn.ReflectionPad2d(3),
            nn.Conv2d(in_channels=3, out_channels=64, kernel_size=7),
            nn.InstanceNorm2d(64),
            nn.ReLU(inplace=True)
        ]

        # 下采样 四倍下采样
        print("——————Generator下采样——————")
        in_channel = 64
        out_channel = in_channel * 2

        for _ in range(2):
            net += [
                nn.Conv2d(in_channels=in_channel, out_channels=out_channel, kernel_size=3, stride=2, padding=1),
                nn.InstanceNorm2d(out_channel),
                nn.ReLU(inplace=True)
            ]
            in_channel = out_channel
            out_channel = in_channel * 2

        for _ in range(9):
            net += [resBlock(in_channel)]

        # 上采样
        print("——————Generator上采样——————")
        out_channel = in_channel // 2
        for _ in range(2):
            # 保证图片进行整数倍的上采样
            net += [nn.ConvTranspose2d(in_channels=in_channel, out_channels=out_channel,
                                       kernel_size=3,
                                       stride=2, padding=1, output_padding=1),
                    nn.InstanceNorm2d(out_channel),
                    nn.ReLU(inplace=True)]
            in_channel = out_channel
            out_channel = in_channel // 2

        net += [
            nn.ReflectionPad2d(3),
            nn.Conv2d(64, 3, 7),
            nn.Tanh()
        ]

        self.model = nn.Sequential(*net)

    def forward(self, x):
        return self.model(x)

# 判别器
class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()
        print("——————Discriminator下采样——————")
        # 下采样
        model = [
            nn.Conv2d(3, 64, 4, stride=2, padding=1),
            nn.LeakyReLU(0.2, inplace=True)
        ]
        model += [
            nn.Conv2d(64, 128, 4, stride=2, padding=1),
            nn.InstanceNorm2d(128),
            nn.LeakyReLU(0.2, inplace=True)
        ]
        model += [
            nn.Conv2d(128, 256, 4, stride=2, padding=1),
            nn.InstanceNorm2d(256),
            nn.LeakyReLU(0.2, inplace=True)
        ]
        model += [
            nn.Conv2d(256, 512, 4, stride=2, padding=1),
            nn.InstanceNorm2d(512),
            nn.LeakyReLU(0.2, inplace=True)
        ]

        # 将输出映射到1维tensor
        model += [nn.Conv2d(512, 1, 4, padding=1)]

        self.model = nn.Sequential(*model)

        # 注：即使经过了4倍下采样，特征图仍不是1x1的大小

    def forward(self, x):
        x = self.model(x)
        # 将W和H维度处理为1x1 或者 batchsize x 1
        # 可以直接 return F.avg_pool2d(x, x.size()[2:]).view(x.size()[0], -1)
        # 但UserWarning: Using a target size (torch.Size([1])) that is different to the input size (torch.Size([1, 1]))
        # 所以需要降维在return x之前，加一句x = x.squeeze(-1)以达到降低一维的目的；
        y = F.avg_pool2d(x, x.size()[2:]).view(x.size()[0], -1)
        return y.squeeze(-1)


if __name__ == '__main__':
    G = Generator()
    D = Discriminator()
    import torch
    # 测试一下
    input_tensor = torch.ones((1, 3, 256, 256), dtype=torch.float)

    out = G(input_tensor)
    print("Generator的输出: {}".format(out.size()))

    out = D(input_tensor)
    print("Discriminator的输出: {}".format(out.size()))