cyclegan马变斑马代码

### 使用 PyTorch 实现 CycleGAN 马到斑马图像转换

为了实现从马到斑马的图像转换，可以采用基于 PyTorch 的 CycleGAN 模型。以下是完整的代码实现：

#### 1. 导入必要的库

import torch
from torchvision import transforms, datasets
from torch.utils.data import DataLoader
import itertools
import os
from PIL import Image

#### 2. 定义数据加载器

class ImageDataset(torch.utils.data.Dataset):
    def __init__(self, root, transform=None, mode='train'):
        self.transform = transform
        self.files_A = sorted(os.listdir(os.path.join(root, 'A')))
        self.files_B = sorted(os.listdir(os.path.join(root, 'B')))

    def __getitem__(self, index):
        item_A = self.transform(Image.open(self.files_A[index % len(self.files_A)]))
        item_B = self.transform(Image.open(self.files_B[index % len(self.files_B)]))

        if self.mode == 'test':
            return {'A': item_A, 'B': item_B}
        else:
            return {'A': item_A, 'B': item_B}

    def __len__(self):
        return max(len(self.files_A), len(self.files_B))

#### 3. 构建网络结构

import torch.nn as nn

class Generator(nn.Module):
    def __init__(self, input_nc=3, output_nc=3, n_residual_blocks=9):
        super(Generator, self).__init__()

        model = [nn.ReflectionPad2d(3),
                 nn.Conv2d(input_nc, 64, kernel_size=7, padding=0, bias=True),
                 nn.InstanceNorm2d(64),
                 nn.ReLU(True)]

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

        # ResNet 块
        for _ in range(n_residual_blocks):
            model += [ResidualBlock(in_features)]

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

        # 输出层
        model += [nn.ReflectionPad2d(3),
                  nn.Conv2d(64, output_nc, kernel_size=7, padding=0),
                  nn.Tanh()]

        self.model = nn.Sequential(*model)

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


class Discriminator(nn.Module):
    def __init__(self, input_nc=3):
        super(Discriminator, self).__init__()
        
        model = [
            nn.Conv2d(input_nc, 64, kernel_size=4, stride=2, padding=1),
            nn.LeakyReLU(0.2, inplace=True),

            nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1),
            nn.InstanceNorm2d(128),
            nn.LeakyReLU(0.2, inplace=True),

            nn.Conv2d(128, 256, kernel_size=4, stride=2, padding=1),
            nn.InstanceNorm2d(256),
            nn.LeakyReLU(0.2, inplace=True),

            nn.Conv2d(256, 512, kernel_size=4, padding=1),
            nn.InstanceNorm2d(512),
            nn.LeakyReLU(0.2, inplace=True),

            nn.Conv2d(512, 1, kernel_size=4, padding=1)
        ]

        self.model = nn.Sequential(*model)

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

#### 4. 训练过程

def train(dataloader, G_AB, G_BA, D_A, D_B, criterion_GAN, criterion_cycle, optimizer_G, optimizer_D_A, optimizer_D_B):

    for epoch in range(num_epochs):
        for i, batch in enumerate(dataloader):
            
            real_A = batch['A'].to(device)
            real_B = batch['B'].to(device)

            valid = torch.ones((real_A.size(0), *D_A.output_shape)).to(device)
            fake = torch.zeros((real_A.size(0), *D_A.output_shape)).to(device)

            ###### Generators A2B and B2A ######
            optimizer_G.zero_grad()

            # Identity loss
            same_B = G_AB(real_B)
            loss_identity_B = criterion_cycle(same_B, real_B) 

            same_A = G_BA(real_A)
            loss_identity_A = criterion_cycle(same_A, real_A)  

            # GAN loss
            fake_B = G_AB(real_A)
            pred_fake = D_B(fake_B)
            loss_GAN_A2B = criterion_GAN(pred_fake, valid)

            fake_A = G_BA(real_B)
            pred_fake = D_A(fake_A)
            loss_GAN_B2A = criterion_GAN(pred_fake, valid)

            # Cycle loss 
            recovered_A = G_BA(fake_B)
            loss_cycle_ABA = criterion_cycle(recovered_A, real_A)

            recovered_B = G_AB(fake_A)
            loss_cycle_BAB = criterion_cycle(recovered_B, real_B)

            # Total loss
            loss_G = (loss_identity_A + loss_identity_B +
                    lambda_cyc * (loss_cycle_ABA + loss_cycle_BAB) +
                    lambda_adv * (loss_GAN_A2B + loss_GAN_B2A))

            loss_G.backward()
            optimizer_G.step()

            ###### Discriminator A ######
            optimizer_D_A.zero_grad()

            # Real loss
            pred_real = D_A(real_A)
            loss_real = criterion_GAN(pred_real, valid)

            # Fake loss (on batch of previously generated samples)
            fake_A_ = fake_A_buffer.push_and_pop(fake_A)
            pred_fake = D_A(fake_A_.detach())
            loss_fake = criterion_GAN(pred_fake, fake)

            # Total loss
            loss_D_A = (loss_real + loss_fake) / 2

            loss_D_A.backward()
            optimizer_D_A.step()

            ###### Discriminator B ######
            optimizer_D_B.zero_grad()

            # Real loss
            pred_real = D_B(real_B)
            loss_real = criterion_GAN(pred_real, valid)

            # Fake loss (on batch