使用StyleGAN2生成动漫头像(一)

本来在StyleGAN2发布的当天就准备弄来玩玩的,但是当时官方给的pkl文件托管在Google Drive上面,后来因为访问频率过高被暂时封掉。而且自己本身的算力条件有限制(Colab似乎跑太久就会被断开),所以没有训练几个epoch就放弃了,最后效果也没怎么样。

这次有机会长时间使用Tesla显卡来跑,准备重新训练一下模型,这里记录一下过程和可能遇到的问题。

项目地址:StyleGAN2(我的Fork)

数据收集和处理

爬取图片

我们直接从Pixiv上面抓取数据,不同画师的水平可能不一,所以最好选一些收藏数量比较高的图片。这里我租用了一个带有200G HDD的日本VPS,方便存储原图,否则可能传回数据过慢或者存不下。

爬Pixiv图片的工具有很多,可以选择Chrome插件,或者自己利用PixivPy写个脚本。

提取面部

从网上找到的办法,大致看了一下,基本上也是利用opencv级联分类器来做的。这里我直接把源码和xml文件都上传了,查看一下anime_开头的两个文件都知道如何使用。

感觉这个的精度不能算特别高,而且会有一些非面部的图片被检测到,这里你可以将他们删除掉。

需要注意的是,我们需要生成的图片必须满足像素是2的次方,所以我们选择了128x128,你也可以选择256×256或者1024×1024。最后我们可能需要利用Waifu2X来扩展像素。

放个部分数据的截图~

制作数据集

使用官方提供的dataset_tool即可,这里提供了生成tfrecords文件和extract出原图的操作,非常方便。

训练模型

安装依赖

这里我们使用的是Ubuntu系统,conda作为管理工具。根据我的测试,你用Windows来做也行,甚至更为方便。但是我还是按照Linux来展开吧。

官方要求的是Python 3.6,我们发现似乎Python 3.7方便一些。而TensorFlow 1.15使用conda安装的时候可能会出现版本兼容错误,所以我们使用如下命令搭建好环境。(提前把conda换成清华源有益身心健康)

1
2
3
4
5
conda create -n tf python=3.6
#conda init bash
conda activate tf
conda install tensorflow==1.14 tensorflow-gpu==1.14 cudatoolkit==9.0 cudnn #先别复制,注意接下来的说明。
pip install requests pillow

这里请选择你显卡合适的cudatoolkit版本,Tesla对应9.0,而RTX系列似乎可以用10.1。否则会出现兼容错误。然后tensorflow和tensorflow-gpu顺序不能搞反,否则出现无法找到GPU。

显卡的显存不能太低,实际占用会超过9G。

开始训练

训练的参数文档已经写得很清楚了,你也可以选择--mirror-augment=true来extend数据集。所以我这里给出可能遇到的问题,以及解决方案。

  • 下载pkl文件过慢
    官方一开始用的是Google Drive,后来换成了CDN,但是效果还是不佳。你这里可以先从inf.link提供的StyleGAN2缓存下载下来,然后修改metrics文件夹中的源码,把URL换成你的本地路径。在默认参数的条件下,仅需要inception_v3_features.pkl。

  • “C++ versions less than C++11 are not supported.”
    C++标准支持出错,这里修改dnnlib>tflib>custom_ops.py,在编译的参数中加上--std=c++11 -DNDEBUG

  • “undefined symbol: _ZN10tensorflow12OpDefBuilder6OutputESs”
    编译问题,这里修改dnnlib>tflib>custom_ops.py,把-D_GLIBCXX_USE_CXX11_ABI=0修改为1。

训练途中

这里测试的是每18min可以训练12kimg,具体时间受制于你的参数还有配置条件。每10 tick会生成一次snapshot和当前网络生成的图片示例。因为我目前还在训练,所以这里给几张过程中的截图,这是训练到1000kimg的结果。

写在最后

本文还是会持续更新的,这里还是因为算力的原因,我准备只训练4天左右。但是已经足以看出StyleGAN2的能力了。后面有机会给大家讲讲神经网络、迁移学习方面的东西。感谢大家支持。

  • 版权声明: 本博客所有文章除特别声明外,均采用 CC BY-NC-SA 许可协议。转载需要标明作者,并以超链接的方式指向原文。
  • Copyrights © 2020 Kevin Li
  • 访问人数: | 浏览次数:

请我喝杯咖啡吧~