如何使用stylegan训练自己的照片数据集

发表于 2019年9月16日来自隐者艾伦 in 人工智能, 31 comments

今天完成了简单的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

#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×1024	512×512	256×256
1	41天4小时	24天21小时	14天22小时
2	21天22小时	13天7小时	9天5小时
4	11天8小时	7天0小时	4天21小时
8	6天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 = '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日

看过: 7,636

标签：ai, GAN, nvidia stylegan, stylegan, 人工智能, 生成对抗神经网络

31 个评论

carp
％1 $ S
关于执行训练中的下载失败问题，gwern有说明：https://www.gwern.net/Faces#training
Some metrics are computed for logging/reporting. The FID metrics are calculated using an old ImageNet CNN; what is realistic on ImageNet may have little to do with your particular domain and while a large FID like 100 is concerning, FIDs like 20 or even increasing are not necessarily a problem or useful guidance compared to just looking at the generated samples or the loss curves. Given that computing FID metrics is not free & potentially irrelevant or misleading on many image domains, I suggest disabling them entirely. (They are not used in the training for anything, and disabling them is safe.)They can be edited out of the main training loop by commenting out the call to metrics.
也就是说直接注释掉metrics.run(就行了。
回复
一只咖喱
％1 $ S
写得太好了，呜呜呜，早点看到就好了，之前自己在csdn上摸了好久好久才搞明白，这种保姆级教程真的少，对小白来说真的很重要！赞赞赞
回复
- 隐者艾伦
  ％1 $ S
  不客气，我也是新手，最懂新手一脸懵比的痛苦……
  回复
shisheng
％1 $ S
群主你好，请问你训练的是128的图片吗？
回复
- 隐者艾伦
  ％1 $ S
  我的电脑比较糟糕，就没有训练自己的数据集~~只是尝试跑通了一下方法
  回复
zwcccr
％1 $ S
博主你好，我是小白一枚，对style gan还不是很了解。百度的AI studio 最近每天免费送10小时的Tesla V100 GPU使用时间，请问这样一天一运行可不可以完成训练呢
回复
- 隐者艾伦
  ％1 $ S
  主要还是取决于你数据量的大小，能不能完成咱也有点不确定……
  回复
zwcccr
％1 $ S
博主你好，卧室小白一枚。百度的AI studio 最近每天免费送10小时的Tesla V100 GPU使用时间，请问这样一天一运行可不可以完成训练呢
回复
Li
％1 $ S
用V100训练了大概4天还是部分图片质量不行，就写了个CNN来过滤生成的图片。我把要用到的pkl文件缓存下来了，大家可以随意多线程下载：https://ai.inf.link/c/stylegan2/
回复
东坡
％1 $ S
解决了吗，方便告诉我一下方法吗
回复
MessiXiang
％1 $ S
作者，我来揪错了，您说的lod指2的次方数，似乎是不对的，我在学习过程中，lod在逐渐减少，生成的fake图片越来越清晰，说明lod应该是fake与true的区别指数。
回复
- SPL
  ％1 $ S
  lod = Level of Details
  是表示生成圖像的仔細度
  數字越細代表圖像越清 (解像度越高)
  回复
MessiXiang
％1 $ S
楼主······修改pretrained文件那步，运行后总说IndentationError: expected an indented block·······怎么办
回复
- SPL
  ％1 $ S
  Python 縮排錯誤, 檢查清楚縮排時按空白鍵的數目
  eg:
  def name():
  x = 5
  y = 6 (縮排錯誤)
  回复
XI
％1 $ S
“个人理解，当出现这行代码的时候代表运行完成……
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”
博主你这个说错了 tick1 是跑了记录器的第一步第一步会比较慢，结束之后会下载个评估模型做反馈。之后就开始正式跑啦。
train loop.py里面官方设置的是140000图tick1次，所以是140.3k image。 lod4.00 是指你现在跑到2^4=16分辨率，也是你在train.py里设置的initial resolution。之后每600tick会+1 变成32 64 128直到1024.
回复
- 隐者艾伦
  ％1 $ S
  感谢指正！
  回复
Salt
％1 $ S
請問一下,您訓練模型大概花了多少時間？
回复
- 隐者艾伦
  ％1 $ S
  你好，因为手里只有1070单卡，看起来耗时庞大，并没有尝试训练整个数据集，只是把youtube上的教程流程走了一遍，并记录了一下。
  回复
wj
％1 $ S
楼主，如何继续上次中断的训练，我的已经2861了
回复
- 隐者艾伦
  ％1 $ S
  表示不知道orz
  回复
- MessiXiang
  ％1 $ S
  修改training loop.py有几个resume选项，填好之后运行train.py就可以了
  回复
6356
％1 $ S
博主你好，stylegan训练自己的模型，单张gpu就可以运行，多张gpu就报错，请问你知道怎么回事吗？
InvalidArgumentError (see above for traceback): No OpKernel was registered to support Op ‘NcclAllReduce’ used by node TrainD/SumAcrossGPUs/NcclAllReduce (defined at C:\Users\B\stylegan-master\dnnlib\tflib\optimizer.py:135) with these attrs: [reduction=”sum”, T=DT_FLOAT, num_devices=2, shared_name=”c124″]
Registered devices: [CPU, GPU]
Registered kernels:
[[node TrainD/SumAcrossGPUs/NcclAllReduce (defined at C:\Users\B\stylegan-master\dnnlib\tflib\optimizer.py:135) ]]
回复
- 隐者艾伦
  ％1 $ S
  emmm，训练过程中的这个多卡设置好了么了~ 我没有试过多卡环境，不知道多卡是不是还需要在其他地方进行设置呢？
  # 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}
  回复
NEXTINtel
％1 $ S
博主你好,我是一个小白,在执行第三步时遇到了以下问题,请问博主在执行时是否也遇到了类似的问题呢?
Streaming data using training.dataset.TFRecordDataset…
Traceback (most recent call last):
File “train.py”, line 191, in
main()
File “train.py”, line 186, in main
dnnlib.submit_run(**kwargs)
File “/content/drive/stylegan-acg/dnnlib/submission/submit.py”, line 290, in submit_run
run_wrapper(submit_config)
File “/content/drive/stylegan-acg/dnnlib/submission/submit.py”, line 242, in run_wrapper
util.call_func_by_name(func_name=submit_config.run_func_name, submit_config=submit_config, **submit_config.run_func_kwargs)
File “/content/drive/stylegan-acg/dnnlib/util.py”, line 257, in call_func_by_name
return func_obj(*args, **kwargs)
File “/content/drive/stylegan-acg/training/training_loop.py”, line 146, in training_loop
training_set = dataset.load_dataset(data_dir=config.data_dir, verbose=True, **dataset_args)
File “/content/drive/stylegan-acg/training/dataset.py”, line 234, in load_dataset
dataset = dnnlib.util.get_obj_by_name(class_name)(**adjusted_kwargs)
File “/content/drive/stylegan-acg/training/dataset.py”, line 70, in __init__
assert os.path.isdir(self.tfrecord_dir)
AssertionError
回复
- 隐者艾伦
  ％1 $ S
  不好意思我也没有遇到过，只知道这个AssertionError说是断言语句失败，但不知道是什么问题orz
  回复
- 无人
  ％1 $ S
  断言就是当括号内表达式为假时中断程序，这段表示你的数据集路径不存在，你可以检查一下你是否正确设置好了数据集名称，或者你是否有数据集（还有一个就是train.py有一段代码是包在if 0后面的，相当于注释了）
  回复
  - SPL
    ％1 $ S
    本人有此經驗, 查清楚自己的數據庫是否:
    1: 全部都是 jpg 檔案
    2: 相同大小
    您可以刪除 assertion 這一行, 但是之後會有另外一些error
    不建議這樣解決
    回复
- xrh
  ％1 $ S
  请问你解决了吗，我找不到关于我数据集的问题
  回复
阿达里西
％1 $ S
数据训练完并不是执行预训练模型，而是去评估。选的metric方法是Frechet Inception Distance，所以才取下模型。楼主知道怎么增大tick吗？140太少了，
回复
- 隐者艾伦
  ％1 $ S
  感谢反馈！楼主也是一头雾水，对于这个参数的配置不是很熟悉哈。。。还在继续研究
  回复
- gqy
  ％1 $ S
  直接用 –tick=10 就可以了
  回复