用python写一段基于GAN的图像去雾代码

好的，这里是一份基于GAN的图像去雾代码：

首先，我们需要导入一些必要的库：

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import transforms
from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
import numpy as np
from PIL import Image

接着，我们需要定义一些超参数：

batch_size = 32
lr = 0.0002
beta1 = 0.5
epochs = 50
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

然后，我们需要构建生成器和判别器：

class Generator(nn.Module):
    def __init__(self):
        super(Generator, self).__init__()
        self.fc1 = nn.Linear(100, 384 * 8 * 8)
        self.conv1 = nn.ConvTranspose2d(384, 192, kernel_size=5, stride=2, padding=2)
        self.bn1 = nn.BatchNorm2d(192)
        self.conv2 = nn.ConvTranspose2d(192, 96, kernel_size=5, stride=2, padding=2)
        self.bn2 = nn.BatchNorm2d(96)
        self.conv3 = nn.ConvTranspose2d(96, 3, kernel_size=5, stride=2, padding=2)
        self.tanh = nn.Tanh()

    def forward(self, x):
        x = self.fc1(x)
        x = x.view(-1, 384, 8, 8)
        x = nn.functional.relu(self.bn1(self.conv1(x)))
        x = nn.functional.relu(self.bn2(self.conv2(x)))
        x = self.tanh(self.conv3(x))
        return x


class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()
        self.conv1 = nn.Conv2d(3, 96, kernel_size=5, stride=2, padding=2)
        self.conv2 = nn.Conv2d(96, 192, kernel_size=5, stride=2, padding=2)
        self.bn2 = nn.BatchNorm2d(192)
        self.conv3 = nn.Conv2d(192, 384, kernel_size=5, stride=2, padding=2)
        self.bn3 = nn.BatchNorm2d(384)
        self.fc1 = nn.Linear(384 * 8 * 8, 1)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        x = nn.functional.relu(self.conv1(x))
        x = nn.functional.relu(self.bn2(self.conv2(x)))
        x = nn.functional.relu(self.bn3(self.conv3(x)))
        x = x.view(-1, 384 * 8 * 8)
        x = self.sigmoid(self.fc1(x))
        return x

接着，我们需要定义一些辅助函数：

def denorm(x):
    out = (x + 1) / 2
    return out.clamp(0, 1)


def show_img(img):
    img = denorm(img)
    npimg = img.detach().numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()


def save_img(img, path):
    img = denorm(img)
    npimg = img.detach().numpy()
    npimg = np.transpose(npimg, (1, 2, 0))
    npimg = npimg * 255
    npimg = npimg.astype(np.uint8)
    im = Image.fromarray(npimg)
    im.save(path)

然后，我们需要加载数据集：

transform = transforms.Compose([
    transforms.Resize(64),
    transforms.CenterCrop(64),
    transforms.ToTensor(),
    transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
])

train_dataset = ImageFolder('./data/train', transform=transform)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4)

接下来，我们需要初始化生成器和判别器，并定义优化器和损失函数：

G = Generator().to(device)
D = Discriminator().to(device)

G.apply(weights_init_normal)
D.apply(weights_init_normal)

optimizer_G = optim.Adam(G.parameters(), lr=lr, betas=(beta1, 0.999))
optimizer_D = optim.Adam(D.parameters(), lr=lr, betas=(beta1, 0.999))

criterion = nn.BCELoss()

在训练过程中，我们需要先生成一些随机噪声，然后将其输入到生成器中生成图像。接着，我们将生成的图像和真实图像一起输入到判别器中进行判别。然后，我们计算生成器和判别器的损失，并更新参数。最后，我们输出一些图像来观察训练的效果：

for epoch in range(epochs):
    for i, (imgs, _) in enumerate(train_loader):

        # Adversarial ground truths
        valid = torch.ones(imgs.size(0), 1).to(device)
        fake = torch.zeros(imgs.size(0), 1).to(device)

        # Configure input
        real_imgs = imgs.to(device)

        # Train Generator
        optimizer_G.zero_grad()

        # Sample noise as generator input
        z = torch.randn(imgs.size(0), 100).to(device)

        # Generate a batch of images
        gen_imgs = G(z)

        # Loss measures generator's ability to fool the discriminator
        g_loss = criterion(D(gen_imgs), valid)

        g_loss.backward()
        optimizer_G.step()

        # Train Discriminator
        optimizer_D.zero_grad()

        # Measure discriminator's ability to classify real from generated samples
        real_loss = criterion(D(real_imgs), valid)
        fake_loss = criterion(D(gen_imgs.detach()), fake)
        d_loss = (real_loss + fake_loss) / 2

        d_loss.backward()
        optimizer_D.step()

        batches_done = epoch * len(train_loader) + i
        if batches_done % 50 == 0:
            print("[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f]"
                  % (epoch, epochs, i, len(train_loader), d_loss.item(), g_loss.item()))

        if batches_done % 400 == 0:
            save_img(gen_imgs.data[:25], "images/%d.png" % batches_done)

注意，我们在每个 epoch 结束后输出一些图像来观察训练的效果。

完整的代码如下所示：

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import transforms
from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
import numpy as np
from PIL import Image


batch_size = 32
lr = 0.0002
beta1 = 0.5
epochs = 50
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")


class Generator(nn.Module):
    def __init__(self):
        super(Generator, self).__init__()
        self.fc1 = nn.Linear(100, 384 * 8 * 8)
        self.conv1 = nn.ConvTranspose2d(384, 192, kernel_size=5, stride=2, padding=2)
        self.bn1 = nn.BatchNorm2d(192)
        self.conv2 = nn.ConvTranspose2d(192, 96, kernel_size=5, stride=2, padding=2)
        self.bn2 = nn.BatchNorm2d(96)
        self.conv3 = nn.ConvTranspose2d(96, 3, kernel_size=5, stride=2, padding=2)
        self.tanh = nn.Tanh()

    def forward(self, x):
        x = self.fc1(x)
        x = x.view(-1, 384, 8, 8)
        x = nn.functional.relu(self.bn1(self.conv1(x)))
        x = nn.functional.relu(self.bn2(self.conv2(x)))
        x = self.tanh(self.conv3(x))
        return x


class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()
        self.conv1 = nn.Conv2d(3, 96, kernel_size=5, stride=2, padding=2)
        self.conv2 = nn.Conv2d(96, 192, kernel_size=5, stride=2, padding=2)
        self.bn2 = nn.BatchNorm2d(192)
        self.conv3 = nn.Conv2d(192, 384, kernel_size=5, stride=2, padding=2)
        self.bn3 = nn.BatchNorm2d(384)
        self.fc1 = nn.Linear(384 * 8 * 8, 1)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        x = nn.functional.relu(self.conv1(x))
        x = nn.functional.relu(self.bn2(self.conv2(x)))
        x = nn.functional.relu(self.bn3(self.conv3(x)))
        x = x.view(-1, 384 * 8 * 8)
        x = self.sigmoid(self.fc1(x))
        return x


def weights_init_normal(m):
    classname = m.__class__.__name__
    if classname.find("Conv") != -1:
        nn.init.normal_(m.weight.data, 0.0, 0.02)
    elif classname.find("BatchNorm") != -1:
        nn.init.normal_(m.weight.data, 1.0, 0.02)
        nn.init.constant_(m.bias.data, 0.0)


def denorm(x):
    out = (x + 1) / 2
    return out.clamp(0, 1)


def show_img(img):
    img = denorm(img)
    npimg = img.detach().numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()


def save_img(img, path):
    img = denorm(img)
    npimg = img.detach().numpy()
    npimg = np.transpose(npimg, (1, 2, 0))
    npimg = npimg * 255
    npimg = npimg.astype(np.uint8)
    im = Image.fromarray(npimg)
    im.save(path)


transform = transforms.Compose([
    transforms.Resize(64),
    transforms.CenterCrop(64),
    transforms.ToTensor(),
    transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
])

train_dataset = ImageFolder('./data/train', transform=transform)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4)

G = Generator().to(device)
D = Discriminator().to(device)

G.apply(weights_init_normal)
D.apply(weights_init_normal)

optimizer_G = optim.Adam(G.parameters(), lr=lr, betas=(beta1, 0.999))
optimizer_D = optim.Adam(D.parameters(), lr=lr, betas=(beta1, 0.999))

criterion = nn.BCELoss()

for epoch in range(epochs):
    for i, (imgs, _) in enumerate(train_loader):

        # Adversarial ground truths
        valid = torch.ones(imgs.size(0), 1).to(device)
        fake = torch.zeros(imgs.size(0), 1).to(device)

        # Configure input
        real_imgs = imgs.to(device)

        # Train Generator
        optimizer_G.zero_grad()

        # Sample noise as generator input
        z = torch.randn(imgs.size(0), 100).to(device)

        # Generate a batch of images
        gen_imgs = G(z)

        # Loss measures generator's ability to fool the discriminator
        g_loss = criterion(D(gen_imgs), valid)

        g_loss.backward()
        optimizer_G.step()

        # Train Discriminator
        optimizer_D.zero_grad()

        # Measure discriminator's ability to classify real from generated samples
        real_loss = criterion(D(real_imgs), valid)
        fake_loss = criterion(D(gen_imgs.detach()), fake)
        d_loss = (real_loss + fake_loss) / 2

        d_loss.backward()
        optimizer_D.step()

        batches_done = epoch * len(train_loader) + i
        if batches_done % 50 == 0:
            print("[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f]"
                  % (epoch, epochs, i, len(train_loader), d_loss.item(), g_loss.item()))

        if batches_done % 400 == 0:
            save_img(gen_imgs.data[:25], "images/%d.png" % batches_done)