迁移学习——CycleGAN

CycleGAN

CycleGAN(循环一致性对抗网络),用于实现两个域(例如,风格或主题不同的图像)之间的无监督图像到图像转换。
CycleGAN的核心思想是使用生成器(Generator)和判别器(Discriminator)来学习从源域(source

domain)到目标域(target domain)的映射,同时保持循环一致性,即从目标域映射回源域应该尽可能接近原始源域图像。

1.导入需要的包

from random import randint: 从Python的random模块中导入randint函数,用于生成随机整数。

import numpy as np: 导入Numpy库,并将其重命名为np,以便在代码中使用。
import torch:导入PyTorch库。
torch.set_default_tensor_type(torch.FloatTensor):设置PyTorch的默认Tensor类型为torch.FloatTensor。
import torch.nn as nn:导入PyTorch的神经网络模块,并将其重命名为nn。
import torch.optim as optim:导入PyTorch的优化器模块,并将其重命名为optim。
import torchvision.datasets as datasets: 导入PyTorch的图像数据集模块,并将其重命名为datasets。
import torchvision.transforms as transforms:导入PyTorch的图像变换模块,并将其重命名为transforms。
import os:导入Python的操作系统模块,用于处理文件和目录。
import matplotlib.pyplot as plt:导入matplotlib的Pyplot模块,用于绘图。
import torch.nn.functional as F:导入PyTorch的函数模块,并将其重命名为F。
from torch.autograd import Variable:从PyTorch的自动求导模块中导入Variable类。
from torchvision.utils import save_image: 从PyTorch的图像处理模块中导入save_image函数。
import shutil:导入Python的文件操作模块,用于删除文件和目录。
import cv2: 导入OpenCV库,用于图像处理和计算机视觉。
import random: 导入Python的随机模块。
from PIL import Image:从Pillow库中导入Image类。
import itertools: 导入Python的迭代工具模块。

python 复制代码
from random import randint
import numpy as np 
import torch
torch.set_default_tensor_type(torch.FloatTensor)
import torch.nn as nn
import torch.optim as optim
import torchvision.datasets as datasets
import torchvision.transforms as transforms
import os
import matplotlib.pyplot as plt
import torch.nn.functional as F
from torch.autograd import Variable
from torchvision.utils import save_image
import shutil
import cv2
import random
from PIL import Image
import itertools

2.数据加载

(1)to_img 函数

out = 0.5 * (x + 1): 将输入张量 x 的值从 [-1, 1] 范围转换到 [0, 1] 范围。这是因为在训练过程中,图像通常会被归一化到 [-1, 1] 范围,而显示图像时需要将其转换回 [0, 1] 范围。
out = out.clamp(0, 1): 确保所有像素值都在 [0, 1] 范围内。clamp 函数将小于0的值设为0,大于1的值设为1。
out = out.view(-1, 3, 256, 256): 将张量 out 的形状重新调整为批次的形状,其中每个样本是一个 3通道(RGB)的 256x256 图像。-1 表示自动计算批次大小。

python 复制代码
def to_img(x):
    out = 0.5 * (x + 1)
    out = out.clamp(0, 1)  
    out = out.view(-1, 3, 256, 256)  
    return out

(2)数据加载

data_path = os.path.abspath('D:\probject\pythonProject1\pytorch\CycleGAN\data'):定义了数据的路径,使用os.path.abspath()将相对路径转换为绝对路径。
image_size = 256:指定图像的大小为256x256。
batch_size = 1:定义了批处理的大小为1。

python 复制代码
data_path = os.path.abspath('D:\probject\pythonProject1\pytorch\CycleGAN\data')
image_size = 256
batch_size = 1

(3)图像转换

transform = transforms.Compose([: 创建一个由多个图像转换操作组成的管道。
transforms.Resize(int(image_size * 1.12), Image.BICUBIC): 将图像大小调整为原始大小的 1.12 倍。这样做是为了在后续的随机裁剪中提供更多的裁剪选择。
transforms.RandomCrop(image_size): 从调整大小后的图像中随机裁剪出 256x256 像素大小的区域。
transforms.RandomHorizontalFlip(): 以 50% 的概率随机水平翻转图像。
transforms.ToTensor(): 将 PIL 图像转换为 PyTorch 张量。
transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5)):对图像进行归一化处理,将每个通道的像素值从 [0, 1] 范围转换为 [-1, 1] 范围。

python 复制代码
transform = transforms.Compose([transforms.Resize(int(image_size * 1.12), 
                                                  Image.BICUBIC), 
            transforms.RandomCrop(image_size), 
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))])

3.随机读取图像进行预处理

(1)函数参数

batch_size: 一个整数,表示每个批次中图像的数量。默认值为1。

python 复制代码
def _get_train_data(batch_size=1):

(2)数据路径

train_a_filepath: 训练集A的文件路径。
train_b_filepath: 训练集B的文件路径。

python 复制代码
	train_a_filepath = data_path + '\\trainA\\'
    train_b_filepath = data_path + '\\trainB\\'

(3)读取文件列表

train_a_list: 读取训练集A目录中的所有文件名。
train_b_list: 读取训练集B目录中的所有文件名。

python 复制代码
   train_a_list = os.listdir(train_a_filepath)
   train_b_list = os.listdir(train_b_filepath)

(4)初始化结果列表

train_a_result: 存储处理后的训练集A图像。
train_b_result: 存储处理后的训练集B图像。

python 复制代码
    train_a_result = []
    train_b_result = [] 

(5)随机采样

numlist: 从0到训练集A长度之间的范围中随机采样 batch_size 个索引。

python 复制代码
numlist = random.sample(range(0, len(train_a_list)), batch_size)

(6)读取和预处理图像

对于 numlist 中的每个索引 i: 读取训练集A和B中对应的文件名。 使用 PIL.Image.open

打开图像文件,并将其转换为RGB格式。 应用之前定义的 transform 方法对图像进行预处理(包括调整大小、裁剪、翻转和归一化)。

将预处理后的图像添加到 train_a_result 和 train_b_result 列表中。

python 复制代码
	for i in numlist:
        a_filename = train_a_list[i]
        a_img = Image.open(train_a_filepath + a_filename).convert('RGB')
        res_a_img = transform(a_img)
        train_a_result.append(torch.unsqueeze(res_a_img, 0))
        
        b_filename = train_b_list[i]
        b_img = Image.open(train_b_filepath + b_filename).convert('RGB')
        res_b_img = transform(b_img)
        train_b_result.append(torch.unsqueeze(res_b_img, 0))
        

(7)返回结果

使用 torch.cattrain_a_resulttrain_b_result

列表中的图像堆叠成一个批次,并返回这两个批次的图像。

4.残差网络块

残差块是一种常用的构建块,用于深度卷积神经网络,特别是在
ResNet(残差网络)架构中。它允许网络在学习过程中保留和利用之前层的信息,通过跳跃连接(shortcut

connections)来解决深层网络训练过程中的梯度消失问题。

(1)构造函数

def __init__(self, in_features): 构造函数接收一个参数 in_features,表示输入特征图的通道数。
super(ResidualBlock, self).__init__(): 调用父类 nn.Module 的构造函数。
self.block_layer: 定义一个顺序模型 nn.Sequential,包含残差块的所有层。

python 复制代码
class ResidualBlock(nn.Module):
    def __init__(self, in_features):
        super(ResidualBlock, self).__init__()
        self.block_layer = nn.Sequential

(2)残差块层

nn.ReflectionPad2d(1):使用反射填充(padding)来扩展输入张量的边界。这种填充方式在边缘反射输入数据,以保持边缘信息的连续性。
nn.Conv2d(in_features, in_features, 3): 使用 3x3的卷积核进行卷积操作,输入和输出通道数相同。
nn.InstanceNorm2d(in_features):应用实例归一化(Instance Normalization)来对每个样本的特征图进行归一化处理。这与批量归一化(Batch Normalization)不同,它不对整个批次的数据进行归一化,而是对单个样本的特征图进行归一化。
nn.ReLU(inplace=True): 应用 ReLU 激活函数,并设置 inplace=True以便直接修改输入张量,减少内存使用。

python 复制代码
(
            nn.ReflectionPad2d(1),
            nn.Conv2d(in_features, in_features, 3),
            nn.InstanceNorm2d(in_features),
            nn.ReLU(inplace=True),
            nn.ReflectionPad2d(1),
            nn.Conv2d(in_features, in_features, 3),
            nn.InstanceNorm2d(in_features))

(3)跳跃连接

return x + self.block_layer(x): 这是残差块的核心,它将输入张量 x 与

self.block_layer(x) 的输出相加,形成跳跃连接。这样,即使 self.block_layer

的输出为零(即网络未能学习到任何东西),输入 x 仍然可以通过跳跃连接直接传递到下一层,从而保持了信息的流通。

python 复制代码
	def forward(self, x):
        return x + self.block_layer(x)

5.生成器网络

生成器的目的是将输入图像从一个域转换到另一个域。

(1)构造函数

super(Generator, self).__init__(): 调用父类 nn.Module 的构造函数。
model: 初始化一个列表,用于存储生成器网络中的层。

python 复制代码
class Generator(nn.Module):
    def __init__(self):
        super(Generator, self).__init__()

(2)编码器部分

nn.ReflectionPad2d(3): 使用反射填充(padding)来扩展输入张量的边界。
nn.Conv2d(3, 64, 7): 使用 7x7 的卷积核将输入图像(3 通道)转换为 64 通道的特征图。
nn.InstanceNorm2d(64):应用实例归一化。
nn.ReLU(inplace=True): 应用 ReLU 激活函数。
for _ in range(2):重复以下层两次,以逐渐减少特征图的尺寸。
nn.Conv2d(in_features, out_features, 3,stride=2, padding=1): 使用 3x3 的卷积核,步长为 2,进行降采样。
nn.InstanceNorm2d(out_features): 应用实例归一化。
nn.ReLU(inplace=True):应用 ReLU 激活函数。

python 复制代码
model = [nn.ReflectionPad2d(3), 
                 nn.Conv2d(3, 64, 7), 
                 nn.InstanceNorm2d(64), 
                 nn.ReLU(inplace=True)]

        in_features = 64
        out_features = in_features * 2
        for _ in range(2):
            model += [nn.Conv2d(in_features, out_features, 
                                3, stride=2, padding=1), 
            nn.InstanceNorm2d(out_features), 
            nn.ReLU(inplace=True)]
            in_features = out_features
            out_features = in_features*2

(3)残差块部分

for _ in range(9): 重复添加 9 个残差块,这些块是 CycleGAN 生成器的核心,用于学习域之间的映射。

python 复制代码
 	for _ in range(9):
            model += [ResidualBlock(in_features)]

(4)解码器部分

out_features = in_features // 2: 准备进行上采样,将特征图的尺寸加倍。
for _ in range(2): 重复以下层两次,以逐渐增加特征图的尺寸。
nn.ConvTranspose2d(in_features, out_features, 3, stride=2, padding=1, output_padding=1): 使用 3x3 的转置卷积核,步长为 2,进行上采样。
nn.InstanceNorm2d(out_features): 应用实例归一化。
nn.ReLU(inplace=True): 应用 ReLU 激活函数。

python 复制代码
    out_features = in_features // 2
        for _ in range(2):
            model += [nn.ConvTranspose2d(
                    in_features, out_features, 
                    3, stride=2, padding=1, output_padding=1), 
                nn.InstanceNorm2d(out_features), 
                nn.ReLU(inplace=True)]
            in_features = out_features
            out_features = in_features // 2

(5)输出层

nn.ReflectionPad2d(3): 使用反射填充。
nn.Conv2d(64, 3, 7): 使用 7x7的卷积核将特征图转换回 3 通道的图像。
nn.Tanh(): 应用 Tanh 激活函数,将输出值范围映射到 [-1, 1]。

python 复制代码
	model += [nn.ReflectionPad2d(3), 
                  nn.Conv2d(64, 3, 7), 
                  nn.Tanh()]

(6)模型初始化

self.gen = nn.Sequential( * model): 将所有层组合成一个顺序模型。

python 复制代码
self.gen = nn.Sequential( * model)

(7)前向传播

def forward(self, x): 定义前向传播函数。
x = self.gen(x): 通过生成器网络传递输入 x。
return x: 返回生成器的输出。

python 复制代码
	def forward(self, x):
        x = self.gen(x)
        return x 

6.判别器网络

(1)构造函数

super(Discriminator, self).__init__(): 调用父类 nn.Module 的构造函数。
self.dis: 定义一个顺序模型 nn.Sequential,包含判别器网络的所有层。

python 复制代码
class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()
        self.dis = nn.Sequential

(2)判别器层

nn.Conv2d(3, 64, 4, 2, 1, bias=False): 使用 4x4 的卷积核,步长为2,进行降采样,输入通道数为 3(RGB),输出通道数为 64。
nn.LeakyReLU(0.2, inplace=True): 应用Leaky ReLU 激活函数,设置斜率为 0.2。
for _ in range(3): 重复以下层三次,以逐渐减少特征图的尺寸。
nn.Conv2d(in_features, out_features, 4, 2, 1, bias=False): 使用 4x4 的卷积核,步长为 2,进行降采样。
nn.InstanceNorm2d(out_features): 应用实例归一化。
nn.LeakyReLU(0.2, inplace=True): 应用 Leaky ReLU 激活函数。

python 复制代码
(
            nn.Conv2d(3, 64, 4, 2, 1, bias=False),
            nn.LeakyReLU(0.2, inplace=True),

            nn.Conv2d(64, 128, 4, 2, 1, bias=False),
            nn.InstanceNorm2d(128),
            nn.LeakyReLU(0.2, inplace=True),

            nn.Conv2d(128, 256, 4, 2, 1, bias=False),
            nn.InstanceNorm2d(256),
            nn.LeakyReLU(0.2, inplace=True),

(3)全卷积网络部分

nn.Conv2d(256, 512, 4, padding=1): 使用 4x4 的卷积核,不进行降采样,输入通道数为256,输出通道数为 512。
nn.InstanceNorm2d(512): 应用实例归一化。
nn.LeakyReLU(0.2, inplace=True): 应用 Leaky ReLU 激活函数。
nn.Conv2d(512, 1, 4, padding=1):使用 4x4 的卷积核,不进行降采样,输入通道数为 512,输出通道数为 1。

python 复制代码
			nn.Conv2d(256, 512, 4, padding=1),
            nn.InstanceNorm2d(512),
            nn.LeakyReLU(0.2, inplace=True),
            
            nn.Conv2d(512, 1, 4, padding=1))  

(4)输出

return F.avg_pool2d(x, x.size()[2:]).view(x.size()[0], -1):对判别器输出的特征图进行平均池化操作,然后将其展平为一维向量。这个一维向量将作为最终的判别结果,其长度为 1,表示输入图像的真实性(接近 1表示真实,接近 0 表示假)。

python 复制代码
	def forward(self, x):
        x = self.dis(x)
        return F.avg_pool2d(x, x.size()[2:]).view(x.size()[0], -1)

7.缓存生成器

(1)构造函数

def __init__(self, max_size=50): 定义了一个构造函数 init ,用于在创建ReplayBuffer 对象时初始化其属性。
self.max_size = max_size: 初始化缓冲区的大小。
self.data = []: 初始化一个空列表 self.data,用于存储缓存的数据。

python 复制代码
class ReplayBuffer():
#     """
#     缓存队列,若不足则新增,否则随机替换
#     """
    def __init__(self, max_size=50):
        self.max_size = max_size
        self.data = []

(2)push_and_pop 方法

def push_and_pop(self, data): 定义了一个方法,用于将新数据推入缓冲区,并在需要时弹出旧数据。
to_return = []: 初始化一个空列表 to_return,用于存储从缓冲区中弹出的数据。
for element in data.data:: 遍历传入的数据 data.data 中的每个元素。
element = torch.unsqueeze(element, 0):将每个元素展平为一维张量。这通常是为了确保张量的形状与预期的形状匹配,以便后续的操作可以正确执行。
if len(self.data) < self.max_size:: 如果缓冲区中还没有达到最大容量,则将新元素添加到缓冲区。
self.data.append(element): 将新元素添加到缓冲区。
to_return.append(element): 将新元素添加到 to_return 列表中。
else:: 如果缓冲区已满,则随机替换缓冲区中的一个元素。
if random.uniform(0,1) > 0.5:: 如果随机数大于 0.5,则从缓冲区中随机选择一个元素替换。
i = random.randint(0, self.max_size-1): 随机选择一个索引。
to_return.append(self.data[i].clone()): 将缓冲区中的元素复制并添加到 to_return列表中。
self.data[i] = element: 用新元素替换缓冲区中的元素。
else:: 如果随机数小于或等于 0.5,则直接添加新元素到 to_return 列表中。
to_return.append(element): 将新元素添加到 to_return 列表中。
return Variable(torch.cat(to_return)): 返回 to_return 列表中所有元素的拼接张量。Variable 是一个 PyTorch 类,用于表示可变的张量。torch.cat 函数用于将多个张量拼接在一起。

python 复制代码
def push_and_pop(self, data):
        to_return = []
        for element in data.data:
            element = torch.unsqueeze(element, 0)
            if len(self.data) < self.max_size:
                self.data.append(element)
                to_return.append(element)
            else:
                if random.uniform(0,1) > 0.5:
                    i = random.randint(0, self.max_size-1)
                    to_return.append(self.data[i].clone())
                    self.data[i] = element
                else:
                    to_return.append(element)
        return Variable(torch.cat(to_return))

8.训练生成对抗网络(GAN)

fake_A_buffer = ReplayBuffer(): 创建了一个名为 fake_A_buffer 的 ReplayBuffer实例。ReplayBuffer是一个用于缓存和随机替换数据的结构,在训练循环中用于缓存生成器生成的假图像,以便在后续的训练步骤中用于训练判别器。
fake_B_buffer = ReplayBuffer(): 创建了一个名为 fake_B_buffer 的 ReplayBuffer实例。这个缓冲区的作用与 fake_A_buffer 类似,用于缓存从生成器 netG_B2A 生成的假图像。

python 复制代码
fake_A_buffer = ReplayBuffer()
fake_B_buffer = ReplayBuffer()

netG_A2B = Generator(): 创建了一个名为 netG_A2B 的 Generator 实例。Generator是一个用于生成新图像的神经网络,在这里,它将从域 A 生成域 B 的图像。
netG_B2A = Generator(): 创建了一个名为 netG_B2A 的 Generator 实例。这个生成器将从域 B生成域 A 的图像。
netD_A = Discriminator(): 创建了一个名为 netD_A 的 Discriminator实例。Discriminator 是一个用于判断图像是否真实的神经网络,在这里,它用于判断 A 类图像是否真实。
netD_B = Discriminator(): 创建了一个名为 netD_B 的 Discriminator实例。这个判别器用于判断 B 类图像是否真实。

python 复制代码
netG_A2B = Generator()
netG_B2A = Generator()
netD_A = Discriminator()
netD_B = Discriminator()

criterion_GAN = torch.nn.MSELoss(): 定义了一个名为 criterion_GAN 的 MSELoss

损失函数。这个损失函数用于计算 GAN 损失,即判别器对真实图像和假图像的预测之间的差异。
criterion_cycle = torch.nn.L1Loss(): 定义了一个名为 criterion_cycle 的 L1Loss损失函数。这个损失函数用于计算循环一致性损失,即生成器生成的图像与其输入图像之间的差异。
criterion_identity = torch.nn.L1Loss(): 定义了一个名为 criterion_identity 的 L1Loss损失函数。这个损失函数用于计算身份损失,即生成器生成的图像与其输入图像之间的差异。

python 复制代码
criterion_GAN = torch.nn.MSELoss()
criterion_cycle = torch.nn.L1Loss()
criterion_identity = torch.nn.L1Loss()

d_learning_rate = 3e-4 : 定义了判别器的学习率。
g_learning_rate = 3e-4:定义了生成器的 learning rate。
optim_betas = (0.5, 0.999): 定义了优化器的超参数betas,这是用于计算梯度下降的动量项的值。

python 复制代码
d_learning_rate = 3e-4  
g_learning_rate = 3e-4
optim_betas = (0.5, 0.999)

9.优化器

g_optimizer = optim.Adam(itertools.chain(netG_A2B.parameters(), netG_B2A.parameters()), lr=d_learning_rate): 创建了一个名为 g_optimizer 的Adam 优化器实例。Adam 是一种常用的优化算法,用于调整神经网络的权重。这里,itertools.chain函数用于将两个生成器的参数合并为一个单一的迭代器,以便于一起优化。lr 参数指定了学习率,它用于控制权重更新的速度。

da_optimizer = optim.Adam(netD_A.parameters(), lr=d_learning_rate):创建了一个名为 da_optimizer 的 Adam 优化器实例,用于训练判别器 netD_A。

db_optimizer = optim.Adam(netD_B.parameters(), lr=d_learning_rate):创建了一个名为 db_optimizer 的 Adam 优化器实例,用于训练判别器 netD_B。

python 复制代码
g_optimizer = optim.Adam(itertools.chain(netG_A2B.parameters(), 
                                         netG_B2A.parameters()), 
            lr=d_learning_rate)
da_optimizer = optim.Adam(netD_A.parameters(), lr=d_learning_rate)
db_optimizer = optim.Adam(netD_B.parameters(), lr=d_learning_rate)

10.训练循环的迭代次数

num_epochs = 100: 定义了训练循环的迭代次数。epoch是一个训练周期,在这个周期内,所有数据都会被遍历一次。在这里,训练循环将执行 100 个周期。

python 复制代码
num_epochs = 100

11.训练循环

for epoch in range(num_epochs):: 开始一个循环,该循环将执行指定的次数(由 num_epochs定义)。
real_a, real_b = _get_train_data(batch_size): 从数据集中获取一批真实图像real_a 和 real_b。
target_real = torch.full((batch_size,), 1).float():创建一个全为 1 的张量 target_real,用于指示真实图像。
target_fake =torch.full((batch_size,), 0).float(): 创建一个全为 0 的张量target_fake,用于指示假图像。
g_optimizer.zero_grad():清除生成器的梯度,以便于下一次前向传播和反向传播时不会累积梯度。

python 复制代码
for epoch in range(num_epochs): 

    real_a, real_b = _get_train_data(batch_size)
    target_real = torch.full((batch_size,), 1).float()
    target_fake = torch.full((batch_size,), 0).float()
    
	g_optimizer.zero_grad()

12.训练生成器

same_B = netG_A2B(real_b).float(): 使用生成器 netG_A2B 从真实图像 real_b生成相似的图像 same_B。

loss_identity_B = criterion_identity(same_B, real_b) * 5.0: 计算same_B 和 real_b 之间的身份损失,并乘以 5.0 以增加其权重。

same_A = netG_B2A(real_a).float(): 使用生成器 netG_B2A 从真实图像 real_a生成相似的图像 same_A。

loss_identity_A = criterion_identity(same_A, real_a) * 5.0: 计算same_A 和 real_a 之间的身份损失,并乘以 5.0 以增加其权重。

fake_B = netG_A2B(real_a).float(): 使用生成器 netG_A2B 从真实图像 real_a 生成假图像fake_B。

pred_fake = netD_B(fake_B).float(): 使用判别器 netD_B 判断 fake_B 是否为假图像。

loss_GAN_A2B = criterion_GAN(pred_fake, target_real): 计算判别器对 fake_B的预测和真实图像的损失,即 GAN 损失。

fake_A = netG_B2A(real_b).float(): 使用生成器 netG_B2A 从真实图像 real_b 生成假图像fake_A。

pred_fake = netD_A(fake_A).float(): 使用判别器 netD_A 判断 fake_A 是否为假图像。

loss_GAN_B2A = criterion_GAN(pred_fake, target_real): 计算判别器对 fake_A的预测和真实图像的损失,即 GAN 损失。

recovered_A = netG_B2A(fake_B).float(): 使用生成器 netG_B2A 从假图像 fake_B生成恢复的图像 recovered_A。

loss_cycle_ABA = criterion_cycle(recovered_A, real_a) * 10.0: 计算recovered_A 和 real_a 之间的循环一致性损失,并乘以 10.0 以增加其权重。

recovered_B = netG_A2B(fake_A).float(): 使用生成器 netG_A2B 从假图像 fake_A生成恢复的图像 recovered_B。

loss_cycle_BAB = criterion_cycle(recovered_B, real_b) * 10.0: 计算recovered_B 和 real_b 之间的循环一致性损失,并乘以 10.0 以增加其权重。

loss_G = (loss_identity_A + loss_identity_B + loss_GAN_A2B + loss_GAN_B2A + loss_cycle_ABA + loss_cycle_BAB): 将所有损失加在一起,得到生成器的总损失。

loss_G.backward(): 对总损失进行反向传播,计算每个参数的梯度。

g_optimizer.step():会对生成器的所有参数进行梯度更新,以最小化生成器损失函数。

python 复制代码
# 第一步:训练生成器
    same_B = netG_A2B(real_b).float()
    loss_identity_B = criterion_identity(same_B, real_b) * 5.0   
    same_A = netG_B2A(real_a).float()
    loss_identity_A = criterion_identity(same_A, real_a) * 5.0
    
    fake_B = netG_A2B(real_a).float()
    pred_fake = netD_B(fake_B).float()
    loss_GAN_A2B = criterion_GAN(pred_fake, target_real)
    fake_A = netG_B2A(real_b).float()
    pred_fake = netD_A(fake_A).float()
    loss_GAN_B2A = criterion_GAN(pred_fake, target_real)
    recovered_A = netG_B2A(fake_B).float()
    loss_cycle_ABA = criterion_cycle(recovered_A, real_a) * 10.0
    recovered_B = netG_A2B(fake_A).float()
    loss_cycle_BAB = criterion_cycle(recovered_B, real_b) * 10.0  
    loss_G = (loss_identity_A + loss_identity_B + loss_GAN_A2B + 
              loss_GAN_B2A + loss_cycle_ABA + loss_cycle_BAB)
    loss_G.backward()    
    g_optimizer.step()

13.训练判别器

da_optimizer.zero_grad(): 清除判别器 A 的梯度,以便于下一次前向传播和反向传播时不会累积梯度。

pred_real = netD_A(real_a).float(): 使用判别器 A 来判断真实图像 real_a 是否为真实图像。

loss_D_real = criterion_GAN(pred_real, target_real): 计算判别器 A对真实图像的预测和真实图像的损失,即 GAN 损失。

fake_A = fake_A_buffer.push_and_pop(fake_A): 从 fake_A_buffer 中获取一批fake_A 图像,这些图像是从生成器 A 生成的假图像。

pred_fake = netD_A(fake_A.detach()).float(): 使用判别器 A 来判断 fake_A是否为假图像。由于 fake_A 是从 fake_A_buffer 中获取的,它已经与生成器的梯度解耦,因此不需要梯度信息。

loss_D_fake = criterion_GAN(pred_fake, target_fake): 计算判别器 A 对fake_A 的预测和假图像的损失,即 GAN 损失。

loss_D_A = (loss_D_real + loss_D_fake) * 0.5: 将判别器 A的真实图像损失和假图像损失加在一起,得到判别器 A 的总损失。

loss_D_A.backward(): 对判别器 A 的总损失进行反向传播,计算每个参数的梯度。
da_optimizer.step(): 使用之前计算的梯度来更新判别器 A 的参数。

python 复制代码
   # 第二步:训练判别器
    # 训练判别器A
    da_optimizer.zero_grad()
    pred_real = netD_A(real_a).float()
    loss_D_real = criterion_GAN(pred_real, target_real)
    fake_A = fake_A_buffer.push_and_pop(fake_A)
    pred_fake = netD_A(fake_A.detach()).float()
    loss_D_fake = criterion_GAN(pred_fake, target_fake)
    loss_D_A = (loss_D_real + loss_D_fake) * 0.5
    loss_D_A.backward()
    da_optimizer.step()
    # 训练判别器B
    db_optimizer.zero_grad()
    pred_real = netD_B(real_b)
    loss_D_real = criterion_GAN(pred_real, target_real)
    fake_B = fake_B_buffer.push_and_pop(fake_B)
    pred_fake = netD_B(fake_B.detach())
    loss_D_fake = criterion_GAN(pred_fake, target_fake)
    loss_D_B = (loss_D_real + loss_D_fake) * 0.5
    loss_D_B.backward()
	db_optimizer.step()

14.损失打印,存储伪造图片

print('Epoch[{}],loss_G:{:.6f} ,loss_D_A:{:.6f},loss_D_B:{:.6f}' .format(epoch, loss_G.data.item(), loss_D_A.data.item(), loss_D_B.data.item())):打印当前训练周期(epoch)的损失,包括生成器损失(loss_G)和两个判别器损失(loss_D_A 和 loss_D_B)。
if (epoch + 1) % 20 == 0 or epoch == 0:: 检查当前训练周期是否是 20的倍数,或者是否是第一个周期。如果是,则执行以下操作。
b_fake = to_img(fake_B.data): 将判别器 B的输入(fake_B)转换回图像格式。
a_fake = to_img(fake_A.data): 将判别器 A的输入(fake_A)转换回图像格式。
a_real = to_img(real_a.data): 将真实图像 A 转换回图像格式。
b_real = to_img(real_b.data): 将真实图像 B 转换回图像格式。
save_image(a_fake,'../tmp/a_fake.png'): 将 a_fake 图像保存到文件 .../tmp/a_fake.png。
save_image(b_fake, '../tmp/b_fake.png'): 将 b_fake 图像保存到文件.../tmp/b_fake.png。
save_image(a_real, '../tmp/a_real.png'): 将 a_real图像保存到文件 .../tmp/a_real.png。
save_image(b_real, '../tmp/b_real.png'):将 b_real 图像保存到文件 .../tmp/b_real.png。

python 复制代码
 #损失打印,存储伪造图片
    print('Epoch[{}],loss_G:{:.6f} ,loss_D_A:{:.6f},loss_D_B:{:.6f}'
      .format(epoch, loss_G.data.item(), loss_D_A.data.item(), 
              loss_D_B.data.item()))
    if (epoch + 1) % 20 == 0 or epoch == 0:  
        b_fake = to_img(fake_B.data)
        a_fake = to_img(fake_A.data)
        a_real = to_img(real_a.data)
        b_real = to_img(real_b.data)
        save_image(a_fake, '../tmp/a_fake.png') 
        save_image(b_fake, '../tmp/b_fake.png') 
        save_image(a_real, '../tmp/a_real.png') 
        save_image(b_real, '../tmp/b_real.png') 


全部代码

python 复制代码
from random import randint
import numpy as np 
import torch
torch.set_default_tensor_type(torch.FloatTensor)
import torch.nn as nn
import torch.optim as optim
import torchvision.datasets as datasets
import torchvision.transforms as transforms
import os
import matplotlib.pyplot as plt
import torch.nn.functional as F
from torch.autograd import Variable
from torchvision.utils import save_image
import shutil
import cv2
import random
from PIL import Image
import itertools   
python 复制代码
 def to_img(x):
    out = 0.5 * (x + 1)
    out = out.clamp(0, 1)  
    out = out.view(-1, 3, 256, 256)  
    return out

# 数据加载 
data_path = os.path.abspath('D:\probject\pythonProject1\pytorch\CycleGAN\data')
image_size = 256
batch_size = 1

transform = transforms.Compose([transforms.Resize(int(image_size * 1.12), 
                                                  Image.BICUBIC), 
            transforms.RandomCrop(image_size), 
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))])
python 复制代码
def _get_train_data(batch_size=1):
    
    train_a_filepath = data_path + '\\trainA\\'
    train_b_filepath = data_path + '\\trainB\\'
    
    train_a_list = os.listdir(train_a_filepath)
    train_b_list = os.listdir(train_b_filepath)
    
    train_a_result = []
    train_b_result = [] 
    
    numlist = random.sample(range(0, len(train_a_list)), batch_size)
    
    for i in numlist:
        a_filename = train_a_list[i]
        a_img = Image.open(train_a_filepath + a_filename).convert('RGB')
        res_a_img = transform(a_img)
        train_a_result.append(torch.unsqueeze(res_a_img, 0))
        
        b_filename = train_b_list[i]
        b_img = Image.open(train_b_filepath + b_filename).convert('RGB')
        res_b_img = transform(b_img)
        train_b_result.append(torch.unsqueeze(res_b_img, 0))
        
    return torch.cat(train_a_result, dim=0), torch.cat(train_b_result, dim=0)

# """
# 残差网络block
python 复制代码
class ResidualBlock(nn.Module):
    def __init__(self, in_features):
        super(ResidualBlock, self).__init__()
        self.block_layer = nn.Sequential(
            nn.ReflectionPad2d(1),
            nn.Conv2d(in_features, in_features, 3),
            nn.InstanceNorm2d(in_features),
            nn.ReLU(inplace=True),
            nn.ReflectionPad2d(1),
            nn.Conv2d(in_features, in_features, 3),
            nn.InstanceNorm2d(in_features))
        
    def forward(self, x):
        return x + self.block_layer(x)
python 复制代码
# 生成器
class Generator(nn.Module):
    def __init__(self):
        super(Generator, self).__init__()
             
        model = [nn.ReflectionPad2d(3), 
                 nn.Conv2d(3, 64, 7), 
                 nn.InstanceNorm2d(64), 
                 nn.ReLU(inplace=True)]

        in_features = 64
        out_features = in_features * 2
        for _ in range(2):
            model += [nn.Conv2d(in_features, out_features, 
                                3, stride=2, padding=1), 
            nn.InstanceNorm2d(out_features), 
            nn.ReLU(inplace=True)]
            in_features = out_features
            out_features = in_features*2

        for _ in range(9):
            model += [ResidualBlock(in_features)]

        out_features = in_features // 2
        for _ in range(2):
            model += [nn.ConvTranspose2d(
                    in_features, out_features, 
                    3, stride=2, padding=1, output_padding=1), 
                nn.InstanceNorm2d(out_features), 
                nn.ReLU(inplace=True)]
            in_features = out_features
            out_features = in_features // 2

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

        self.gen = nn.Sequential( * model)
        
    def forward(self, x):
        x = self.gen(x)
        return x 
python 复制代码
# 判别器 

class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()
        self.dis = nn.Sequential(
            nn.Conv2d(3, 64, 4, 2, 1, bias=False),
            nn.LeakyReLU(0.2, inplace=True),

            nn.Conv2d(64, 128, 4, 2, 1, bias=False),
            nn.InstanceNorm2d(128),
            nn.LeakyReLU(0.2, inplace=True),

            nn.Conv2d(128, 256, 4, 2, 1, bias=False),
            nn.InstanceNorm2d(256),
            nn.LeakyReLU(0.2, inplace=True),

            nn.Conv2d(256, 512, 4, padding=1),
            nn.InstanceNorm2d(512),
            nn.LeakyReLU(0.2, inplace=True),
            
            nn.Conv2d(512, 1, 4, padding=1))        
        
    def forward(self, x):
        x = self.dis(x)
        return F.avg_pool2d(x, x.size()[2:]).view(x.size()[0], -1)
python 复制代码
class ReplayBuffer():
#     """
#     缓存队列,若不足则新增,否则随机替换
#     """
    def __init__(self, max_size=50):
        self.max_size = max_size
        self.data = []
        
    def push_and_pop(self, data):
        to_return = []
        for element in data.data:
            element = torch.unsqueeze(element, 0)
            if len(self.data) < self.max_size:
                self.data.append(element)
                to_return.append(element)
            else:
                if random.uniform(0,1) > 0.5:
                    i = random.randint(0, self.max_size-1)
                    to_return.append(self.data[i].clone())
                    self.data[i] = element
                else:
                    to_return.append(element)
        return Variable(torch.cat(to_return))
    
python 复制代码
fake_A_buffer = ReplayBuffer()
fake_B_buffer = ReplayBuffer()

netG_A2B = Generator()
netG_B2A = Generator()
netD_A = Discriminator()
netD_B = Discriminator()

criterion_GAN = torch.nn.MSELoss()
criterion_cycle = torch.nn.L1Loss()
criterion_identity = torch.nn.L1Loss()

d_learning_rate = 3e-4  # 3e-4
g_learning_rate = 3e-4
optim_betas = (0.5, 0.999)
python 复制代码
g_optimizer = optim.Adam(itertools.chain(netG_A2B.parameters(), 
                                         netG_B2A.parameters()), 
            lr=d_learning_rate)
da_optimizer = optim.Adam(netD_A.parameters(), lr=d_learning_rate)
db_optimizer = optim.Adam(netD_B.parameters(), lr=d_learning_rate)

num_epochs = 100
python 复制代码
for epoch in range(num_epochs): 

    real_a, real_b = _get_train_data(batch_size)
    target_real = torch.full((batch_size,), 1).float()
    target_fake = torch.full((batch_size,), 0).float()
    
    g_optimizer.zero_grad()
    
    # 第一步:训练生成器
    same_B = netG_A2B(real_b).float()
    loss_identity_B = criterion_identity(same_B, real_b) * 5.0   
    same_A = netG_B2A(real_a).float()
    loss_identity_A = criterion_identity(same_A, real_a) * 5.0
    
    fake_B = netG_A2B(real_a).float()
    pred_fake = netD_B(fake_B).float()
    loss_GAN_A2B = criterion_GAN(pred_fake, target_real)
    fake_A = netG_B2A(real_b).float()
    pred_fake = netD_A(fake_A).float()
    loss_GAN_B2A = criterion_GAN(pred_fake, target_real)
    recovered_A = netG_B2A(fake_B).float()
    loss_cycle_ABA = criterion_cycle(recovered_A, real_a) * 10.0
    recovered_B = netG_A2B(fake_A).float()
    loss_cycle_BAB = criterion_cycle(recovered_B, real_b) * 10.0  
    loss_G = (loss_identity_A + loss_identity_B + loss_GAN_A2B + 
              loss_GAN_B2A + loss_cycle_ABA + loss_cycle_BAB)
    loss_G.backward()    
    g_optimizer.step()
    
    
    # 第二步:训练判别器
    # 训练判别器A
    da_optimizer.zero_grad()
    pred_real = netD_A(real_a).float()
    loss_D_real = criterion_GAN(pred_real, target_real)
    fake_A = fake_A_buffer.push_and_pop(fake_A)
    pred_fake = netD_A(fake_A.detach()).float()
    loss_D_fake = criterion_GAN(pred_fake, target_fake)
    loss_D_A = (loss_D_real + loss_D_fake) * 0.5
    loss_D_A.backward()
    da_optimizer.step()
    # 训练判别器B
    db_optimizer.zero_grad()
    pred_real = netD_B(real_b)
    loss_D_real = criterion_GAN(pred_real, target_real)
    fake_B = fake_B_buffer.push_and_pop(fake_B)
    pred_fake = netD_B(fake_B.detach())
    loss_D_fake = criterion_GAN(pred_fake, target_fake)
    loss_D_B = (loss_D_real + loss_D_fake) * 0.5
    loss_D_B.backward()
    db_optimizer.step()
    
    
    #损失打印,存储伪造图片
    print('Epoch[{}],loss_G:{:.6f} ,loss_D_A:{:.6f},loss_D_B:{:.6f}'
      .format(epoch, loss_G.data.item(), loss_D_A.data.item(), 
              loss_D_B.data.item()))
    if (epoch + 1) % 20 == 0 or epoch == 0:  
        b_fake = to_img(fake_B.data)
        a_fake = to_img(fake_A.data)
        a_real = to_img(real_a.data)
        b_real = to_img(real_b.data)
        save_image(a_fake, '../tmp/a_fake.png') 
        save_image(b_fake, '../tmp/b_fake.png') 
        save_image(a_real, '../tmp/a_real.png') 
        save_image(b_real, '../tmp/b_real.png') 
    
相关推荐
古希腊掌管学习的神14 分钟前
[机器学习]XGBoost(3)——确定树的结构
人工智能·机器学习
ZHOU_WUYI42 分钟前
4.metagpt中的软件公司智能体 (ProjectManager 角色)
人工智能·metagpt
靴子学长1 小时前
基于字节大模型的论文翻译(含免费源码)
人工智能·深度学习·nlp
梧桐树04292 小时前
python常用内建模块:collections
python
AI_NEW_COME2 小时前
知识库管理系统可扩展性深度测评
人工智能
Dream_Snowar2 小时前
速通Python 第三节
开发语言·python
海棠AI实验室3 小时前
AI的进阶之路:从机器学习到深度学习的演变(一)
人工智能·深度学习·机器学习
hunteritself3 小时前
AI Weekly『12月16-22日』:OpenAI公布o3,谷歌发布首个推理模型,GitHub Copilot免费版上线!
人工智能·gpt·chatgpt·github·openai·copilot
IT古董3 小时前
【机器学习】机器学习的基本分类-强化学习-策略梯度(Policy Gradient,PG)
人工智能·机器学习·分类
centurysee3 小时前
【最佳实践】Anthropic:Agentic系统实践案例
人工智能