前两次的训练都不满意，用mmgen库做通用和线稿的，效果看起来都不好，我怀疑是不是mmgen不行。

我重新找了一个仓库来做这个事情。

仓库

[email protected]:junyanz/pytorch-CycleGAN-and-pix2pix.git

环境

docker run --gpus all --shm-size=32g -it --net host \
-v ./pytorch-CycleGAN-and-pix2pix/:/pytorch-CycleGAN-and-pix2pix/ \
-v ./anime-sketch-colorization-pair/data:/data \
pytorch/pytorch:2.5.1-cuda12.1-cudnn9-devel bash

pip install -r requirements.txt

已经保存为这个镜像：

kevinchina/deeplearning:2.5.1-cuda12.1-cudnn9-devel-pix2pix

可以直接使用此环境：

docker run --gpus all --shm-size=32g -it --net host \
-v ./pytorch-CycleGAN-and-pix2pix/:/pytorch-CycleGAN-and-pix2pix/ \
-v ./anime-sketch-colorization-pair/data:/data \
kevinchina/deeplearning:2.5.1-cuda12.1-cudnn9-devel-pix2pix bash

训练

先启动一个visdom服务器：

python -m visdom.server -p 8097

要么这么开始训练：

python train.py --dataroot /data/ --name anime --gpu_ids 0,1,2,3 --model pix2pix --direction BtoA --batch_size 192 --lr_policy cosine --num_threads 32 --init_type xavier --norm instance --netG unet_256 

或者用visdom服务器训练：

-display_server http://remote_ip --display_port 8097

浏览器访问8097的服务就可以看到训练效果：

200轮的训练，看起来还是没那么优秀：

上色模型

仔细阅读一些相关信息，才发现之前黑白图转色彩图可以这么训练，输入单通道黑白图，输出为LAB空间图的AB通道。

这么开始训练：

docker run --gpus all --shm-size=32g -it --net host \
-v ./pytorch-CycleGAN-and-pix2pix/:/pytorch-CycleGAN-and-pix2pix/ \
-v ./unpaired_self_datasets/:/data \
kevinchina/deeplearning:2.5.1-cuda12.1-cudnn9-devel-pix2pix bash

nohup python -m visdom.server -p 8097 &

pip install -U scikit-image

python train.py --dataroot /data/ --name tongyong_l2ab --gpu_ids 0,1,2,3 --model colorization --direction AtoB --batch_size 192 --lr_policy cosine --num_threads 32 --init_type xavier --norm instance --netG unet_256 --dataset_mode colorization --input_nc 1 --output_nc 2 --phase trainA

还是不错的：