AIGC从入门到实战：众里寻他千百度：ChatGPT 及其他 AIGC 赋能个人 - qidao123.com技术社区-IT企服评测·应用市场

import torch
import torch.nn as nn
import torchvision.transforms as transforms
from torchvision.utils import save_image
from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder
# 定义生成器和判别器
class Generator(nn.Module):
def __init__(self):
super(Generator, self).__init__()
self.model = nn.Sequential(
nn.Linear(100, 512),
nn.ReLU(True),
nn.Linear(512, 1024),
nn.ReLU(True),
nn.Linear(1024, 784),
nn.Tanh()
)
def forward(self, x):
x = self.model(x)
x = x.view(-1, 1, 28, 28)
return x
class Discriminator(nn.Module):
def __init__(self):
super(Discriminator, self).__init__()
self.model = nn.Sequential(
nn.Linear(784, 1024),
nn.LeakyReLU(0.2, inplace=True),
nn.Linear(1024, 512),
nn.LeakyReLU(0.2, inplace=True),
nn.Linear(512, 1),
nn.Sigmoid()
)
def forward(self, x):
x = x.view(-1, 784)
x = self.model(x)
return x
# 训练参数
batch_size = 64
learning_rate = 0.0002
epochs = 100
# 加载数据
transform = transforms.Compose([
transforms.Resize((64, 64)),
transforms.ToTensor()
])
dataset = ImageFolder(root='data', transform=transform)
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
# 初始化模型
generator = Generator()
discriminator = Discriminator()
# 定义优化器
optimizer_g = torch.optim.Adam(generator.parameters(), lr=learning_rate)
optimizer_d = torch.optim.Adam(discriminator.parameters(), lr=learning_rate)
# 训练过程
for epoch in range(epochs):
for i, (imgs, _) in enumerate(dataloader):
# 训练判别器
real_imgs = imgs.to(device)
z = torch.randn(imgs.size(0), 100).to(device)
fake_imgs = generator(z).to(device)
real_loss = nn.BCELoss()(discriminator(real_imgs), torch.ones_like(discriminator(real_imgs)))
fake_loss = nn.BCELoss()(discriminator(fake_imgs.detach()), torch.zeros_like(discriminator(fake_imgs.detach())))
d_loss = 0.5 * (real_loss + fake_loss)
optimizer_d.zero_grad()
d_loss.backward()
optimizer_d.step()
# 训练生成器
z = torch.randn(imgs.size(0), 100).to(device)
fake_imgs = generator(z).to(device)
g_loss = nn.BCELoss()(discriminator(fake_imgs), torch.ones_like(discriminator(fake_imgs)))
optimizer_g.zero_grad()
g_loss.backward()
optimizer_g.step()
# 打印训练信息
if i % 100 == 0:
print(f"Epoch {epoch}, Step {i}, D_loss: {d_loss.item()}, G_loss: {g_loss.item()}")
# 保存生成图像
save_image(fake_imgs, 'fake_images.png')
# 关闭模型
generator.eval()
discriminator.eval()
# 生成图像
with torch.no_grad():
z = torch.randn(1, 100).to(device)
fake_img = generator(z)
save_image(fake_img, 'generated_image.png')

复制代码