使用像素级标签通过 pix2pix 实现图像翻译

1import torch.nn as nn
2
3class UnetGenerator(nn.Module):
4    def __init__(self):
5        super(UnetGenerator, self).__init__()
6        unet_block = UnetSkipConnectionBlock(64 * 8, 64 * 8, submodule=None, innermost=True)
7        for i in range(8 - 5):
8            unet_block = UnetSkipConnectionBlock(64 * 8, 64 * 8, submodule=unet_block, use_dropout=True)
9        unet_block = UnetSkipConnectionBlock(64 * 4, 64 * 8, submodule=unet_block)
10        unet_block = UnetSkipConnectionBlock(64 * 2, 64 * 4, submodule=unet_block)
11        unet_block = UnetSkipConnectionBlock(64, 64 * 2, submodule=unet_block)
12        self.model = UnetSkipConnectionBlock(3, 64, input_nc=3, submodule=unet_block, outermost=True)
13
14    def forward(self, input):
15        return self.model(input)

1class UnetSkipConnectionBlock(nn.Module):
2    def __init__(self):
3        down = [nn.LeakyReLU(0.2, inplace=True),
4                nn.Conv2d(64 * 8, 64 * 8, kernel_size=4, stride=2, padding=1, bias=False)]
5        up = [nn.ReLU(inplace=True),
6              nn.ConvTranspose2d(64 * 8, 64 * 8, kernel_size=4, stride=2, padding=1, bias=False),
7              nn.BatchNorm2d(64 * 8)]
8        model = down + up
9        self.model = nn.Sequential(*model)
10
11    def forward(self, x):
12        return torch.cat([x, self.model(x)], 1)

1class UnetSkipConnectionBlock(nn.Module):
2    def __init__(self, out_channels, in_channels, submodule, use_dropout):
3        down = [nn.LeakyReLU(0.2, inplace=True),
4                nn.Conv2d(out_channels, in_channels, kernel_size=4, stride=2, padding=1, bias=False),
5                nn.BatchNorm2d(in_channels)]
6        up = [nn.ReLU(inplace=True),
7              nn.ConvTranspose2d(in_channels * 2, out_channels, kernel_size=4, stride=2, padding=1, bias=False),
8              nn.BatchNorm2d(out_channels)]
9        if use_dropout:
10            model = down + [submodule] + up + [nn.Dropout(0.5)]
11        else:
12            model = down + [submodule] + up
13        self.model = nn.Sequential(*model)
14
15    def forward(self, x):
16        return torch.cat([x, self.model(x)], 1)

1class UnetSkipConnectionBlock(nn.Module):
2    def __init__(self):
3        down = [nn.Conv2d(3, 64, kernel_size=4, stride=2, padding=1, bias=False)]
4        up = [nn.ReLU(inplace=True),
5              nn.ConvTranspose2d(64 * 2, 3, kernel_size=4, stride=2, padding=1), nn.Tanh()]
6        model = down + [submodule] + up
7        self.model = nn.Sequential(*model)
8
9    def forward(self, x):
10        return self.model(x)

1import torch
2import torch.nn as nn
3from torch.nn import init
4import functools
5from torch.optim import lr_scheduler
6
7
8###############################################################################
9# 辅助函数
10###############################################################################
11
12
13class Identity(nn.Module):
14    def forward(self, x):
15        return x
16
17
18def get_norm_layer(norm_type='instance'):
19    """返回一个归一化层
20
21    参数：
22        norm_type (str) -- 归一化层的名称：batch | instance | none
23
24    对于 BatchNorm，我们使用可学习的仿射参数并跟踪运行统计信息（均值/标准差）。
25    对于 InstanceNorm，我们不使用可学习的仿射参数，也不跟踪运行统计信息。
26    """
27    if norm_type == 'batch':
28        norm_layer = functools.partial(nn.BatchNorm2d, affine=True, track_running_stats=True)
29    elif norm_type == 'instance':
30        norm_layer = functools.partial(nn.InstanceNorm2d, affine=False, track_running_stats=False)
31    elif norm_type == 'none':
32        def norm_layer(x):
33            return Identity()
34    else:
35        raise NotImplementedError('未找到归一化层 [%s]' % norm_type)
36    return norm_layer
37
38
39def get_scheduler(optimizer, opt):
40    """返回一个学习率调度器
41
42    参数：
43        optimizer          -- 网络的优化器
44        opt (选项类) -- 存储所有实验标志；需要是 BaseOptions 的子类。
45                         opt.lr_policy 是学习率策略的名称：linear | step | plateau | cosine
46
47    对于 'linear'，在前 <opt.n_epochs> 个 epoch 中保持相同的学习率，
48    然后在接下来的 <opt.n_epochs_decay> 个 epoch 中线性衰减至零。
49    对于其他调度器（step、plateau 和 cosine），我们使用 PyTorch 默认调度器。
50    更多详情请参见 https://pytorch.org/docs/stable/optim.html。
51    """
52    if opt.lr_policy == 'linear':
53        def lambda_rule(epoch):
54            lr_l = 1.0 - max(0, epoch + opt.epoch_count - opt.n_epochs) / float(opt.n_epochs_decay + 1)
55            return lr_l
56        scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule)
57    elif opt.lr_policy == 'step':
58        scheduler = lr_scheduler.StepLR(optimizer, step_size=opt.lr_decay_iters, gamma=0.1)
59    elif opt.lr_policy == 'plateau':
60        scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.2, threshold=0.01, patience=5)
61    elif opt.lr_policy == 'cosine':
62        scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=opt.n_epochs, eta_min=0)
63    else:
64        return NotImplementedError('未实现的学习率策略 [%s]' % opt.lr_policy)
65    return scheduler
66
67
68def init_weights(net, init_type='normal', init_gain=0.02):
69    """初始化网络权重。
70
71    参数：
72        net (network)   -- 要初始化的网络
73        init_type (str) -- 初始化方法名称：normal | xavier | kaiming | orthogonal
74        init_gain (float)    -- normal、xavier 和 orthogonal 方法的缩放因子。
75
76    我们在原始 pix2pix 和 CycleGAN 论文中使用 'normal'。但 xavier 和 kaiming 可能在某些应用中效果更好。
77    欢迎自行尝试。
78    """
79    def init_func(m):  # 定义初始化函数
80        classname = m.__class__.__name__
81        if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):
82            if init_type == 'normal':
83                init.normal_(m.weight.data, 0.0, init_gain)
84            elif init_type == 'xavier':
85                init.xavier_normal_(m.weight.data, gain=init_gain)
86            elif init_type == 'kaiming':
87                init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
88            elif init_type == 'orthogonal':
89                init.orthogonal_(m.weight.data, gain=init_gain)
90            else:
91                raise NotImplementedError('未实现的初始化方法 [%s]' % init_type)
92            if hasattr(m, 'bias') and m.bias is not None:
93                init.constant_(m.bias.data, 0.0)
94        elif classname.find('BatchNorm2d') != -1:  # BatchNorm 层的权重不是矩阵；仅适用正态分布。
95            init.normal_(m.weight.data, 1.0, init_gain)
96            init.constant_(m.bias.data, 0.0)
97
98    print('使用 %s 初始化网络' % init_type)
99    net.apply(init_func)  # 应用初始化函数 <init_func>
100
101
102def init_net(net, init_type='normal', init_gain=0.02, gpu_ids=[]):
103    """初始化网络：1. 注册 CPU/GPU 设备（支持多 GPU）；2. 初始化网络权重
104    参数：
105        net (network)      -- 要初始化的网络
106        init_type (str)    -- 初始化方法名称：normal | xavier | kaiming | orthogonal
107        gain (float)       -- normal、xavier 和 orthogonal 方法的缩放因子。
108        gpu_ids (int list) -- 网络运行的 GPU 编号列表：例如 0,1,2
109
110    返回一个已初始化的网络。
111    """
112    if len(gpu_ids) > 0:
113        assert(torch.cuda.is_available())
114        net.to(gpu_ids[0])
115        net = torch.nn.DataParallel(net, gpu_ids)  # 多 GPU
116    init_weights(net, init_type, init_gain=init_gain)
117    return net
118
119
120def define_G(input_nc, output_nc, ngf, netG, norm='batch', use_dropout=False, init_type='normal', init_gain=0.02, gpu_ids=[]):
121    """创建一个生成器
122
123    参数：
124        input_nc (int) -- 输入图像的通道数
125        output_nc (int) -- 输出图像的通道数
126        ngf (int) -- 最后一个卷积层中的滤波器数量
127        netG (str) -- 架构名称：resnet_9blocks | resnet_6blocks | unet_256 | unet_128
128        norm (str) -- 网络中使用的归一化层类型：batch | instance | none
129        use_dropout (bool) -- 是否使用 dropout 层。
130        init_type (str)    -- 初始化方法名称。
131        init_gain (float)  -- normal、xavier 和 orthogonal 方法的缩放因子。
132        gpu_ids (int list) -- 网络运行的 GPU 列表：例如 0,1,2
133
134    返回一个生成器
135
136    当前实现提供两种类型的生成器：
137        U-Net: [unet_128]（用于 128x128 输入图像）和 [unet_256]（用于 256x256 输入图像）
138        原始 U-Net 论文：https://arxiv.org/abs/1505.04597
139
140        基于 Resnet 的生成器: [resnet_6blocks]（包含 6 个 Resnet 模块）和 [resnet_9blocks]（包含 9 个 Resnet 模块）
141        基于 Resnet 的生成器由若干 Resnet 模块和一些下采样/上采样操作组成。
142        我们改编自 Justin Johnson 的神经风格迁移项目代码（https://github.com/jcjohnson/fast-neural-style）。
143
144    生成器已通过 <init_net> 初始化。它使用 RELU 作为非线性激活函数。
145    """
146    net = None
147    norm_layer = get_norm_layer(norm_type=norm)
148
149    if netG == 'resnet_9blocks':
150        net = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=9)
151    elif netG == 'resnet_6blocks':
152        net = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=6)
153    elif netG == 'unet_128':
154        net = UnetGenerator(input_nc, output_nc, 7, ngf, norm_layer=norm_layer, use_dropout=use_dropout)
155    elif netG == 'unet_256':
156        net = UnetGenerator(input_nc, output_nc, 8, ngf, norm_layer=norm_layer, use_dropout=use_dropout)
157    else:
158        raise NotImplementedError('未识别的生成器模型名称 [%s]' % netG)
159    return init_net(net, init_type, init_gain, gpu_ids)
160
161
162def define_D(input_nc, ndf, netD, n_layers_D=3, norm='batch', init_type='normal', init_gain=0.02, gpu_ids=[]):
163    """创建一个判别器
164
165    参数：
166        input_nc (int)     -- 输入图像的通道数
167        ndf (int)          -- 第一个卷积层中的滤波器数量
168        netD (str)         -- 架构名称：basic | n_layers | pixel
169        n_layers_D (int)   -- 判别器中卷积层的数量；当 netD=='n_layers' 时有效
170        norm (str)         -- 网络中使用的归一化层类型。
171        init_type (str)    -- 初始化方法名称。
172        init_gain (float)  -- normal、xavier 和 orthogonal 方法的缩放因子。
173        gpu_ids (int list) -- 网络运行的 GPU 列表：例如 0,1,2
174
175    返回一个判别器
176
177    当前实现提供三种类型的判别器：
178        [basic]: 原始 pix2pix 论文中描述的 'PatchGAN' 分类器。
179        它可以判断 70×70 重叠图像块是真实还是伪造。
180        这种基于图像块的判别器架构比全图判别器参数更少，
181        并且可以以全卷积方式处理任意尺寸的图像。
182
183        [n_layers]: 使用此模式，你可以通过参数 <n_layers_D> 指定判别器中卷积层的数量
184        （默认为 3，与 [basic]（PatchGAN）中使用的相同）。
185
186        [pixel]: 1x1 PixelGAN 判别器可以判断每个像素是否真实。
187        它鼓励更高的颜色多样性，但对空间统计无影响。
188
189    判别器已通过 <init_net> 初始化。它使用 Leakly RELU 作为非线性激活函数。
190    """
191    net = None
192    norm_layer = get_norm_layer(norm_type=norm)
193
194    if netD == 'basic':  # 默认 PatchGAN 分类器
195        net = NLayerDiscriminator(input_nc, ndf, n_layers=3, norm_layer=norm_layer)
196    elif netD == 'n_layers':  # 更多选项
197        net = NLayerDiscriminator(input_nc, ndf, n_layers_D, norm_layer=norm_layer)
198    elif netD == 'pixel':     # 判断每个像素是真实还是伪造
199        net = PixelDiscriminator(input_nc, ndf, norm_layer=norm_layer)
200    else:
201        raise NotImplementedError('未识别的判别器模型名称 [%s]' % netD)
202    return init_net(net, init_type, init_gain, gpu_ids)
203
204
205##############################################################################
206# 类定义
207##############################################################################
208class GANLoss(nn.Module):
209    """定义不同的 GAN 目标函数。
210
211    GANLoss 类抽象了创建与输入同尺寸目标标签张量的需求。
212    """
213
214    def __init__(self, gan_mode, target_real_label=1.0, target_fake_label=0.0):
215        """ 初始化 GANLoss 类。
216
217        参数：
218            gan_mode (str) - - GAN 目标类型。目前支持 vanilla、lsgan 和 wgangp。
219            target_real_label (bool) - - 真实图像的标签
220            target_fake_label (bool) - - 伪造图像的标签
221
222        注意：不要在判别器的最后一层使用 sigmoid。
223        LSGAN 不需要 sigmoid。vanilla GAN 将通过 BCEWithLogitsLoss 处理。
224        """
225        super(GANLoss, self).__init__()
226        self.register_buffer('real_label', torch.tensor(target_real_label))
227        self.register_buffer('fake_label', torch.tensor(target_fake_label))
228        self.gan_mode = gan_mode
229        if gan_mode == 'lsgan':
230            self.loss = nn.MSELoss()
231        elif gan_mode == 'vanilla':
232            self.loss = nn.BCEWithLogitsLoss()
233        elif gan_mode in ['wgangp']:
234            self.loss = None
235        else:
236            raise NotImplementedError('未实现的 gan 模式 %s' % gan_mode)
237
238    def get_target_tensor(self, prediction, target_is_real):
239        """创建与输入同尺寸的标签张量。
240
241        参数：
242            prediction (tensor) - - 通常是判别器的预测输出
243            target_is_real (bool) - - 真实标签对应的是真实图像还是伪造图像
244
245        返回：
246            一个填充了真实标签且与输入尺寸相同的标签张量
247        """
248
249        if target_is_real:
250            target_tensor = self.real_label
251        else:
252            target_tensor = self.fake_label
253        return target_tensor.expand_as(prediction)
254
255    def __call__(self, prediction, target_is_real):
256        """根据判别器输出和真实标签计算损失。
257
258        参数：
259            prediction (tensor) - - 通常是判别器的预测输出
260            target_is_real (bool) - - 真实标签对应的是真实图像还是伪造图像
261
262        返回：
263            计算出的损失值。
264        """
265        if self.gan_mode in ['lsgan', 'vanilla']:
266            target_tensor = self.get_target_tensor(prediction, target_is_real)
267            loss = self.loss(prediction, target_tensor)
268        elif self.gan_mode == 'wgangp':
269            if target_is_real:
270                loss = -prediction.mean()
271            else:
272                loss = prediction.mean()
273        return loss
274
275
276def cal_gradient_penalty(netD, real_data, fake_data, device, type='mixed', constant=1.0, lambda_gp=10.0):
277    """计算梯度惩罚损失，用于 WGAN-GP 论文 https://arxiv.org/abs/1704.00028
278
279    参数：
280        netD (network)              -- 判别器网络
281        real_data (tensor array)    -- 真实图像
282        fake_data (tensor array)    -- 生成器生成的图像
283        device (str)                -- GPU / CPU: 来自 torch.device('cuda:{}'.format(self.gpu_ids[0])) if self.gpu_ids else torch.device('cpu')
284        type (str)                  -- 是否混合真实和伪造数据 [real | fake | mixed]。
285        constant (float)            -- 公式中使用的常数 (||gradient||_2 - constant)^2
286        lambda_gp (float)           -- 此损失的权重
287
288    返回梯度惩罚损失
289    """
290    if lambda_gp > 0.0:
291        if type == 'real':   # 使用真实图像、伪造图像或两者的线性插值。
292            interpolatesv = real_data
293        elif type == 'fake':
294            interpolatesv = fake_data
295        elif type == 'mixed':
296            alpha = torch.rand(real_data.shape[0], 1, device=device)
297            alpha = alpha.expand(real_data.shape[0], real_data.nelement() // real_data.shape[0]).contiguous().view(*real_data.shape)
298            interpolatesv = alpha * real_data + ((1 - alpha) * fake_data)
299        else:
300            raise NotImplementedError('{} 未实现'.format(type))
301        interpolatesv.requires_grad_(True)
302        disc_interpolates = netD(interpolatesv)
303        gradients = torch.autograd.grad(outputs=disc_interpolates, inputs=interpolatesv,
304                                        grad_outputs=torch.ones(disc_interpolates.size()).to(device),
305                                        create_graph=True, retain_graph=True, only_inputs=True)
306        gradients = gradients[0].view(real_data.size(0), -1)  # 展平数据
307        gradient_penalty = (((gradients + 1e-16).norm(2, dim=1) - constant) ** 2).mean() * lambda_gp        # 添加 eps
308        return gradient_penalty, gradients
309    else:
310        return 0.0, None
311
312
313class ResnetGenerator(nn.Module):
314    """基于 Resnet 的生成器，在几次下采样/上采样操作之间包含多个 Resnet 模块。
315
316    我们借鉴了 Justin Johnson 的神经风格迁移项目（https://github.com/jcjohnson/fast-neural-style）的 Torch 代码和思想。
317    """
318
319    def __init__(self, input_nc, output_nc, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False, n_blocks=6, padding_type='reflect'):
320        """构造一个基于 Resnet 的生成器
321
322        参数：
323            input_nc (int)      -- 输入图像的通道数
324            output_nc (int)     -- 输出图像的通道数
325            ngf (int)           -- 最后一个卷积层中的滤波器数量
326            norm_layer          -- 归一化层
327            use_dropout (bool)  -- 是否使用 dropout 层
328            n_blocks (int)      -- ResNet 模块的数量
329            padding_type (str)  -- 卷积层中填充层的名称：reflect | replicate | zero
330        """
331        assert(n_blocks >= 0)
332        super(ResnetGenerator, self).__init__()
333        if type(norm_layer) == functools.partial:
334            use_bias = norm_layer.func == nn.InstanceNorm2d
335        else:
336            use_bias = norm_layer == nn.InstanceNorm2d
337
338        model = [nn.ReflectionPad2d(3),
339                 nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0, bias=use_bias),
340                 norm_layer(ngf),
341                 nn.ReLU(True)]
342
343        n_downsampling = 2
344        for i in range(n_downsampling):  # 添加下采样层
345            mult = 2 ** i
346            model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1, bias=use_bias),
347                      norm_layer(ngf * mult * 2),
348                      nn.ReLU(True)]
349
350        mult = 2 ** n_downsampling
351        for i in range(n_blocks):       # 添加 ResNet 模块
352
353            model += [ResnetBlock(ngf * mult, padding_type=padding_type, norm_layer=norm_layer, use_dropout=use_dropout, use_bias=use_bias)]
354
355        for i in range(n_downsampling):  # 添加上采样层
356            mult = 2 ** (n_downsampling - i)
357            model += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2),
358                                         kernel_size=3, stride=2,
359                                         padding=1, output_padding=1,
360                                         bias=use_bias),
361                      norm_layer(int(ngf * mult / 2)),
362                      nn.ReLU(True)]
363        model += [nn.ReflectionPad2d(3)]
364        model += [nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0)]
365        model += [nn.Tanh()]
366
367        self.model = nn.Sequential(*model)
368
369    def forward(self, input):
370        """标准前向传播"""
371        return self.model(input)
372
373
374class ResnetBlock(nn.Module):
375    """定义一个 Resnet 模块"""
376
377    def __init__(self, dim, padding_type, norm_layer, use_dropout, use_bias):
378        """初始化 Resnet 模块
379
380        Resnet 模块是一个带有跳跃连接的卷积模块。
381        我们使用 build_conv_block 函数构建卷积模块，
382        并在 <forward> 函数中实现跳跃连接。
383        原始 Resnet 论文：https://arxiv.org/pdf/1512.03385.pdf
384        """
385        super(ResnetBlock, self).__init__()
386        self.conv_block = self.build_conv_block(dim, padding_type, norm_layer, use_dropout, use_bias)
387
388    def build_conv_block(self, dim, padding_type, norm_layer, use_dropout, use_bias):
389        """构建一个卷积模块。
390
391        参数：
392            dim (int)           -- 卷积层中的通道数。
393            padding_type (str)  -- 填充层名称：reflect | replicate | zero
394            norm_layer          -- 归一化层
395            use_dropout (bool)  -- 是否使用 dropout 层。
396            use_bias (bool)     -- 卷积层是否使用偏置
397
398        返回一个卷积模块（包含卷积层、归一化层和非线性层（ReLU））
399        """
400        conv_block = []
401        p = 0
402        if padding_type == 'reflect':
403            conv_block += [nn.ReflectionPad2d(1)]
404        elif padding_type == 'replicate':
405            conv_block += [nn.ReplicationPad2d(1)]
406        elif padding_type == 'zero':
407            p = 1
408        else:
409            raise NotImplementedError('未实现的填充方式 [%s]' % padding_type)
410
411        conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias), norm_layer(dim), nn.ReLU(True)]
412        if use_dropout:
413            conv_block += [nn.Dropout(0.5)]
414
415        p = 0
416        if padding_type == 'reflect':
417            conv_block += [nn.ReflectionPad2d(1)]
418        elif padding_type == 'replicate':
419            conv_block += [nn.ReplicationPad2d(1)]
420        elif padding_type == 'zero':
421            p = 1
422        else:
423            raise NotImplementedError('未实现的填充方式 [%s]' % padding_type)
424        conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias), norm_layer(dim)]
425
426        return nn.Sequential(*conv_block)
427
428    def forward(self, x):
429        """前向传播函数（含跳跃连接）"""
430        out = x + self.conv_block(x)  # 添加跳跃连接
431        return out
432
433
434class UnetGenerator(nn.Module):
435    """创建一个基于 U-Net 的生成器"""
436
437    def __init__(self, input_nc, output_nc, num_downs, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False):
438        """构造一个 U-Net 生成器
439        参数：
440            input_nc (int)  -- 输入图像的通道数
441            output_nc (int) -- 输出图像的通道数
442            num_downs (int) -- UNet 中的下采样次数。例如，若 |num_downs| == 7，
443                               128x128 的图像在瓶颈处将变为 1x1
444            ngf (int)       -- 最后一个卷积层中的滤波器数量
445            norm_layer      -- 归一化层
446
447        我们从最内层到最外层构建 U-Net。
448        这是一个递归过程。
449        """
450        super(UnetGenerator, self).__init__()
451        # 构建 unet 结构
452        unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=None, norm_layer=norm_layer, innermost=True)  # 添加最内层
453        for i in range(num_downs - 5):          # 添加中间层，使用 ngf * 8 个滤波器
454            unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer, use_dropout=use_dropout)
455        # 逐步减少滤波器数量，从 ngf * 8 到 ngf
456        unet_block = UnetSkipConnectionBlock(ngf * 4, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
457        unet_block = UnetSkipConnectionBlock(ngf * 2, ngf * 4, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
458        unet_block = UnetSkipConnectionBlock(ngf, ngf * 2, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
459        self.model = UnetSkipConnectionBlock(output_nc, ngf, input_nc=input_nc, submodule=unet_block, outermost=True, norm_layer=norm_layer)  # 添加最外层
460
461    def forward(self, input):
462        """标准前向传播"""
463        return self.model(input)
464
465
466class UnetSkipConnectionBlock(nn.Module):
467    """定义带跳跃连接的 Unet 子模块。
468        X -------------------identity----------------------
469        |-- 下采样 -- |submodule| -- 上采样 --|
470    """
471
472    def __init__(self, outer_nc, inner_nc, input_nc=None,
473                 submodule=None, outermost=False, innermost=False, norm_layer=nn.BatchNorm2d, use_dropout=False):
474        """构造一个带跳跃连接的 Unet 子模块。
475
476        参数：
477            outer_nc (int) -- 外层卷积的滤波器数量
478            inner_nc (int) -- 内层卷积的滤波器数量
479            input_nc (int) -- 输入图像/特征的通道数
480            submodule (UnetSkipConnectionBlock) -- 之前定义的子模块
481            outermost (bool)    -- 该模块是否为最外层模块
482            innermost (bool)    -- 该模块是否为最内层模块
483            norm_layer          -- 归一化层
484            use_dropout (bool)  -- 是否使用 dropout 层。
485        """
486        super(UnetSkipConnectionBlock, self).__init__()
487        self.outermost = outermost
488        if type(norm_layer) == functools.partial:
489            use_bias = norm_layer.func == nn.InstanceNorm2d
490        else:
491            use_bias = norm_layer == nn.InstanceNorm2d
492        if input_nc is None:
493            input_nc = outer_nc
494        downconv = nn.Conv2d(input_nc, inner_nc, kernel_size=4,
495                             stride=2, padding=1, bias=use_bias)
496        downrelu = nn.LeakyReLU(0.2, True)
497        downnorm = norm_layer(inner_nc)
498        uprelu = nn.ReLU(True)
499        upnorm = norm_layer(outer_nc)
500
501        if outermost:
502            upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,
503                                        kernel_size=4, stride=2,
504                                        padding=1)
505            down = [downconv]
506            up = [uprelu, upconv, nn.Tanh()]
507            model = down + [submodule] + up
508        elif innermost:
509            upconv = nn.ConvTranspose2d(inner_nc, outer_nc,
510                                        kernel_size=4, stride=2,
511                                        padding=1, bias=use_bias)
512            down = [downrelu, downconv]
513            up = [uprelu, upconv, upnorm]
514            model = down + up
515        else:
516            upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,
517                                        kernel_size=4, stride=2,
518                                        padding=1, bias=use_bias)
519            down = [downrelu, downconv, downnorm]
520            up = [uprelu, upconv, upnorm]
521
522            if use_dropout:
523                model = down + [submodule] + up + [nn.Dropout(0.5)]
524            else:
525                model = down + [submodule] + up
526
527        self.model = nn.Sequential(*model)
528
529    def forward(self, x):
530        if self.outermost:
531            return self.model(x)
532        else:   # 添加跳跃连接
533            return torch.cat([x, self.model(x)], 1)
534
535
536class NLayerDiscriminator(nn.Module):
537    """定义一个 PatchGAN 判别器"""
538
539    def __init__(self, input_nc, ndf=64, n_layers=3, norm_layer=nn.BatchNorm2d):
540        """构造一个 PatchGAN 判别器
541
542        参数：
543            input_nc (int)  -- 输入图像的通道数
544            ndf (int)       -- 最后一个卷积层中的滤波器数量
545            n_layers (int)  -- 判别器中卷积层的数量
546            norm_layer      -- 归一化层
547        """
548        super(NLayerDiscriminator, self).__init__()
549        if type(norm_layer) == functools.partial:  # BatchNorm2d 有仿射参数，因此无需使用偏置
550            use_bias = norm_layer.func == nn.InstanceNorm2d
551        else:
552            use_bias = norm_layer == nn.InstanceNorm2d
553
554        kw = 4
555        padw = 1
556        sequence = [nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw), nn.LeakyReLU(0.2, True)]
557        nf_mult = 1
558        nf_mult_prev = 1
559        for n in range(1, n_layers):  # 逐步增加滤波器数量
560            nf_mult_prev = nf_mult
561            nf_mult = min(2 ** n, 8)
562            sequence += [
563                nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=2, padding=padw, bias=use_bias),
564                norm_layer(ndf * nf_mult),
565                nn.LeakyReLU(0.2, True)
566            ]
567
568        nf_mult_prev = nf_mult
569        nf_mult = min(2 ** n_layers, 8)
570        sequence += [
571            nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=1, padding=padw, bias=use_bias),
572            norm_layer(ndf * nf_mult),
573            nn.LeakyReLU(0.2, True)
574        ]
575
576        sequence += [nn.Conv2d(ndf * nf_mult, 1, kernel_size=kw, stride=1, padding=padw)]  # 输出单通道预测图
577        self.model = nn.Sequential(*sequence)
578
579    def forward(self, input):
580        """标准前向传播。"""
581        return self.model(input)
582
583
584class PixelDiscriminator(nn.Module):
585    """定义一个 1x1 PatchGAN 判别器（pixelGAN）"""
586
587    def __init__(self, input_nc, ndf=64, norm_layer=nn.BatchNorm2d):
588        """构造一个 1x1 PatchGAN 判别器
589
590        参数：
591            input_nc (int)  -- 输入图像的通道数
592            ndf (int)       -- 最后一个卷积层中的滤波器数量
593            norm_layer      -- 归一化层
594        """
595        super(PixelDiscriminator, self).__init__()
596        if type(norm_layer) == functools.partial:  # BatchNorm2d 有仿射参数，因此无需使用偏置
597            use_bias = norm_layer.func == nn.InstanceNorm2d
598        else:
599            use_bias = norm_layer == nn.InstanceNorm2d
600
601        self.net = [
602            nn.Conv2d(input_nc, ndf, kernel_size=1, stride=1, padding=0),
603            nn.LeakyReLU(0.2, True),
604            nn.Conv2d(ndf, ndf * 2, kernel_size=1, stride=1, padding=0, bias=use_bias),
605            norm_layer(ndf * 2),
606            nn.LeakyReLU(0.2, True),
607            nn.Conv2d(ndf * 2, 1, kernel_size=1, stride=1, padding=0, bias=use_bias)]
608
609        self.net = nn.Sequential(*self.net)
610
611    def forward(self, input):
612        """标准前向传播。"""
613        return self.net(input)

1class SomeModel(nn.Module):
2    def __init__(self):
3        super(SomeModel, self).__init__()
4        sequence = [layer1, layer2, ...]
5        self.model = nn.Sequential(*sequence)
6    def forward(self, input):
7        return self.model(input)

1def forward(self, input):
2    x = input
3    for i in range(len(self.model)):
4        print(x.shape)
5        x = self.model[i](x)
6        print(x.shape)
7    return x

input_patch_size = k + s ⋅ (output_patch_size - 1)