CVPR 2020 风格迁移之NICE-GAN

🔔 基础信息

• Reusing Discriminators for Encoding: Towards Unsupervised Image-to-Image Translation
• 《重用判别器进行编码：实现无监督的图像到图像翻译》
• https://arxiv.org/pdf/2003.00273.pdf
• 论文精简翻译
• 官方 PyTorch 实现 -- https://github.com/alpc91/NICE-GAN-pytorch

🎉 git 官方 环境依赖如下：

gitHub 官方

📢 环境搭建

• 服务器：ubuntu1~18.04 Quadro RTX 5000 16G
• CUDA版本 V10.0.130

conda create -n nice python=3.6.9

conda activate nice

pip install torch==1.1.0

pip install torchvision==0.3.0
或者（建议使用 conda 来安装）
conda install pytorch==1.1.0 torchvision==0.3.0 cudatoolkit=10.0 -c pytorch

pip install pillow==5.2.0

pip install opencv-python

pip install scipy
pip install thop

🎄 项目结构

 git clone https://github.com/alpc91/NICE-GAN-pytorch.git
 
 cd NICE-GAN-pytorch
 
 cd dataset/

💜 加入数据集

💜 训练参数设置

🎄 训练和测试

💜 Train

# 显卡16G不够用，所以设置  --light True 

python main.py --dataset horse2zebra --light True

# 指定多卡训练的方式如下

CUDA_VISIBLE_DEVICES=2,3 python main.py --dataset horse2zebra --light True

# nohup 把训练进程放到后台

nohup python main.py --dataset horse2zebra --light True &

90分钟 训练 5000 个epoch ，训练效果如下

💜 Test

python main.py --dataset horse2zebra --phase test  --light True

测试运行输出如下

python main.py --dataset horse2zebra --phase test  --light True

# 测试运行输出如下
 
number of epochs must be larger than or equal to one

##### Information #####
# light :  True
# dataset :  horse2zebra
# batch_size :  1
# iteration per epoch :  300000
# the size of image :  256
# the size of image channel :  3
# base channel number per layer :  64

##### Generator #####
# residual blocks :  6

##### Discriminator #####
# discriminator layers :  7

##### Weight #####
# adv_weight :  1
# cycle_weight :  10
# recon_weight :  10
-----------------------------------------------
[INFO] Register count_linear() for <class 'torch.nn.modules.linear.Linear'>.
[INFO] Register count_convNd() for <class 'torch.nn.modules.conv.Conv2d'>.
[INFO] Register count_relu() for <class 'torch.nn.modules.activation.LeakyReLU'>.
[Network disA] Total number of parameters:  93.749M
[Network disA] Total number of FLOPs:  12.013G
-----------------------------------------------
[INFO] Register zero_ops() for <class 'torch.nn.modules.activation.ReLU'>.
[INFO] Register count_linear() for <class 'torch.nn.modules.linear.Linear'>.

[Network gen2B] Total number of parameters:  16.192M
[Network gen2B] Total number of FLOPs:  67.552G
-----------------------------------------------
5000 [测试使用模型的迭代次数]

dataset/horse2zebra/testA/n02381460_1000.jpg
dataset/horse2zebra/testA/n02381460_1010.jpg
...

💬 备注

• 该模型训练中一次保存，大小 1.2 个G
• --batch_size == 1 时，GPU 占用 7679MiB
• --batch_size == 2 时，GPU 占用 12473MiB
• --batch_size == 4 时，16G 的GPU就不够用啦

👋 可能遇到的报错

单卡训练，NICE-GAN full version 训练， 【16G】GPU 不够用

• CUDA_VISIBLE_DEVICES=3 python main.py --dataset horse2zebra

RuntimeError: CUDA out of memory. Tried to allocate 1024.00 MiB (GPU 0; 15.75 GiB total capacity; 13.78 GiB already allocated; 782.44 MiB free; 499.74 MiB cached)

解决方法如下

• 经验证换用 24 G 以上更大内存的 GPU 进行训练，不会遇到该错误，即进行【NICE-GAN full version】训练
• 训练时，追加参数 --light True 即进行【NICE-GAN light version】训练

🚀 给学妹换个风格 🚀

这里使用 【梵高油画和现实风景图像】构成的数据集进行 50万 iteration 训练

nohup python main.py --dataset vangogh2photo --light True &

# 训练时长，最终定格在 5天

训练结束输出如下

[499995/500000] time: 431994.8614 d_loss: 2.32049370, g_loss: 3.90420341
[499996/500000] time: 431995.4892 d_loss: 2.04860115, g_loss: 3.71788430
[499997/500000] time: 431996.1223 d_loss: 2.22972393, g_loss: 4.25117588
[499998/500000] time: 431996.7609 d_loss: 2.36788273, g_loss: 5.42507362
[499999/500000] time: 431997.3897 d_loss: 2.53704023, g_loss: 4.05708218
[500000/500000] time: 431998.0172 d_loss: 2.21194649, g_loss: 4.53932619
current D_learning rate:-1.7075100951256253e-16
current G_learning rate:-1.7075100951256253e-16
 [*] Training finished!

💜 训练分析

可以看到，默认每隔10w iteration 保存一次模型，50w 训练只保存了 5次模型 这 5 次保留的模型经测试发现，效果都无法让人心动

---

意外翻车，原因分析如下

• Gan 网络的模型训练本身就有很多不稳定性、不确定性存在
• 改进策略，数据集 trainA 和 trainB 风格尽量一致
• 训练时，建议每隔 1 w 或者 5k iteration 保存一次模型
• 这样可以得到多个生成模型，然后一一测试，也许会有惊喜
• 对于 NiceGAN 而言，实测应用，模型训练基本是这样一个思路
• 只能说，iteration 次数越多，理论上训练效果相对会更好，但是实际应用中，效果好的模型未必就是训练次数最大的那次得到的模型

💜 学妹还是有的

论数据集风格统一的重要性

---

【❤️下次，一定，让你更美❤️】

🍖 训练代码+数据分享

💜【把训练数据 + 代码分享给大家】💜

************************************************************

🚀🚀 文末专栏推荐部分 🚀🚀

🎄如果感觉文章看完了不过瘾，还想更进一步，那么可以来我的其他 专栏 看一下哦~

• ❤️ 图像风格转换 —— 代码环境搭建 实战教程【关注即可阅】！
• 💜 图像修复-代码环境搭建-知识总结 实战教程 【据说还行】
• 💙 超分重建-代码环境搭建-知识总结 解秘如何让白月光更清晰【脱单神器】
• 💛 YOLO专栏，只有实战，不讲道理 图像分类【建议收藏】！

• 🎄 个人整理的Cuda系列 Linux安装教程【适合小白进阶】

• 💜 ubuntu18给当前用户安装cuda11.2 图文教程 | 配置cuDNN8.1 |
• 💜 Linux服务器下给当前用户安装自己的CUDA10.0
• 💜 Linux下cuda10.0安装Pytorch和Torchvision
• 💜 Linux 可以安装多个版本的Cuda 吗 | 给我一台新的服务器，我会怎么安排 Cuda
• 💜 查看CUDA和cuDNN的版本号

❤️ 当生成器和判别器趋于平衡，那便是爱情 💜

博主简介：墨理，2020年 硕士毕业，目前从事图像算法，AI工程化 相关工作

• 🍊 计算机视觉：超分重建、图像修复、目标检测、风格迁移 等领域 稍有所学
• 🍊 AI 工程化：Ncnn、MNN、TensorRT 正在 学习
• 🍊 C++、Python、Java 略懂一二
• 🍊 全网 10 w 读者粉丝
• 🍊 博文主要创作平台：CSDN + 公--中--号 + 掘金 + 知乎

喜欢请关注 墨理学AI 及其 同名 公---中--- 号 墨理学AI

取经路上，让墨理学AI 陪你畅享更多有趣AI