训练horse2zebra CycleGAN模型时发生中断，如何用保存的四个.pth文件继续训练，给出具体代码

如果训练CycleGAN模型时发生中断，可以通过使用保存的.pth文件来恢复训练。以下是使用保存的四个.pth文件继续训练的具体代码：

import torch
from models import Generator, Discriminator, CycleGAN
from datasets import ImageDataset
from torch.utils.data import DataLoader

# 定义设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 加载数据集
dataset = ImageDataset(root='path/to/data', mode='train')
dataloader = DataLoader(dataset, batch_size=1, shuffle=True)

# 初始化生成器和判别器
G_AB = Generator().to(device)
G_BA = Generator().to(device)
D_A = Discriminator().to(device)
D_B = Discriminator().to(device)

# 加载保存的.pth文件
G_AB.load_state_dict(torch.load('path/to/G_AB.pth'))
G_BA.load_state_dict(torch.load('path/to/G_BA.pth'))
D_A.load_state_dict(torch.load('path/to/D_A.pth'))
D_B.load_state_dict(torch.load('path/to/D_B.pth'))

# 定义损失函数和优化器
criterion_GAN = torch.nn.MSELoss()
criterion_cycle = torch.nn.L1Loss()
optimizer_G = torch.optim.Adam(
    itertools.chain(G_AB.parameters(), G_BA.parameters()), lr=0.0002, betas=(0.5, 0.999))
optimizer_D_A = torch.optim.Adam(D_A.parameters(), lr=0.0002, betas=(0.5, 0.999))
optimizer_D_B = torch.optim.Adam(D_B.parameters(), lr=0.0002, betas=(0.5, 0.999))

# 初始化CycleGAN模型
model = CycleGAN(G_AB, G_BA, D_A, D_B, criterion_GAN, criterion_cycle, optimizer_G, optimizer_D_A, optimizer_D_B, device)

# 设置开始的epoch和iteration
start_epoch = 0
start_iteration = 0

# 加载保存的训练状态
checkpoint = torch.load('path/to/checkpoint.pth')
model.load_state_dict(checkpoint['model_state_dict'])
optimizer_G.load_state_dict(checkpoint['optimizer_G_state_dict'])
optimizer_D_A.load_state_dict(checkpoint['optimizer_D_A_state_dict'])
optimizer_D_B.load_state_dict(checkpoint['optimizer_D_B_state_dict'])
start_epoch = checkpoint['epoch']
start_iteration = checkpoint['iteration']

# 继续训练
model.train(start_epoch, start_iteration, dataloader, num_epochs=100)

其中，`models`和`datasets`是自定义的模型和数据集，需要根据具体情况进行更改。`CycleGAN`是一个自定义的CycleGAN模型，包含训练函数`train`。在恢复训练时，需要加载保存的模型权重和优化器状态，并设置开始的epoch和iteration。最后调用`train`函数，继续训练模型。