如何使用自定义数据集训练 StyleGAN2-ADA-了解如何训练 AI 生成您想要的任何图像

gingo 2021-11-18 16:17:33 资料仓库 GAN 收藏

0 / 1848

How to Train StyleGAN2-ADA with Custom Dataset

如何使用自定义数据集训练 StyleGAN2-ADA

Learn how to train an AI to generate any images you want

了解如何训练 AI 生成您想要的任何图像

您是否一直想知道如何训练自己的生成模型？我第一次发现像这个人脸生成网站这样的 GAN 应用程序时，我一直想知道如何在其他事情上训练 GAN。幸运的是，作为我研究的一部分，我最近有机会训练一个自行车生成模型。在本文中，我将记录我如何在您自己的图像上训练 StyleGAN2-ADA 的经验。

Generated Bike Interpolation [Image by Author]Have you always wondered how to train your own generative model? The first time I discovered GAN applications like this face generating website , I always wondered how to train GAN on other things. Luckily, I recently had the opportunity to train a bike generating model as part of my research. In this article, I will document my experience on how to train StyleGAN2-ADA on your own images.

StyleGAN 是NVIDIA最受欢迎的生成模型之一。StlyeGAN 的多个版本已经发布，我们将使用最新版本 StyleGAN2-ADA。为了避免冗余，我不会解释 StyleGAN，因为有很多文章已经很好地解释了它。

StyleGAN is one of the most popular generative models by NVIDIA. Multiple version of StlyeGAN has been released and we will be using the latest version which is StyleGAN2-ADA. To avoid redundancy, I won’t explain StyleGAN as there are many articles that have explained it really well.

GAN — StyleGAN & StyleGAN2

Do you know your style? Most GAN models don’t. In the vanilla GAN, we generate an image from a latent factor z.

jonathan-hui.medium.com

Generating Anime Characters with StyleGAN2

Learn how to generate this cool anime faces interpolation

towardsdatascience.com

训练 StyleGAN 的计算成本很高。因此，如果您没有像样的 GPU，您可能希望在云上进行训练。如果您决定在 Google Colab 上训练（免费），有人为此制作了一个不错的笔记本。

Training StyleGAN is computationally expensive. Hence, if you don’t have a decent GPU, you may want to train on the cloud. If you decide to train on Google Colab (it’s free), someone has made a nice notebook for this.
在本教程中，我将使用自行车数据集BIKED。随意使用您自己的数据集。只要确保所有的训练图像都是方形的，并将它们放在同一个文件夹中。

在本文中，我将使用 StyleGAN2-ADA 的 Tensorflow 实现。确保您使用 Tensorflow 版本 1，因为代码与 Tensorflow 2 不兼容。或者，如果您更喜欢 PyTorch，您可以使用最近发布的 PyTorch 版本。PyTorch 代码的性能似乎稍快一些。如果您使用 PyTorch，您仍然可以遵循本教程，但在数据集准备方面略有不同。

In the tutorial, I will be using the bike dataset BIKED. Feel free to use your own dataset. Just make sure all the training images are square and put them inside the same folder.

In this article, I will be using the Tensorflow implementation of StyleGAN2-ADA. Make sure you use Tensorflow version 1, as the code is not compatible with Tensorflow 2. Alternatively, if you prefer PyTorch, you can use the PyTorch version that has been recently released. The PyTorch code seems to be slightly faster in performance. If you use PyTorch, you can still follow this tutorial with a slight difference in the dataset preparation.

Requirement

64 位 Python 3.6 或 3.7。建议使用 numpy 1.14.3 或更新版本的 Anaconda3。
推荐使用 TensorFlow 1.14，但 Linux 也支持 TensorFlow 1.15。不支持 TensorFlow 2.x。
在 Windows 上，您需要使用 TensorFlow 1.14，因为标准 1.15 安装不包括必要的 C++ 头文件。
1-8 个具有至少 12 GB GPU 内存、NVIDIA 驱动程序、CUDA 10.0 工具包和 cuDNN 7.5 的高端 NVIDIA GPU
64-bit Python 3.6 or 3.7. Anaconda3 with numpy 1.14.3 or newer is recommended.
TensorFlow 1.14 is recommended, but TensorFlow 1.15 is also supported on Linux. TensorFlow 2.x is not supported.
On Windows you need to use TensorFlow 1.14, as the standard 1.15 installation does not include necessary C++ headers.
1–8 high-end NVIDIA GPUs with at least 12 GB of GPU memory, NVIDIA drivers, CUDA 10.0 toolkit and cuDNN 7.5

Steps

Clone the StyleGAN2-ADA repository and go inside the directory
克隆 StyleGAN2-ADA 存储库并进入目录

git clone https://github.com/NVlabs/stylegan2-ada.git
cd styelgan2-ada

Download or create your own dataset. I will be using BIKED dataset that I already preprocessed. You can download my preprocessed version from dropbox.
下载或创建您自己的数据集。我将使用我已经预处理过的BIKED数据集。您可以从dropbox下载我的预处理版本。

Sample Image of BIKED Dataset [CreativeGAN]

# Dowload dataset
wget "https://www.dropbox.com/s/0ybsudabqscstf7/biked_dataset.tar.gz" -q -O biked_dataset.tar.gz# extract dataset
tar -zxvf biked_dataset.tar.gz# Delete the tar.gz file
rm biked_dataset.tar.gz

After extracting the content, you will have a folder named BIKED that contains 4510 square images of bike designs.
提取内容后，您将拥有一个名为 BIKED 的文件夹，其中包含 4510 个自行车设计的方形图像。

Note*: If you are using your own dataset, create a folder and put all training images inside the folder. Make sure all the images are square and the same size.*

注意*：如果您使用自己的数据集，请创建一个文件夹并将所有训练图像放入该文件夹中。确保所有图像都是正方形且大小相同。*

Preparing Dataset
准备数据集

As the code needs the dataset to be in .tfrecords format. We first need to convert our dataset to this format. StyleGAN2-ADA has made a script that makes this conversion easy.
由于代码需要数据集为 .tfrecords 格式。我们首先需要将我们的数据集转换为这种格式。StyleGAN2-ADA 制作了一个脚本，使这种转换变得容易。

# first argument is output and second arg is path to dataset
python dataset_tool.py create_from_images ./datasets/biked biked

This will create a multi-resolution .tfrecord file in /datasets/biked/ folder.
这将在 /datasets/biked/ 文件夹中创建一个多分辨率 .tfrecord 文件。

Training StyleGAN2-ADA

# snap is how often you want to save the model and sample results
# res is what image resolution you want to train on
# augpipe is augmentation pipes, such as 'blit', 'geom', 'color', 'filter', 'noise', 'cutout' or combination of thesepython train.py --outdir ./results --snap=10 --data=./datasets/biked --augpipe=bgcfnc --res=512

您可以修改许多其他参数，请随时查看 train.py 代码以了解有关参数的更多信息。

运行该命令后，它将开始训练并根据您提供的 snap 参数（在本例中为每 10kimg）定期保存结果和模型文件 (.pkl)。一旦您认为结果足够好或 FID 开始趋于稳定，您可以停止训练并使用上次保存的 .pkl 文件。

拥有模型文件后，您可以使用此命令生成图像。

There are many other arguments that you can modify, feel free to check the train.py code to learn more about the arguments.

Once you run the command, it will start training and periodically save the result and the model file (.pkl) based on the snap arguments that you provided (In this case, every 10kimg). Once you think that the result is good enough or the FID starts to plateau, you can stop training and use the last saved .pkl file.

Once you have the model file you can generate images using this command.

python generate.py --outdir=out --trunc=0.5 --seeds=600-605 --network={path_to_pkl_model_file}

您可以为种子提供范围或逗号分隔值。trunc 是截断技巧的值。截断值越高，输出越多样化或越极端，但可能会降低图像质量。该值越低，图像质量越高，但多样性可能越少。最大值为 1。

但是，如果您想生成插值视频或图像网格。可以参考我之前的文章。

You can provide a range or a comma-separated value for the seeds. The trunc is the value for the truncation trick. The higher the truncation value the more diverse or extreme the output, but might lower the image quality. The lower the value, the higher the image quality but might be less diverse. The maximum value is 1.

However, if you want to generate interpolation videos or a grid of images. You can refer to my previous article.

Generating Anime Characters with StyleGAN2

Learn how to generate this cool anime faces interpolation

towardsdatascience.com

Transfer Learning or Resume Training
迁移学习或恢复训练

如果您的训练因某种原因停止或崩溃。您仍然可以从上次保存的进度继续训练。您只需要添加 - resume 参数和模型 (.pkl) 文件的路径。

此外，您还可以将此参数用于迁移学习。通常最好从预训练模型之一开始，而不是从头开始训练，即使数据集本身并不相似。只需将 .pkl 路径替换为StyleGAN-ADA提供的预训练模型之一。

在这个例子中，我将在自行车数据集上从我的预训练模型恢复训练。

If your training stopped or crashed for some reason. You can still resume training from the last saved progress. You just need to add — resume argument with the path to the model (.pkl) file.

Additionally, you can also use this argument for transfer learning. Instead of training from scratch, it is usually best to start with one of the pre-trained models, even if the dataset itself not similar. Just replace the .pkl path with one of the pre-trained models provided by StyleGAN-ADA.

In this example, I will resume training from my pre-trained model on the biked dataset.

python train.py --outdir ./results --snap=10 --data=./datasets/biked --augpipe=bgcfnc --res=512 --resume=full-bike-network-snapshot-004096.pkl

Training Result

Here is an animation of the training result of 256x256 resolution on the Tesla P100 GPU after training for a day.
这是训练一天后在Tesla P100 GPU上256x256分辨率的训练结果的动画。