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 , 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.
Generating Anime Characters with StyleGAN2
Learn how to generate this cool anime faces interpolation
训练 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 th