今天完成了简单的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}
解释一下需要修改的这一段哈,#代表这行是注释,并不参与运行,如果你需要启用此行代码,就要删除掉# 使得这行代码可用
下面是需要修改的两行解释
# Dataset.
    desc += '-custom_dataset';     dataset = EasyDict(tfrecord_dir='custom_dataset(这个是刚刚打包好的训练集的路径)', resolution=128(这个我理解为图片分辨率)); 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}
这里是取决于你插了几块GPU,如果插了一块就使用第一个代码,如果是多GPU用其他几行

修改完成之后,保存。

  • 第三步:执行训练

在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的默认配置的预计培训时间:

图形处理器1024×1024512×512256×256
141天4小时24天21小时14天22小时
221天22小时13天7小时9天5小时
411天8小时7天0小时4天21小时
86天14小时4天10小时3天8小时

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

接下来:

原文:个人理解,当出现这行代码的时候代表运行完成……(但是理解错了!这里修改一下)

注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日