工设里世界

今天完成了简单的stylegan自定义数据集训练过程，感谢国外老哥的教程，原来很简单……其实对于专门做人工智能方面的人来说训练应该是很简单的，但是对于我这种跨专业跨领域的小白来说，还是需要一个详细教程的，这里给大家介绍一下如何使用自己的数据集来测试stylegan。

首先，大家需要在自己的电脑上配置好tensorflow环境，注意：tensorflow2.0可能无法运行stylegan的原始代码，建议使用1.13版本。

个人运行环境：WIN10+CUDA10.0 +CUDNN7.5+ TENSORFLOW-GPU1.13.1 + python3.7

环境配置教程请先阅读参照之前我写的两篇stylegan文章：

WIN10+CUDA10 +CUDNN7.5+ TENSORFLOW-GPU1.13.1 + python3.7 运行NVIDIA STYLEGAN 的安装过程和踩坑实录

如何使用 NVIDIA StyleGAN 生成自己的动漫（老婆）头像

NVIDIA StyleGan ：https://github.com/NVlabs/stylegan

NVIDIA StyleGan论文地址：https://arxiv.org/pdf/1812.04948.pdf

本文参照的国外的一个教程：

视频：https://www.youtube.com/watch?v=QBUFrep8ljI&t=194s

原文档：https://evigio.com/post/how-to-use-custom-datasets-with-stylegan-tensorFlow-implementation

首先开始时，默认你已经装好了上述环境配置，所以我们直接从这一步开始。

• 第一步：转换数据集为TFRecords格式

首先，将自己的自定义数据集命名为：dataset文件夹，并且将其放在StyleGAN的主目录

现在，我们需要将这些图像转换为TFRecords。而stylegan自带了一个转换TFRecords数据集的工具，也就是stylegan目录下的dataset_tool.py程序，运行它以转换数据集：

python dataset_tool.py create_from_images datasets/LOCATION_TO_SAVE_TF_RECORDS（此处是你要存放生成的TFRecords的文件目录） ./NAME_OF_YOUR_DATASET_DIRECTORY（此处修改为你存放数据集照片的文件夹路径）

例如：我的数据集在zhaopian文件夹下，生成TFRecords文件放在datasets文件夹下的custom_dataset文件夹下，输入的代码如下

python dataset_tool.py create_from_images datasets/custom_dataset ./zhaopian

运行之后，出现如图所示的载入信息

完成后会在stylegan目录下生成一个custom_dataset文件夹，内部应该有好几个TFRecords文件，如下图所示：

这就完成了打包图片数据集的过程~~~

• 第二步：修改train.py文件

之后需要修改训练的python文件：train.py

其中标红的这两行需要修改，首先

    # Dataset.
    desc += '-custom_dataset';     dataset = EasyDict(tfrecord_dir='custom_dataset', resolution=128); train.mirror_augment = False
    #desc += '-ffhq';     dataset = EasyDict(tfrecord_dir='ffhq');                 train.mirror_augment = True
    #desc += '-ffhq512';  dataset = EasyDict(tfrecord_dir='ffhq', resolution=512); train.mirror_augment = True
    #desc += '-ffhq256';  dataset = EasyDict(tfrecord_dir='ffhq', resolution=256); train.mirror_augment = True
    #desc += '-celebahq'; dataset = EasyDict(tfrecord_dir='celebahq');             train.mirror_augment = True
    #desc += '-bedroom';  dataset = EasyDict(tfrecord_dir='lsun-bedroom-full');    train.mirror_augment = False
    #desc += '-car';      dataset = EasyDict(tfrecord_dir='lsun-car-512x384');     train.mirror_augment = False
    #desc += '-cat';      dataset = EasyDict(tfrecord_dir='lsun-cat-full');        train.mirror_augment = False

    # Number of GPUs.
    desc += '-1gpu'; submit_config.num_gpus = 1; sched.minibatch_base = 4; sched.minibatch_dict = {4: 128, 8: 128, 16: 128, 32: 64, 64: 32, 128: 16, 256: 8, 512: 4}
    #desc += '-2gpu'; submit_config.num_gpus = 2; sched.minibatch_base = 8; sched.minibatch_dict = {4: 256, 8: 256, 16: 128, 32: 64, 64: 32, 128: 16, 256: 8}
    #desc += '-4gpu'; submit_config.num_gpus = 4; sched.minibatch_base = 16; sched.minibatch_dict = {4: 512, 8: 256, 16: 128, 32: 64, 64: 32, 128: 16}
    #desc += '-8gpu'; submit_config.num_gpus = 8; sched.minibatch_base = 32; sched.minibatch_dict = {4: 512, 8: 256, 16: 128, 32: 64, 64: 32}

修改完成之后，保存。

• 第三步：执行训练

在anaconda prompt中执行

python train.py

会出现这么一堆代码

Building TensorFlow graph...

WARNING: The TensorFlow contrib module will not be included in TensorFlow 2.0.
For more information, please see:
* https://github.com/tensorflow/community/blob/master/rfcs/20180907-contrib-sunset.md
* https://github.com/tensorflow/addons
If you depend on functionality not listed there, please file an issue.

Setting up snapshot image grid...
Setting up run dir...
Training...

出现这个training...的时候，恭喜你，就代表正在训练啦！

官方文档表示：默认情况下，train.py配置为使用8个GPU以1024×1024分辨率为FFHQ数据集训练最高质量的StyleGAN（表1中的配置F）。请注意，我们在所有实验中都使用了8个GPU。使用较少GPU进行培训可能无法产生相同的结果 - 如果您希望与我们的技术进行比较，我们强烈建议您使用相同数量的GPU。

使用Tesla V100 GPU的默认配置的预计培训时间：

所以你懂的……这个训练还是比较耗费计算资源的……

接下来：

原文：个人理解，当出现这行代码的时候代表运行完成……（但是理解错了！这里修改一下）

注1.（评论中XI同学指正）：

评论里有朋友指正：

“tick1 是跑了记录器的第一步 第一步会比较慢，结束之后会下载个评估模型做反馈。之后就开始正式跑啦。
train loop.py里面官方设置的是140000图tick1次，所以是140.3k image。 lod4.00 是指你现在跑到2^4=16分辨率，也是你在train.py里设置的initial resolution。之后每600tick会+1 变成32 64 128直到1024.”

注2.（评论中MessiXiang同学指正）：

作者，我来揪错了，您说的lod指2的次方数，似乎是不对的，我在学习过程中，lod在逐渐减少，生成的fake图片越来越清晰，说明lod应该是fake与true的区别指数。

20200205补充：作者不是本专业，有点懵，你们辩证看待哈orz

tick 1 kimg 140.3 lod 4.00 minibatch 128 time 4m 59s sec/tick 254.7 sec/kimg 1.82 maintenance 44.5 gpumem 3.0

但是有个奇怪的一点，就是当数据集训练完成的时候，会进行评估，需要下载inception_v3_features.pkl文件。

Downloading https://drive.google.com/uc?id=1MzTY44rLToO5APn8TZmfR7_ENSe5aZUn ............ failed
Traceback (most recent call last):
File "C:\ProgramData\Anaconda3\envs\tensorflow-gpu\lib\site-packages\urllib3\connection.py", line 159, in _new_conn
(self._dns_host, self.port), self.timeout, **extra_kw)
File "C:\ProgramData\Anaconda3\envs\tensorflow-gpu\lib\site-packages\urllib3\util\connection.py", line 80, in create_connection
raise err
File "C:\ProgramData\Anaconda3\envs\tensorflow-gpu\lib\site-packages\urllib3\util\connection.py", line 70, in create_connection
sock.connect(sa)
TimeoutError: [WinError 10060] 由于连接方在一段时间后没有正确答复或连接的主机没有反应，连接尝试失败。

然后运行了这个……看了下这个链接是下载inception_v3_features.pkl这个文件的，当然大家也懂得谷歌这个网址肯定是链接不上的，可以直接关闭，修改pretrained_example.py

20190930修改，评论区有朋友指出，“数据训练完并不是执行预训练模型，而是去评估。选的metric方法是Frechet Inception Distance，所以才取下模型”这里是我理解错了，纠正一下，谢谢指正！

• 第四步：修改pretrained_example.py文件

打开根目录下的pretrained_example.py文件

    # 需要修改的内容
    url = 'https://drive.google.com/uc?id=1MEGjdvVpUsu1jB4zrXZN7Y4kBBOzizDQ' # karras2019stylegan-ffhq-1024x1024.pkl
    with dnnlib.util.open_url(url, cache_dir=config.cache_dir) as f:
        _G, _D, Gs = pickle.load(f)

    # 改成这个内容
    url = os.path.abspath("results/00000-sgan-custom-dataset-4gpu/network-snapshot-000140.pkl")
    with open(url, 'rb') as f:
        _G, _D, Gs = pickle.load(f)url = os.path.abspath("results/00000-sgan-custom-dataset-4gpu/YOUR_LAST_NETWORK_SNAPSHOT")

注意！这个需要根据自己的训练结果修改的！！！不能复制哈

url = os.path.abspath("results/00000-sgan-custom-dataset-4gpu（这个是reaults文件夹下，你刚刚训练好的模型文件的目录名）/network-snapshot-000140.pkl（这个名字要改成前面那个文件夹下的PKL文件的名字）")

和开始的两篇文章一样，

在anaconda prompt中执行

python pretrained_example.py

这样训练好的网络就会运行，并且在result文件夹下生成照片了。

本人也是简单测试这个过程，训练好的模型非常简陋- -下一步的工作就是要对模型进行深入研究了。万里长征第一步，走的不容易呀

祝愿大家成功训练出自己的stylegan啦！

隐者艾伦

2019年9月16日