你也能用英伟达GAN造脸了:官方实现有了,高清数据集也开源了

时间:2022-06-20
本文章向大家介绍你也能用英伟达GAN造脸了:官方实现有了,高清数据集也开源了,主要内容包括其使用实例、应用技巧、基本知识点总结和需要注意事项,具有一定的参考价值,需要的朋友可以参考一下。
栗子 假装发自 凹非寺 量子位 出品 | 公众号 QbitAI

你大概还没忘记,英伟达去年年底推出的GAN,它合成的人脸甚至骗得过肉眼。

如今,它终于有了自己的名字,叫StyleGAN。顾名思义,GAN的生成器,是借用风格迁移的思路重新发明的。

更重要的是,你现在也能自己养一只这样的GAN了:

官方实现的代码开源了,提供了许多预训练好的模型,自然也支持自己训练模型。

另外,Flicker人像照片的高清数据集 (FFHQ) 也开源了,包含70,000张高清人脸。

官方实现

首先,送上来自英伟达的友情提示:

这个实现可以用Linux跑,也可以用Windows跑,但墙裂推荐大家用Linux跑,为性能和兼容性着想。

除此之外,必需品还有Python 3.6,和TensorFlow 1.10以上 (支持GPU) 。

预训练模型

这里,用pretrained_example.py举个简易的栗子。执行的时候,脚本会从Google Drive下载一个预训练的StyleGAN生成器,然后用它来生成图像:

 1> python pretrained_example.py
 2Downloading https://drive.google.com/uc?id=1MEGjdvVpUsu1jB4zrXZN7Y4kBBOzizDQ .... done
 3
 4Gs                              Params    OutputShape          WeightShape
 5---                             ---       ---                  ---
 6latents_in                      -         (?, 512)             -
 7...
 8images_out                      -         (?, 3, 1024, 1024)   -
 9---                             ---       ---                  ---
10Total                           26219627
11
12> ls results
13example.png # https://drive.google.com/uc?id=1UDLT_zb-rof9kKH0GwiJW_bS9MoZi8oP

预训练的生成器有三种食用方法:

一是Gs.run()快速模式,这里的输入和输出都是numpy阵列:

1# Pick latent vector.
2rnd = np.random.RandomState(5)
3latents = rnd.randn(1, Gs.input_shape[1])
4
5# Generate image.
6fmt = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True)
7images = Gs.run(latents, None, truncation_psi=0.7, randomize_noise=True, output_transform=fmt)

二是用Gs.get_output_for()把生成器整合到一个更大的TensorFlow表达式里面:

1latents = tf.random_normal([self.minibatch_per_gpu] + Gs_clone.input_shape[1:])
2images = Gs_clone.get_output_for(latents, None, is_validation=True, randomize_noise=True)
3images = tflib.convert_images_to_uint8(images)
4result_expr.append(inception_clone.get_output_for(images))

三是查找Gs.components.mapping和Gs.components.synthesis,访问生成器的各个子网络。与Gs相似,这鞋子网络也表示为dnnlib.tflib.Network的独立示例。

1src_latents = np.stack(np.random.RandomState(seed).randn(Gs.input_shape[1]) for seed in src_seeds)
2src_dlatents = Gs.components.mapping.run(src_latents, None) # [seed, layer, component]
3src_images = Gs.components.synthesis.run(src_dlatents, randomize_noise=False, **synthesis_kwargs)

自己训练模型

如果不满足于预训练的模型,想自己训练模型,就要先:

准备好数据集

把数据集储存为多重分辨率的TFRecords,训练和评估脚本都是在这上面跑:

1> python dataset_tool.py create_lsun datasets/lsun-bedroom-full ~/lsun/bedroom_lmdb --resolution 256
2> python dataset_tool.py create_lsun_wide datasets/lsun-car-512x384 ~/lsun/car_lmdb --width 512 --height 384
3> python dataset_tool.py create_lsun datasets/lsun-cat-full ~/lsun/cat_lmdb --resolution 256
4> python dataset_tool.py create_cifar10 datasets/cifar10 ~/cifar10
5> python dataset_tool.py create_from_images datasets/custom-dataset ~/custom-images

数据集表示为一个目录,里面的每张图像都有多种不同的分辨率,用于高效的streaming。每个分辨率都有一个自己的*.tfrecords文件。数据有标注的话,也是用一个分开的文件来储存的。

训练吧

官方提供的训练过程分四步:

1. 编辑train.py,通过取消注释或者修改某些行,来指定数据集和训练配置; 2. 用train.py来运行训练脚本; 3. 结果会写在一个新目录里,叫results/-; 4. 训练直至完成,几天时间可能是要的。

(最好有一台英伟达高端GPU,至少11GB的DRAM,再开始训练。有很多台的话,当然更好。)

至于训练好的模型该怎样评估,大家可以去项目页自行观察。

高清无码数据集

和官方代码实现一起发布的,就是Flickr高清人脸数据集了。

那些几可乱真的人脸,就是StyleGAN吃了这个数据集,才生成的。

数据集里包含7万张1024 x 1024高清人像。英伟达说,这些照片在年龄、种族、以及图片背景上,都有很强的多样性。

并且,眼镜墨镜帽子这些元素,也是应有尽有。

团队说,图像是直接从Flickr上面扒下来,自动对齐自动裁剪而成。并且,数据集里收录的图片都有使用许可,无须担心。

还不去试试

有大胆想法的同学们,可以去试试了。毕竟,不是只有生成人脸这一种功能,猫片,汽车,房间……

你还想生成一些什么?

代码实现传送门: https://github.com/NVlabs/stylegan

FFHQ数据集传送门: https://github.com/NVlabs/ffhq-dataset

论文传送门: https://arxiv.org/abs/1812.04948