pytorch对抗生成网络间隔5个epoch保存模型

可以使用PyTorch的checkpoint机制来定期保存模型。以下是一个示例代码，每5个epoch保存一次模型：

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.transforms as transforms
import torchvision.datasets as datasets
from torch.utils.data import DataLoader

# 定义GAN模型
class Generator(nn.Module):
    def __init__(self):
        super(Generator, self).__init__()
        # 省略生成器的定义

class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()
        # 省略判别器的定义

# 定义训练函数
def train_gan(generator, discriminator, train_loader, num_epochs, device):
    # 定义损失函数和优化器
    criterion = nn.BCELoss()
    gen_optimizer = optim.Adam(generator.parameters(), lr=0.0002, betas=(0.5, 0.999))
    dis_optimizer = optim.Adam(discriminator.parameters(), lr=0.0002, betas=(0.5, 0.999))

    # 开始训练
    for epoch in range(num_epochs):
        for i, (real_images, _) in enumerate(train_loader):
            # 训练判别器
            real_labels = torch.ones(real_images.size(0)).to(device)
            fake_labels = torch.zeros(real_images.size(0)).to(device)
            real_images = real_images.to(device)
            fake_images = generator(torch.randn(real_images.size(0), 100).to(device))

            real_outputs = discriminator(real_images).view(-1)
            real_loss = criterion(real_outputs, real_labels)
            fake_outputs = discriminator(fake_images).view(-1)
            fake_loss = criterion(fake_outputs, fake_labels)
            dis_loss = real_loss + fake_loss

            dis_optimizer.zero_grad()
            dis_loss.backward()
            dis_optimizer.step()

            # 训练生成器
            fake_images = generator(torch.randn(real_images.size(0), 100).to(device))
            fake_outputs = discriminator(fake_images).view(-1)
            gen_loss = criterion(fake_outputs, real_labels)

            gen_optimizer.zero_grad()
            gen_loss.backward()
            gen_optimizer.step()

        # 保存模型
        if epoch % 5 == 0:
            checkpoint = {
                'epoch': epoch,
                'generator_state_dict': generator.state_dict(),
                'discriminator_state_dict': discriminator.state_dict(),
                'gen_optimizer_state_dict': gen_optimizer.state_dict(),
                'dis_optimizer_state_dict': dis_optimizer.state_dict(),
            }
            torch.save(checkpoint, f'gan_checkpoint_epoch_{epoch}.pt')

# 加载数据集
train_dataset = datasets.MNIST(root='data/', train=True, transform=transforms.ToTensor(), download=True)
train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True)

# 创建模型和设备
generator = Generator().to('cuda')
discriminator = Discriminator().to('cuda')

# 训练GAN模型
train_gan(generator, discriminator, train_loader, num_epochs=100, device='cuda')

在每个epoch结束时，我们都会检查当前epoch是否是5的倍数，如果是就保存一个checkpoint，其中包含了当前epoch的状态以及优化器的状态。这样我们就可以在需要的时候恢复模型的状态。