做消融实验 只保留一层unet把完整代码复制粘贴

为了进行消融实验并只保留一层UNet，我们需要简化 `UNet` 和 `AttentionUNet` 模型的结构。以下是修改后的代码：

### 修改后的 `UNet` 类

class SimpleUNet(nn.Module):
    def __init__(self):
        super(SimpleUNet, self).__init__()
        self.encoder = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(64, 3, kernel_size=2, stride=2),
            nn.Sigmoid()
        )

    def forward(self, x):
        x = self.encoder(x)
        x = self.decoder(x)
        return x

### 修改后的 `AttentionUNet` 类

由于 `AttentionUNet` 是一个更复杂的模型，我们也可以将其简化为只有一个编码器和解码器层的版本：

class SimpleAttentionUNet(nn.Module):
    def __init__(self):
        super(SimpleAttentionUNet, self).__init__()
        self.encoder = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(64, 3, kernel_size=2, stride=2),
            nn.BatchNorm2d(3),
            nn.ReLU(inplace=True),
            nn.Sigmoid()
        )

    def forward(self, x):
        x = self.encoder(x)
        x = self.decoder(x)
        return x

### 修改后的 `Generator` 类

同样地，我们可以将 `Generator` 类也简化为只有一个编码器和解码器层的版本：

class SimpleGenerator(nn.Module):
    def __init__(self):
        super(SimpleGenerator, self).__init__()
        self.encoder = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(64, 3, kernel_size=2, stride=2),
            nn.BatchNorm2d(3),
            nn.ReLU(inplace=True),
            nn.Sigmoid()
        )

    def forward(self, x):
        x = self.encoder(x)
        x = self.decoder(x)
        return x

### 完整代码

下面是完整的代码，包括上述修改后的类和其他部分保持不变：

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image
import os
import argparse
import glob
import matplotlib.pyplot as plt
from tqdm import tqdm
import numpy as np

class SimpleUNet(nn.Module):
    def __init__(self):
        super(SimpleUNet, self).__init__()
        self.encoder = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(64, 3, kernel_size=2, stride=2),
            nn.Sigmoid()
        )

    def forward(self, x):
        x = self.encoder(x)
        x = self.decoder(x)
        return x

class SimpleAttentionUNet(nn.Module):
    def __init__(self):
        super(SimpleAttentionUNet, self).__init__()
        self.encoder = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(64, 3, kernel_size=2, stride=2),
            nn.BatchNorm2d(3),
            nn.ReLU(inplace=True),
            nn.Sigmoid()
        )

    def forward(self, x):
        x = self.encoder(x)
        x = self.decoder(x)
        return x

class SimpleGenerator(nn.Module):
    def __init__(self):
        super(SimpleGenerator, self).__init__()
        self.encoder = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(64, 3, kernel_size=2, stride=2),
            nn.BatchNorm2d(3),
            nn.ReLU(inplace=True),
            nn.Sigmoid()
        )

    def forward(self, x):
        x = self.encoder(x)
        x = self.decoder(x)
        return x

class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()
        self.main = nn.Sequential(
            nn.Conv2d(3, 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.BatchNorm2d(128),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(128, 256, kernel_size=4, stride=2, padding=1),
            nn.BatchNorm2d(256),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(256, 512, kernel_size=4, stride=2, padding=1),
            nn.BatchNorm2d(512),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(512, 1, kernel_size=16),
        )

    def forward(self, x):
        return self.main(x).view(-1)

def compute_iou(outputs, targets, threshold=0.5):
    outputs = (outputs > threshold).float()
    targets = (targets > threshold).float()
    intersection = (outputs * targets).sum(dim=(1, 2, 3))
    union = outputs.sum(dim=(1, 2, 3)) + targets.sum(dim=(1, 2, 3)) - intersection
    iou = (intersection + 1e-6) / (union + 1e-6)
    return iou.mean().item()

from skimage.metrics import peak_signal_noise_ratio as psnr_metric
from skimage.metrics import structural_similarity as ssim_metric

def compute_psnr(outputs, targets):
    outputs = outputs.cpu().detach().numpy()
    targets = targets.cpu().detach().numpy()
    psnr = 0
    for i in range(outputs.shape[0]):
        psnr += psnr_metric(targets[i], outputs[i], data_range=1.0)
    return psnr / outputs.shape[0]

def compute_ssim(outputs, targets):
    outputs = outputs.cpu().detach().numpy()
    targets = targets.cpu().detach().numpy()
    ssim = 0
    for i in range(outputs.shape[0]):
        output_img = outputs[i].transpose(1, 2, 0)
        target_img = targets[i].transpose(1, 2, 0)
        H, W, _ = output_img.shape
        min_dim = min(H, W)
        win_size = min(7, min_dim if min_dim % 2 == 1 else min_dim - 1)
        win_size = max(win_size, 3)
        ssim += ssim_metric(
            target_img, output_img, data_range=1.0, channel_axis=-1, win_size=win_size
        )
    return ssim / outputs.shape[0]

def wasserstein_loss(pred, target):
    return torch.mean(pred * target)

from torch.autograd import grad

def compute_gradient_penalty(discriminator, real_samples, fake_samples, device):
    alpha = torch.rand(real_samples.size(0), 1, 1, 1, device=device)
    interpolates = (alpha * real_samples + ((1 - alpha) * fake_samples)).requires_grad_(True)
    d_interpolates = discriminator(interpolates)
    fake = torch.ones(real_samples.size(0), device=device)
    gradients = grad(
        outputs=d_interpolates,
        inputs=interpolates,
        grad_outputs=fake,
        create_graph=True,
        retain_graph=True,
        only_inputs=True,
    )[0]
    gradients = gradients.view(gradients.size(0), -1)
    gradient_penalty = ((gradients.norm(2, dim=1) - 1) ** 2).mean()
    return gradient_penalty

def train_correction_model(generator, discriminator, dataloader, optimizer_G, optimizer_D, device, lambda_gp, lambda_pixel, n_critic):
    generator.train()
    discriminator.train()
    running_g_loss = 0.0
    running_d_loss = 0.0
    running_iou = 0.0
    running_psnr = 0.0
    running_ssim = 0.0
    for batch_idx, (inputs, targets) in enumerate(tqdm(dataloader, desc="Training")):
        inputs = inputs.to(device)
        targets = targets.to(device)
        # ---------------------
        # 训练判别器
        # ---------------------
        optimizer_D.zero_grad()
        corrected_images = generator(inputs)
        real_validity = discriminator(targets)
        fake_validity = discriminator(corrected_images.detach())
        gp = compute_gradient_penalty(discriminator, targets.data, corrected_images.data, device)
        d_loss = -torch.mean(real_validity) + torch.mean(fake_validity) + lambda_gp * gp
        d_loss.backward()
        optimizer_D.step()
        # ---
        if batch_idx % n_critic == 0:
            optimizer_G.zero_grad()
            corrected_images = generator(inputs)
            fake_validity = discriminator(corrected_images)
            g_adv_loss = -torch.mean(fake_validity)
            pixelwise_loss = nn.L1Loss()
            g_pixel_loss = pixelwise_loss(corrected_images, targets)
            g_loss = g_adv_loss + lambda_pixel * g_pixel_loss
            g_loss.backward()
            optimizer_G.step()
        else:
            g_loss = torch.tensor(0.0)
        running_g_loss += g_loss.item()
        running_d_loss += d_loss.item()
        iou = compute_iou(corrected_images, targets)
        psnr = compute_psnr(corrected_images, targets)
        ssim = compute_ssim(corrected_images, targets)
        running_iou += iou
        running_psnr += psnr
        running_ssim += ssim
    epoch_g_loss = running_g_loss / len(dataloader)
    epoch_d_loss = running_d_loss / len(dataloader)
    epoch_iou = running_iou / len(dataloader)
    epoch_psnr = running_psnr / len(dataloader)
    epoch_ssim = running_ssim / len(dataloader)
    return epoch_g_loss, epoch_d_loss, epoch_iou, epoch_psnr, epoch_ssim

def validate_correction_model(generator, discriminator, dataloader, device, lambda_gp):
    generator.eval()
    discriminator.eval()
    running_g_loss = 0.0
    running_d_loss = 0.0
    running_iou = 0.0
    running_psnr = 0.0
    running_ssim = 0.0
    with torch.no_grad():
        for inputs, targets in tqdm(dataloader, desc="Validation"):
            inputs = inputs.to(device)
            targets = targets.to(device)
            corrected_images = generator(inputs)
            real_validity = discriminator(targets)
            fake_validity = discriminator(corrected_images)
            d_loss = -torch.mean(real_validity) + torch.mean(fake_validity)
            g_adv_loss = -torch.mean(fake_validity)
            g_loss = g_adv_loss
            running_g_loss += g_loss.item()
            running_d_loss += d_loss.item()
            iou = compute_iou(corrected_images, targets)
            psnr = compute_psnr(corrected_images, targets)
            ssim = compute_ssim(corrected_images, targets)
            running_iou += iou
            running_psnr += psnr
            running_ssim += ssim
    epoch_g_loss = running_g_loss / len(dataloader)
    epoch_d_loss = running_d_loss / len(dataloader)
    epoch_iou = running_iou / len(dataloader)
    epoch_psnr = running_psnr / len(dataloader)
    epoch_ssim = running_ssim / len(dataloader)
    return epoch_g_loss, epoch_d_loss, epoch_iou, epoch_psnr, epoch_ssim

def visualize_results(generator, dataloader, device, num_images=10, save_path='./results'):
    generator.eval()
    inputs, targets = next(iter(dataloader))
    inputs = inputs.to(device)
    targets = targets.to(device)
    with torch.no_grad():
        corrected_images = generator(inputs)
    inputs = inputs.cpu().numpy()
    targets = targets.cpu().numpy()
    corrected_images = corrected_images.cpu().numpy()
    if not os.path.exists(save_path):
        os.makedirs(save_path)
    plt.figure(figsize=(20, 10))
    for i in range(num_images):
        plt.subplot(3, num_images, i + 1)
        plt.imshow(targets[i].transpose(1, 2, 0))
        plt.title("Original")
        plt.axis('off')
        plt.subplot(3, num_images, i + 1 + num_images)
        plt.imshow(inputs[i].transpose(1, 2, 0))
        plt.title("Simulated Colorblind")
        plt.axis('off')
        plt.subplot(3, num_images, i + 1 + 2 * num_images)
        plt.imshow(corrected_images[i].transpose(1, 2, 0))
        plt.title("Corrected")
        plt.axis('off')
    plt.tight_layout()
    plt.savefig(f'{save_path}_visualization.png')
    plt.show()

def plot_and_save_metrics(train_metrics, val_metrics, epoch, path='./metrics_plots'):
    if not os.path.exists(path):
        os.makedirs(path)
    epochs = np.arange(1, epoch + 1)
    train_g_losses, train_d_losses, train_ious, train_psnrs, train_ssims = zip(*train_metrics)
    val_g_losses, val_d_losses, val_ious, val_psnrs, val_ssims = zip(*val_metrics)
    plt.figure()
    plt.plot(epochs, train_g_losses, label='Training Generator Loss')
    plt.plot(epochs, val_g_losses, label='Validation Generator Loss')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.title('Generator Loss over Epochs')
    plt.legend()
    plt.savefig(f'{path}/generator_loss