【GAN】pytorch-CycleGAN-and-pix2pix开源项目训练自定义数据集

Github：https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix

一、概念简介

图像转换：将一张图片转换为不同风格的图片，如将一张灰度图变成彩色图，素描图换成实物图。

pix2pix：需要成对的数据集，两个不同域的图像一一对应。

CycleGAN：不需要成对的数据集，只用两个不同域（domain）的图像集合X和Y，两个集合中的图像不是一一对应的。

二、训练pix2pix模型

2.1 combine_A_and_B

datasets/combine_A_and_B.py

将A和B不同场景域下相同场景进行拼接，将A和B拼接成AB。

|--A
|  |--train
|  |--val
|
| --B
|  |--train
|  |--val
|
| --test_AB（生成的）
|  |--train
|  |--val

修改A、B、AB三个目录：

并把主函数放在main中运行，否则可能会报错。

合成后的数据集如下图：

2.2 base_options

通常要修改的内容包含：

--dataroot：数据集的目录，设置为刚刚拼接后的目录../data/test_AB
--use_wandb：设置wandb，填入default=False设置不适用wandb；
--model：选择模型pix2pix ，可以选择cycle_gan | pix2pix | test | colorization；
--input_nc：输入通道数，3是RGB，1是灰度
--output_nc：输出通道数，3是RGB，1是灰度
--dataset_mode：数据模式，unaligned | aligned | single | colorization，pix2pix 的默认会设置为aligned，修不修改都可以。
--direction：数据生成方向，是A生成B，还是B生成A

2.3 train_options

包含了保存模型的频率和训练过程的参数设置，根据情况进行设置，就不一一列举了。

2.4 训练

启动visdom：python -m visdom.server

打开浏览器： http://localhost:8097

运行train.py进行训练。