PyTorch 实现 Conditional DCGAN(条件深度卷积生成对抗网络)进行图像到图像转换的示例代码

以下是一个使用 PyTorch 实现 Conditional DCGAN(条件深度卷积生成对抗网络)进行图像到图像转换的示例代码。该代码包含训练和可视化部分,假设输入为图片和 4 个工艺参数,根据这些输入生成相应的图片。

1. 导入必要的库

python 复制代码
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader, Dataset
import numpy as np
import matplotlib.pyplot as plt

# 检查是否有可用的 GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

2. 定义数据集类

python 复制代码
class ImagePairDataset(Dataset):
    def __init__(self, image_pairs, params):
        self.image_pairs = image_pairs
        self.params = params

    def __len__(self):
        return len(self.image_pairs)

    def __getitem__(self, idx):
        input_image, target_image = self.image_pairs[idx]
        param = self.params[idx]
        return input_image, target_image, param

3. 定义生成器和判别器

python 复制代码
# 生成器
class Generator(nn.Module):
    def __init__(self, z_dim=4, img_channels=3):
        super(Generator, self).__init__()
        self.gen = nn.Sequential(
            # 输入: [batch_size, z_dim + 4, 1, 1]
            self._block(z_dim + 4, 1024, 4, 1, 0),  # [batch_size, 1024, 4, 4]
            self._block(1024, 512, 4, 2, 1),  # [batch_size, 512, 8, 8]
            self._block(512, 256, 4, 2, 1),  # [batch_size, 256, 16, 16]
            self._block(256, 128, 4, 2, 1),  # [batch_size, 128, 32, 32]
            nn.ConvTranspose2d(128, img_channels, kernel_size=4, stride=2, padding=1),
            nn.Tanh()
        )

    def _block(self, in_channels, out_channels, kernel_size, stride, padding):
        return nn.Sequential(
            nn.ConvTranspose2d(
                in_channels, out_channels, kernel_size, stride, padding, bias=False
            ),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(True)
        )

    def forward(self, z, params):
        params = params.view(params.size(0), 4, 1, 1)
        x = torch.cat([z, params], dim=1)
        return self.gen(x)

# 判别器
class Discriminator(nn.Module):
    def __init__(self, img_channels=3):
        super(Discriminator, self).__init__()
        self.disc = nn.Sequential(
            # 输入: [batch_size, img_channels + 4, 64, 64]
            nn.Conv2d(img_channels + 4, 64, kernel_size=4, stride=2, padding=1),
            nn.LeakyReLU(0.2),
            self._block(64, 128, 4, 2, 1),  # [batch_size, 128, 16, 16]
            self._block(128, 256, 4, 2, 1),  # [batch_size, 256, 8, 8]
            self._block(256, 512, 4, 2, 1),  # [batch_size, 512, 4, 4]
            nn.Conv2d(512, 1, kernel_size=4, stride=2, padding=0),
            nn.Sigmoid()
        )

    def _block(self, in_channels, out_channels, kernel_size, stride, padding):
        return nn.Sequential(
            nn.Conv2d(
                in_channels, out_channels, kernel_size, stride, padding, bias=False
            ),
            nn.BatchNorm2d(out_channels),
            nn.LeakyReLU(0.2)
        )

    def forward(self, img, params):
        params = params.view(params.size(0), 4, 1, 1).repeat(1, 1, img.size(2), img.size(3))
        x = torch.cat([img, params], dim=1)
        return self.disc(x)

4. 训练代码

python 复制代码
def train_conditional_dcgan(image_pairs, params, batch_size=32, epochs=10, lr=0.0002, z_dim=4):
    dataset = ImagePairDataset(image_pairs, params)
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

    gen = Generator(z_dim).to(device)
    disc = Discriminator().to(device)

    criterion = nn.BCELoss()
    opt_gen = optim.Adam(gen.parameters(), lr=lr, betas=(0.5, 0.999))
    opt_disc = optim.Adam(disc.parameters(), lr=lr, betas=(0.5, 0.999))

    for epoch in range(epochs):
        for i, (input_images, target_images, param) in enumerate(dataloader):
            input_images = input_images.to(device)
            target_images = target_images.to(device)
            param = param.to(device)

            # 训练判别器
            opt_disc.zero_grad()
            real_labels = torch.ones((target_images.size(0), 1, 1, 1)).to(device)
            fake_labels = torch.zeros((target_images.size(0), 1, 1, 1)).to(device)

            # 计算判别器对真实图像的损失
            real_output = disc(target_images, param)
            d_real_loss = criterion(real_output, real_labels)

            # 生成假图像
            z = torch.randn(target_images.size(0), z_dim, 1, 1).to(device)
            fake_images = gen(z, param)

            # 计算判别器对假图像的损失
            fake_output = disc(fake_images.detach(), param)
            d_fake_loss = criterion(fake_output, fake_labels)

            # 总判别器损失
            d_loss = d_real_loss + d_fake_loss
            d_loss.backward()
            opt_disc.step()

            # 训练生成器
            opt_gen.zero_grad()
            output = disc(fake_images, param)
            g_loss = criterion(output, real_labels)
            g_loss.backward()
            opt_gen.step()

        print(f'Epoch [{epoch+1}/{epochs}] D_loss: {d_loss.item():.4f} G_loss: {g_loss.item():.4f}')

    return gen

5. 可视化代码

python 复制代码
def visualize_generated_images(gen, input_images, params, z_dim=4):
    input_images = input_images.to(device)
    params = params.to(device)
    z = torch.randn(input_images.size(0), z_dim, 1, 1).to(device)
    fake_images = gen(z, params).cpu().detach()

    fig, axes = plt.subplots(1, input_images.size(0), figsize=(15, 3))
    for i in range(input_images.size(0)):
        img = fake_images[i].permute(1, 2, 0).numpy()
        img = (img + 1) / 2  # 从 [-1, 1] 转换到 [0, 1]
        axes[i].imshow(img)
        axes[i].axis('off')
    plt.show()

6. 示例使用

python 复制代码
# 假设 image_pairs 是一个包含图像对的列表,params 是一个包含 4 个工艺参数的列表
image_pairs = []  # 这里需要替换为实际的图像对数据
params = []  # 这里需要替换为实际的工艺参数数据

# 训练模型
gen = train_conditional_dcgan(image_pairs, params)

# 可视化生成的图像
test_input_images, test_target_images, test_params = image_pairs[:5], image_pairs[:5], params[:5]
test_input_images = torch.stack([torch.tensor(img) for img in test_input_images]).float()
test_params = torch.tensor(test_params).float()
visualize_generated_images(gen, test_input_images, test_params)

代码说明

  1. 数据集类ImagePairDataset 用于加载图像对和工艺参数。
  2. 生成器和判别器GeneratorDiscriminator 分别定义了生成器和判别器的网络结构。
  3. 训练代码train_conditional_dcgan 函数用于训练 Conditional DCGAN 模型。
  4. 可视化代码visualize_generated_images 函数用于可视化生成的图像。
  5. 示例使用:最后部分展示了如何使用上述函数进行训练和可视化。

请注意,你需要将 image_pairsparams 替换为实际的数据集。此外,代码中的超参数(如 batch_sizeepochslr 等)可以根据实际情况进行调整。

相关推荐
千宇宙航3 小时前
闲庭信步使用SV搭建图像测试平台:第三十一课——基于神经网络的手写数字识别
图像处理·人工智能·深度学习·神经网络·计算机视觉·fpga开发
onceco3 小时前
领域LLM九讲——第5讲 为什么选择OpenManus而不是QwenAgent(附LLM免费api邀请码)
人工智能·python·深度学习·语言模型·自然语言处理·自动化
天水幼麟4 小时前
动手学深度学习-学习笔记(总)
笔记·深度学习·学习
jndingxin6 小时前
OpenCV CUDA模块设备层-----高效地计算两个 uint 类型值的带权重平均值
人工智能·opencv·计算机视觉
天水幼麟6 小时前
动手学深度学习-学习笔记【二】(基础知识)
笔记·深度学习·学习
Sweet锦6 小时前
零基础保姆级本地化部署文心大模型4.5开源系列
人工智能·语言模型·文心一言
hie988947 小时前
MATLAB锂离子电池伪二维(P2D)模型实现
人工智能·算法·matlab
晨同学03277 小时前
opencv的颜色通道问题 & rgb & bgr
人工智能·opencv·计算机视觉
蓝婷儿7 小时前
Python 机器学习核心入门与实战进阶 Day 3 - 决策树 & 随机森林模型实战
人工智能·python·机器学习
大千AI助手7 小时前
PageRank:互联网的马尔可夫链平衡态
人工智能·机器学习·贝叶斯·mc·pagerank·条件概率·马尔科夫链