【模型复现】自制数据集上复现风格迁移模型 CUT:Contrastive Unpaired Translation
Code 链接: Contrastive Unpaired Translation (CUT)
Paper 链接: Contrastive Learning for Unpaired Image-to-Image Translation
1. 环境安装
# 创建环境
conda create -n mycut python=3.9
# 激活环境
conda activate mycut
# torch 安装
# 本机 CUDA 为 11.8,故安装了符合要求的 pytorch==1.13,这里需要自行根据 CUDA 版本安装适配的 torch 版本
pip install torch==1.13.1+cu117 torchvision==0.14.1+cu117 torchaudio==0.13.1 --extra-index-url https://download.pytorch.org/whl/cu117
# clone 代码
git clone https://github.com/taesungp/contrastive-unpaired-translation.git
# pip 包
cd contrastive-unpaired-translation
pip install -r requirements.txt
2. 数据准备
- 准备训练数据,数据集的文件结构为:
my_datasets/ ├──trainA │ ├──xxx.jpg │ ├──xxx.jpg └──trainB ├──xxx.jpg └──xxx.jpg
trainA
是在同一分布内的源域数据(如 CityScapes 数据集)trainB
是在另一种分布内的目标域数据(如 CityScapes Foggy 数据集)
3. 模型训练
3.1 训练超参数配置
- 通过
options/base_options.py
进行基础超参数配置,按需进行超参数配置。
- 通过
options/train_options.py
进行训练超参数配置,按需进行超参数配置。
3.2 模型训练
- 训练指令
# 若在脚本中已正确配置所有超参数,直接运行 train.py 即可 python train.py # 若未在脚本中配置超参数,直接运行下述命令也可进行模型训练 python train.py --dataroot ./datasets/mydatasets --name mydatasets_CUT --CUT_mode CUT --gpu_ids 0,1,2,3 --batch_size 4 --preprocess scale_shortside_and_crop --load_size 1024 --crop_size 320 --n_epochs 500 --n_epochs_decay 200
- 在终端中运行训练命令,若看到下述界面,即成功复现!!!
- 同时,在
visdom
中界面显示如下。
4. 模型推理
4.1 推理超参数配置
- 与模型训练同理,通过
options/base_options.py
进行基础超参数配置,按需进行超参数配置。 - 通过
options/test_options.py
进行测试超参数配置,按需进行超参数配置。
4.2 模型推理
- 推理指令
# 若在脚本中已正确配置所有超参数,直接运行 test.py 即可 python test.py # 若未在脚本中配置超参数,直接运行下述命令 python test.py --dataroot ./datasets/mydatasets/ --results_dir ./datasets/mydatasets_test/ --name mydatasets_CUT --CUT_mode CUT --phase train --preprocess none --gpu_ids 0,1,2,3 --num_test 2040
- 推理成功界面如下。
- 同时在设定的
results_dir
路径下可看到目录如下,其中fake_B
即为生成的风格迁移图像。
- 若使用 CityScapes 和 CityScapes Foggy 数据集,则图像生成展示如下:
5. 踩坑记录
5.1【Errno 111】Connection refused,Could not connect to Visdom server.
- 解决方法:
- 在模型训练前安装并开启
visdom
服务,使用python -m visdom.server
或visdom
可开启。此外,visdom
默认端口为 8097,若端口占用会出现下述报错。
- 此时需要指定端口号开启服务,即使用
python -m visdom.server -p xxxx
或visdom -p xxxx
开启(注意!!!指定端口号需与训练脚本中配置的visdom
端口号一致),开启后显示如下:
- 在模型训练前安装并开启
5.2 FileNotFoundError: [Errno 2] No such file or directory: ‘./checkpoints/person_car_CUT/latest_net_G.pth’
- 问题分析:
- 该路径确实存在,但运行时找不到该文件,检查代码发现是由于配置超参数
name
时末尾含空格,导致路径中空格多余而找不到文件。
- 该路径确实存在,但运行时找不到该文件,检查代码发现是由于配置超参数
- 解决方法:
- 在
base_model.py
文件中修改权重加载路径为绝对路径即可开始模型测试推理。 - 代码推理成功界面如下:
- 在