CycleGAN深度学习项目

news2024/11/13 23:25:06

远程仓库

leftthomas/CycleGAN: A PyTorch implementation of CycleGAN based on ICCV 2017 paper "Unpaired Image-to-Image Translation using Cycle-Consistent Adversarial Networks" (github.com)

运行准备

Anaconda

安装需要的库

指令

pip install pandas -i https://pypi.tuna.tsinghua.edu.cn/simple

pip install torch==1.11.0 -i Simple Index

pip install torchvision==0.12.0 -i Simple Index

pip install dominate==2.4.0 -i Simple Index

pip install visdom==0.1.8.8 -i Simple Index

pip install tqdm -i https://pypi.tuna.tsinghua.edu.cn/simple

运行结果

数据集

我当前使用的数据集

leftthomas/CycleGAN: A PyTorch implementation of CycleGAN based on ICCV 2017 paper "Unpaired Image-to-Image Translation using Cycle-Consistent Adversarial Networks" (github.com)

因为数据集太大，训练时间慢所以删掉了很多图片

A-副本和B-副本里面是原始的数据集

A B是我自己删了图片的数据集

如果使用其他数据集也可以训练，例如：从网上随便下载图片

运行结果

程序解读

从main.py的if __name__ == '__main__':开始看

因为程序从这里开始执行

parser = argparse.ArgumentParser(description='Train Model')
# common args
parser.add_argument('--data_root', default='horse2zebra', type=str, help='Dataset root path')
# 文件放的位置
parser.add_argument('--batch_size', default=1, type=int, help='Number of images in each mini-batch')
#每个小批量中的图像数量
parser.add_argument('--epochs', default=2, type=int, help='Number of epochs over the data to train')
# 多少轮训练
parser.add_argument('--lr', default=0.0002, type=float, help='Initial learning rate')
# 开始时学习率
parser.add_argument('--decay', default=2, type=int, help='Epoch to start linearly decaying lr to 0')
# 从第几轮开始学习率逐渐减为0
parser.add_argument('--save_root', default='result', type=str, help='Result saved root path')
# 训练出来的保存在哪里
# args parse
args = parser.parse_args()
data_root, batch_size, epochs, lr = args.data_root, args.batch_size, args.epochs, args.lr
decay, save_root = args.decay, args.save_root

# data prepare
train_data = ImageDataset(data_root, 'train')
# 训练集
print("数据")
print(train_data.__len__())
# 打印出数据集的长度
test_data = ImageDataset(data_root, 'test')
# 验证集
train_loader = DataLoader(train_data, batch_size=batch_size, shuffle=False, num_workers=8)
test_loader = DataLoader(test_data, batch_size=1, shuffle=False, num_workers=8)

使用通义灵码解释

# optimizer setup
optimizer_G = Adam(itertools.chain(G_A.parameters(), G_B.parameters()), lr=lr, betas=(0.5, 0.999))
optimizer_DA = Adam(D_A.parameters(), lr=lr, betas=(0.5, 0.999))
optimizer_DB = Adam(D_B.parameters(), lr=lr, betas=(0.5, 0.999))
lr_scheduler_G = LambdaLR(optimizer_G, lr_lambda=lambda eiter: 1.0 - max(0, eiter - decay) / float(decay))
lr_scheduler_DA = LambdaLR(optimizer_DA, lr_lambda=lambda eiter: 1.0 - max(0, eiter - decay) / float(decay))
lr_scheduler_DB = LambdaLR(optimizer_DB, lr_lambda=lambda eiter: 1.0 - max(0, eiter - decay) / float(decay))

这段代码是用于设置优化器和学习率调度器的。

首先，使用Adam优化器来初始化G_A和G_B的参数以及D_A和D_B的参数。Adam优化器是一种基于梯度的优化算法，它利用了动量和自适应学习率的特性。itertools.chain函数用于将G_A和G_B的参数组合在一起。

然后，使用LambdaLR学习率调度器来设置学习率的衰减。LambdaLR调度器使用给定的函数来计算每个迭代步骤的学习率。这里使用了一个lambda函数，它在迭代次数eiter超过decay后开始衰减学习率，衰减的速度由decay参数控制。

这些优化器和学习率调度器将用于训练生成器和判别器模型。

# training loop
results = {'train_g_loss': [], 'train_da_loss': [], 'train_db_loss': []}
if not os.path.exists(save_root):
    os.makedirs(save_root)
for epoch in range(1, epochs + 1):
    g_loss, da_loss, db_loss = train(G_A, G_B, D_A, D_B, train_loader, optimizer_G, optimizer_DA, optimizer_DB)
    results['train_g_loss'].append(g_loss)
    results['train_da_loss'].append(da_loss)
    results['train_db_loss'].append(db_loss)
    val(G_A, G_B, test_loader)
    lr_scheduler_G.step()
    lr_scheduler_DA.step()
    lr_scheduler_DB.step()
    # save statistics
    data_frame = pd.DataFrame(data=results, index=range(1, epoch + 1))
    data_frame.to_csv('{}/results.csv'.format(save_root), index_label='epoch')
    torch.save(G_A.state_dict(), '{}/GA.pth'.format(save_root))
    torch.save(G_B.state_dict(), '{}/GB.pth'.format(save_root))
    torch.save(D_A.state_dict(), '{}/DA.pth'.format(save_root))
    torch.save(D_B.state_dict(), '{}/DB.pth'.format(save_root))