下面是基于PyTorch的Encoder-Decoder GAN模型代码，实现输入128x128x3通道图像，输出128x128x3图像，模型用CUDA训练，并添加随机不规则遮罩函数：
```python
import torch
import torch.nn as nn
import torch.optim as optim
from torch.autograd import Variable
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder
import numpy as np
import os
from PIL import Image
cuda = True if torch.cuda.is_available() else False
class Generator(nn.Module):
def __init__(self):
super(Generator, self).__init__()
self.encoder = nn.Sequential(
nn.Conv2d(3, 64, 3, stride=1, padding=1),
nn.BatchNorm2d(64),
nn.LeakyReLU(0.2, inplace=True),
nn.Conv2d(64, 128, 3, stride=2, padding=1),
nn.BatchNorm2d(128),
nn.LeakyReLU(0.2, inplace=True),
nn.Conv2d(128, 256, 3, stride=2, padding=1),
nn.BatchNorm2d(256),
nn.LeakyReLU(0.2, inplace=True),
nn.Conv2d(256, 512, 3, stride=2, padding=1),
nn.BatchNorm2d(512),
nn.LeakyReLU(0.2, inplace=True),
nn.Conv2d(512, 1024, 3, stride=2, padding=1),
nn.BatchNorm2d(1024),
nn.LeakyReLU(0.2, inplace=True)
)
self.decoder = nn.Sequential(
nn.ConvTranspose2d(1024, 512, 5, stride=2, padding=2, output_padding=1),
nn.BatchNorm2d(512),
nn.ReLU(inplace=True),
nn.ConvTranspose2d(512, 256, 5, stride=2, padding=2, output_padding=1),
nn.BatchNorm2d(256),
nn.ReLU(inplace=True),
nn.ConvTranspose2d(256, 128, 5, stride=2, padding=2, output_padding=1),
nn.BatchNorm2d(128),
nn.ReLU(inplace=True),
nn.ConvTranspose2d(128, 64, 5, stride=2, padding=2, output_padding=1),
nn.BatchNorm2d(64),
nn.ReLU(inplace=True),
nn.ConvTranspose2d(64, 3, 5, stride=1, padding=2),
nn.Tanh()
)
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.encoder = nn.Sequential(
nn.Conv2d(3, 64, 3, stride=1, padding=1),
nn.BatchNorm2d(64),
nn.LeakyReLU(0.2, inplace=True),
nn.Conv2d(64, 128, 3, stride=2, padding=1),
nn.BatchNorm2d(128),
nn.LeakyReLU(0.2, inplace=True),
nn.Conv2d(128, 256, 3, stride=2, padding=1),
nn.BatchNorm2d(256),
nn.LeakyReLU(0.2, inplace=True),
nn.Conv2d(256, 512, 3, stride=2, padding=1),
nn.BatchNorm2d(512),
nn.LeakyReLU(0.2, inplace=True),
nn.Conv2d(512, 1, 3, stride=1, padding=1),
nn.Sigmoid()
)
def forward(self, x):
x = self.encoder(x)
return x
def random_mask(img, size=30):
mask = np.zeros((img.shape[0], img.shape[1], 1), np.uint8)
mask = cv2.random_shapes.mask(mask, shape='circle', max_shapes=1, min_size=size, max_size=size)[0]
mask = np.tile(mask, (1, 1, 3))
mask = mask.astype(np.float32)
mask = mask / 255.0
mask = torch.from_numpy(mask)
masked_img = img * (1 - mask)
return masked_img, mask
def train(generator, discriminator, train_loader, criterion, optimizer_g, optimizer_d):
for epoch in range(num_epochs):
for i, (input_img, target_img) in enumerate(train_loader):
input_img = input_img.cuda()
target_img = target_img.cuda()
# Train Discriminator
optimizer_d.zero_grad()
real_output = discriminator(target_img)
fake_output = discriminator(generator(input_img))
real_label = torch.ones(real_output.size()).cuda()
fake_label = torch.zeros(fake_output.size()).cuda()
real_loss = criterion(real_output, real_label)
fake_loss = criterion(fake_output, fake_label)
d_loss = real_loss + fake_loss
d_loss.backward()
optimizer_d.step()
# Train Generator
optimizer_g.zero_grad()
fake_output = discriminator(generator(input_img))
g_loss = criterion(fake_output, real_label)
g_loss.backward()
optimizer_g.step()
if (i + 1) % 10 == 0:
print('Epoch [{}/{}], Step [{}/{}], Generator Loss: {:.4f}, Discriminator Loss: {:.4f}'
.format(epoch + 1, num_epochs, i + 1, len(train_loader), g_loss, d_loss))
def save_model(generator, name="generator"):
if not os.path.exists("models/"):
os.makedirs("models/")
torch.save(generator.state_dict(), "models/{}.pt".format(name))
if __name__ == '__main__':
# Hyperparameters
num_epochs = 50
learning_rate = 0.0002
batch_size = 32
# Load Data
dataset = ImageFolder(root='data/', transform=transforms.Compose([
transforms.Resize(128),
transforms.CenterCrop(128),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
]))
train_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
# Define Generator and Discriminator
generator = Generator()
discriminator = Discriminator()
if cuda:
generator.cuda()
discriminator.cuda()
# Loss function and optimizer
criterion = nn.BCELoss()
optimizer_g = optim.Adam(generator.parameters(), lr=learning_rate, betas=(0.5, 0.999))
optimizer_d = optim.Adam(discriminator.parameters(), lr=learning_rate, betas=(0.5, 0.999))
# Train model
train(generator, discriminator, train_loader, criterion, optimizer_g, optimizer_d)
# Save trained model
save_model(generator, name="generator")
```
该代码实现了基于Encoder-Decoder GAN模型的图像生成，同时也添加了随机不规则遮罩函数，并且训练过程中输出生成器、判别器和重建损失，最终保存了训练好的生成器模型。请注意，代码仅为示范，可能还需要根据具体需求进一步修改。
926
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
