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

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

看过: 5,157

29 个评论

一只咖喱
％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
    不建議這樣解決
    回复
阿达里西
％1 $ S
数据训练完并不是执行预训练模型，而是去评估。选的metric方法是Frechet Inception Distance，所以才取下模型。楼主知道怎么增大tick吗？140太少了，
回复
- 隐者艾伦
  ％1 $ S
  感谢反馈！楼主也是一头雾水，对于这个参数的配置不是很熟悉哈。。。还在继续研究
  回复
- gqy
  ％1 $ S
  直接用 –tick=10 就可以了
  回复