今天完成了简单的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
其中标红的这两行需要修改，首先
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 | # 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的默认配置的预计培训时间：
| 图形处理器 | 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文件
1 2 3 4 5 6 7 8 9 | # 需要修改的内容 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日
18 个评论
Li
用V100训练了大概4天还是部分图片质量不行，就写了个CNN来过滤生成的图片。我把要用到的pkl文件缓存下来了，大家可以随意多线程下载：https://ai.inf.link/c/stylegan2/
东坡
解决了吗，方便告诉我一下方法吗
MessiXiang
作者，我来揪错了，您说的lod指2的次方数，似乎是不对的，我在学习过程中，lod在逐渐减少，生成的fake图片越来越清晰，说明lod应该是fake与true的区别指数。
MessiXiang
楼主······修改pretrained文件那步，运行后总说IndentationError: expected an indented block·······怎么办
XI
“个人理解，当出现这行代码的时候代表运行完成……
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.
隐者艾伦
感谢指正！
Salt
請問一下,您訓練模型大概花了多少時間？
隐者艾伦
你好，因为手里只有1070单卡，看起来耗时庞大，并没有尝试训练整个数据集，只是把youtube上的教程流程走了一遍，并记录了一下。
wj
楼主，如何继续上次中断的训练，我的已经2861了
隐者艾伦
表示不知道orz
MessiXiang
修改training loop.py有几个resume选项，填好之后运行train.py就可以了
6356
博主你好，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) ]]
隐者艾伦
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
博主你好,我是一个小白,在执行第三步时遇到了以下问题,请问博主在执行时是否也遇到了类似的问题呢?
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
隐者艾伦
不好意思我也没有遇到过，只知道这个AssertionError说是断言语句失败，但不知道是什么问题orz
无人
断言就是当括号内表达式为假时中断程序，这段表示你的数据集路径不存在，你可以检查一下你是否正确设置好了数据集名称，或者你是否有数据集（还有一个就是train.py有一段代码是包在if 0后面的，相当于注释了）
阿达里西
数据训练完并不是执行预训练模型，而是去评估。选的metric方法是Frechet Inception Distance，所以才取下模型。楼主知道怎么增大tick吗？140太少了，
隐者艾伦
感谢反馈！楼主也是一头雾水，对于这个参数的配置不是很熟悉哈。。。还在继续研究