This is the training code for StyleGAN 2 model.
\n_^_0_^_
\nThese are _^_1_^_ images generated after training for about 80K steps.
\nOur implementation is a minimalistic StyleGAN 2 model training code. Only single GPU training is supported to keep the implementation simple. We managed to shrink it to keep it at less than 500 lines of code, including the training loop.
\nWithout DDP (distributed data parallel) and multi-gpu training it will not be possible to train the model for large resolutions (128+). If you want training code with fp16 and DDP take a look at lucidrains/stylegan2-pytorch.
\nWe trained this on CelebA-HQ dataset. You can find the download instruction in this discussion on fast.ai. Save the images inside _^_2_^_ folder.
\n": "\u8fd9\u662f StyleGan 2 \u6a21\u578b\u7684\u8bad\u7ec3\u4ee3\u7801\u3002
\n_^_0_^_
\n\u8fd9\u4e9b\u662f\u5728\u8bad\u7ec3\u4e86\u5927\u7ea6 80K \u6b65\u4e4b\u540e\u751f\u6210\u7684_^_1_^_\u56fe\u50cf\u3002
\n\u6211\u4eec\u7684\u5b9e\u73b0\u662f\u4e00\u4e2a\u7b80\u7ea6\u7684 StyleGan 2 \u6a21\u578b\u8bad\u7ec3\u4ee3\u7801\u3002\u4ec5\u652f\u6301\u5355\u4e2a GPU \u8bad\u7ec3\uff0c\u4ee5\u4fdd\u6301\u5b9e\u73b0\u7b80\u5355\u3002\u6211\u4eec\u8bbe\u6cd5\u7f29\u5c0f\u4e86\u5b83\uff0c\u4f7f\u5176\u4fdd\u6301\u5728\u4e0d\u5230 500 \u884c\u4ee3\u7801\u4e2d\uff0c\u5305\u62ec\u8bad\u7ec3\u5faa\u73af\u3002
\n\u5982\u679c\u6ca1\u6709 DDP\uff08\u5206\u5e03\u5f0f\u6570\u636e\u5e76\u884c\uff09\u548c\u591a GPU \u8bad\u7ec3\uff0c\u5c06\u65e0\u6cd5\u4e3a\u5927\u5206\u8fa8\u7387\uff08128+\uff09\u8bad\u7ec3\u6a21\u578b\u3002\u5982\u679c\u4f60\u60f3\u7528 fp16 \u548c DDP \u8bad\u7ec3\u4ee3\u7801\uff0c\u53ef\u4ee5\u770b\u770b l ucidrains/stylegan2-pytorch\u3002
\n\u6211\u4eec\u5728 Celeba-HQ \u6570\u636e\u96c6\u4e0a\u8bad\u7ec3\u4e86\u8fd9\u4e2a\u3002\u4f60\u53ef\u4ee5\u5728\u8fd9\u7bc7\u5173\u4e8e fast.ai \u7684\u8ba8\u8bba\u4e2d\u627e\u5230\u4e0b\u8f7d\u8bf4\u660e\u3002\u5c06\u56fe\u50cf\u4fdd\u5b58\u5728_^_2_^_\u6587\u4ef6\u5939\u4e2d\u3002
\n", "This loads the training dataset and resize it to the give image size.
\n": "\u8fd9\u5c06\u52a0\u8f7d\u8bad\u7ec3\u6570\u636e\u96c6\u5e76\u5c06\u5176\u8c03\u6574\u4e3a\u7ed9\u5b9a\u7684\u56fe\u50cf\u5927\u5c0f\u3002
\n", "This generate images using the generator
\n": "\u8fd9\u4f1a\u4f7f\u7528\u751f\u6210\u5668\u751f\u6210\u56fe\u50cf
\n", "This generates noise for each generator block
\n": "\u8fd9\u4f1a\u4e3a\u6bcf\u4e2a\u53d1\u7535\u673a\u7ec4\u4ea7\u751f\u566a\u58f0
\n", "Instead of calculating the regularization losses, the paper proposes lazy regularization where the regularization terms are calculated once in a while. This improves the training efficiency a lot.
\n": "\u6587\u6ca1\u6709\u8ba1\u7b97\u6b63\u5219\u5316\u635f\u5931\uff0c\u800c\u662f\u63d0\u51fa\u4e86\u61d2\u60f0\u7684\u6b63\u5219\u5316\uff0c\u5373\u5076\u5c14\u8ba1\u7b97\u4e00\u6b21\u6b63\u5219\u5316\u9879\u3002\u8fd9\u5927\u5927\u63d0\u9ad8\u4e86\u8bad\u7ec3\u6548\u7387\u3002
\n", "This samples _^_1_^_ randomly and get _^_2_^_ from the mapping network.
\nWe also apply style mixing sometimes where we generate two latent variables _^_3_^_ and _^_4_^_ and get corresponding _^_5_^_ and _^_6_^_. Then we randomly sample a cross-over point and apply _^_7_^_ to the generator blocks before the cross-over point and _^_8_^_ to the blocks after.
\n": "\u8fd9\u662f_^_1_^_\u968f\u673a\u91c7\u6837\u5e76_^_2_^_\u4ece\u6620\u5c04\u7f51\u7edc\u4e2d\u83b7\u53d6\u3002
\n\u6709\u65f6\u6211\u4eec\u8fd8\u4f1a\u5e94\u7528\u6837\u5f0f\u6df7\u5408\uff0c\u6211\u4eec\u751f\u6210\u4e24\u4e2a\u6f5c\u5728\u53d8\u91cf_^_3_^_\u548c_^_4_^_\u5e76\u5f97\u5230\u76f8\u5e94\u7684_^_5_^_\u548c_^_6_^_\u3002\u7136\u540e\u6211\u4eec\u968f\u673a\u91c7\u6837\u4e00\u4e2a\u4ea4\u53c9\u70b9\uff0c\u7136\u540e\u5e94\u7528_^_7_^_\u4e8e\u4ea4\u53c9\u70b9\u4e4b\u524d\u7684\u751f\u6210\u5668\u65b9\u5757\u548c_^_8_^_\u4e4b\u540e\u7684\u533a\u5757\u3002
\n", "\n": "\n", "\n": "\n", "\n": "\n", "
Gradient Penalty Regularization Loss
\n": "\u68af\u5ea6\u60e9\u7f5a\u6b63\u5219\u5316\u635f\u5931
\n", "\n": "\n", "\n": "\u8def\u5f84\u957f\u5ea6\u60e9\u7f5a
\n", "We trained this on CelebA-HQ dataset. You can find the download instruction in this discussion on fast.ai. Save the images inside _^_0_^_ folder.
\n": "\u6211\u4eec\u5728 Celeba-HQ \u6570\u636e\u96c6\u4e0a\u8bad\u7ec3\u4e86\u8fd9\u4e2a\u3002\u4f60\u53ef\u4ee5\u5728\u8fd9\u7bc7\u5173\u4e8e fast.ai \u7684\u8ba8\u8bba\u4e2d\u627e\u5230\u4e0b\u8f7d\u8bf4\u660e\u3002\u5c06\u56fe\u50cf\u4fdd\u5b58\u5728_^_0_^_\u6587\u4ef6\u5939\u4e2d\u3002
\n", "_^_0_^_ and _^_1_^_ for Adam optimizer
\n": "_^_0_^_\u5bf9_^_1_^_\u4e8e Adam \u4f18\u5316\u5668\u6765\u8bf4
\n", "_^_0_^_ of image resolution
\n": "_^_0_^_\u7684\u56fe\u50cf\u5206\u8fa8\u7387
\n", "Accumulate gradients for _^_0_^_
\n": "\u7d2f\u79ef\u68af\u5ea6_^_0_^_
\n", "Add gradient penalty
\n": "\u6dfb\u52a0\u6e10\u53d8\u60e9\u7f5a
\n", "Add model hooks to monitor layer outputs
\n": "\u6dfb\u52a0\u6a21\u578b\u6302\u63a5\u4ee5\u76d1\u89c6\u5c42\u8f93\u51fa
\n", "Add noise tensors to the list
\n": "\u5c06\u566a\u58f0\u5f20\u91cf\u6dfb\u52a0\u5230\u5217\u8868\u4e2d
\n", "Add path length penalty
\n": "\u589e\u52a0\u8def\u5f84\u957f\u5ea6\u60e9\u7f5a
\n", "Batch size
\n": "\u6279\u91cf\u5927\u5c0f
\n", "Calculate and log gradient penalty
\n": "\u8ba1\u7b97\u5e76\u8bb0\u5f55\u68af\u5ea6\u635f\u5931
\n", "Calculate gradients
\n": "\u8ba1\u7b97\u68af\u5ea6
\n", "Calculate path length penalty
\n": "\u8ba1\u7b97\u8def\u5f84\u957f\u5ea6\u635f\u5931
\n", "Clip gradients for stabilization
\n": "\u7528\u4e8e\u7a33\u5b9a\u7684\u526a\u8f91\u6e10\u53d8
\n", "Compute gradients
\n": "\u8ba1\u7b97\u68af\u5ea6
\n", "Continuous cyclic loader
\n": "\u8fde\u7eed\u5faa\u73af\u88c5\u8f7d\u673a
\n", "Convert to PyTorch tensor
\n": "\u8f6c\u6362\u4e3a pyTorch \u5f20\u91cf
\n", "Create an experiment
\n": "\u521b\u5efa\u5b9e\u9a8c
\n", "Create configurations object
\n": "\u521b\u5efa\u914d\u7f6e\u5bf9\u8c61
\n", "Create data loader
\n": "\u521b\u5efa\u6570\u636e\u52a0\u8f7d\u5668
\n", "Create dataset
\n": "\u521b\u5efa\u6570\u636e\u96c6
\n", "Create discriminator and generator
\n": "\u521b\u5efa\u9274\u522b\u5668\u548c\u751f\u6210\u5668
\n", "Create mapping network
\n": "\u521b\u5efa\u6d4b\u7ed8\u7f51\u7edc
\n", "Create optimizers
\n": "\u521b\u5efa\u4f18\u5316\u5668
\n", "Create path length penalty loss
\n": "\u521b\u5efa\u8def\u5f84\u957f\u5ea6\u60e9\u7f5a\u635f\u5931
\n", "Data loader
\n": "\u6570\u636e\u52a0\u8f7d\u5668
\n", "Device to train the model on. _^_0_^_ picks up an available CUDA device or defaults to CPU.
\n": "\u7528\u4e8e\u8bad\u7ec3\u6a21\u578b\u7684\u8bbe\u5907\u3002_^_0_^_\u9009\u62e9\u53ef\u7528\u7684 CUDA \u8bbe\u5907\u6216\u9ed8\u8ba4\u4e3a CPU\u3002
\n", "Dimensionality of _^_0_^_ and _^_1_^_
\n": "_^_0_^_\u548c\u7684\u7ef4\u5ea6_^_1_^_
\n", "Discriminator and generator loss functions. We use Wasserstein loss
\n": "\u9274\u522b\u5668\u548c\u53d1\u751f\u5668\u635f\u8017\u51fd\u6570\u3002\u6211\u4eec\u4f7f\u7528 Wasserstein \u7684\u635f\u5931
\n", "Discriminator and generator losses
\n": "\u9274\u522b\u5668\u548c\u53d1\u7535\u673a\u635f\u8017
\n", "Discriminator classification for generated images
\n": "\u751f\u6210\u56fe\u50cf\u7684\u9274\u522b\u5668\u5206\u7c7b
\n", "Discriminator classification for real images
\n": "\u771f\u5b9e\u56fe\u50cf\u7684\u9274\u522b\u5668\u5206\u7c7b
\n", "Expand _^_0_^_ and _^_1_^_ for the generator blocks and concatenate
\n": "_^_0_^_\u5c55_^_1_^_\u5f00 and for \u751f\u6210\u5668\u5757\u5e76\u8fde\u63a5
\n", "Expand _^_0_^_ for the generator blocks
\n": "_^_0_^_\u4e3a\u53d1\u7535\u673a\u7ec4\u5c55\u5f00
\n", "Flush tracker
\n": "\u51b2\u6d17\u8ffd\u8e2a\u5668
\n", "Generate images
\n": "\u751f\u6210\u56fe\u50cf
\n", "Generate noise for each generator block
\n": "\u4e3a\u6bcf\u4e2a\u53d1\u7535\u673a\u7ec4\u751f\u6210\u566a\u58f0
\n", "Generate noise to add after the first convolution layer
\n": "\u751f\u6210\u8981\u5728\u7b2c\u4e00\u4e2a\u5377\u79ef\u5c42\u4e4b\u540e\u6dfb\u52a0\u7684\u566a\u6ce2
\n", "Generate noise to add after the second convolution layer
\n": "\u751f\u6210\u8981\u5728\u7b2c\u4e8c\u4e2a\u5377\u79ef\u5c42\u4e4b\u540e\u6dfb\u52a0\u7684\u566a\u6ce2
\n", "Generator & Discriminator learning rate
\n": "\u751f\u6210\u5668\u548c\u9274\u522b\u5668\u5b66\u4e60\u901f\u7387
\n", "Get _^_0_^_
\n": "\u5f97\u5230_^_0_^_
\n", "Get _^_0_^_ and _^_1_^_
\n": "\u83b7\u53d6_^_0_^_\u548c_^_1_^_
\n", "Get discriminator loss
\n": "\u83b7\u5f97\u9274\u522b\u5668\u635f\u5931
\n", "Get generator loss
\n": "\u83b7\u5f97\u53d1\u7535\u673a\u635f\u5931
\n", "Get noise
\n": "\u5f97\u5230\u566a\u97f3
\n", "Get number of generator blocks for creating style and noise inputs
\n": "\u83b7\u53d6\u7528\u4e8e\u521b\u5efa\u6837\u5f0f\u548c\u566a\u58f0\u8f93\u5165\u7684\u751f\u6210\u5668\u6a21\u5757\u7684\u6570\u91cf
\n", "Get real images from the data loader
\n": "\u4ece\u6570\u636e\u52a0\u8f7d\u5668\u83b7\u53d6\u771f\u5b9e\u56fe\u50cf
\n", "Get the paths of all _^_0_^_ files
\n": "\u83b7\u53d6\u6240\u6709_^_0_^_\u6587\u4ef6\u7684\u8def\u5f84
\n", "Get the the _^_0_^_-th image
\n": "\u83b7\u53d6\u7b2c_^_0_^_-th \u5f20\u56fe\u7247
\n", "Gradient penalty coefficient _^_0_^_
\n": "\u68af\u5ea6\u60e9\u7f5a\u7cfb\u6570_^_0_^_
\n", "Height/width of the image
\n": "\u56fe\u50cf\u7684\u9ad8\u5ea6/\u5bbd\u5ea6
\n", "How often to log generated images
\n": "\u8bb0\u5f55\u751f\u6210\u7684\u56fe\u50cf\u7684\u9891\u7387
\n", "How often to save model checkpoints
\n": "\u4fdd\u5b58\u6a21\u578b\u68c0\u67e5\u70b9\u7684\u9891\u7387
\n", "Ignore if _^_0_^_
\n": "\u5ffd\u7565\u5982\u679c_^_0_^_
\n", "Initialize
\n": "\u521d\u59cb\u5316
\n", "List to store noise
\n": "\u5b58\u50a8\u566a\u97f3\u7684\u5217\u8868
\n", "Log discriminator loss
\n": "\u65e5\u5fd7\u9274\u522b\u5668\u4e22\u5931
\n", "Log discriminator model parameters occasionally
\n": "\u5076\u5c14\u8bb0\u5f55\u9274\u522b\u5668\u6a21\u578b\u53c2\u6570
\n", "Log generated images
\n": "\u65e5\u5fd7\u751f\u6210\u7684\u56fe\u50cf
\n", "Log generator loss
\n": "\u65e5\u5fd7\u751f\u6210\u5668\u4e22\u5931
\n", "Loop for _^_0_^_
\n": "\u5faa\u73af\u5bfb\u56de_^_0_^_
\n", "Mapping network learning rate (_^_0_^_ lower than the others)
\n": "\u6620\u5c04\u7f51\u7edc\u5b66\u4e60\u7387\uff08_^_0_^_\u4f4e\u4e8e\u5176\u4ed6\uff09
\n", "Mix styles
\n": "\u6df7\u5408\u98ce\u683c
\n", "Multiply by coefficient and add gradient penalty
\n": "\u4e58\u4ee5\u7cfb\u6570\u5e76\u6dfb\u52a0\u68af\u5ea6\u60e9\u7f5a
\n", "Next block has _^_0_^_ resolution
\n": "\u4e0b\u4e00\u4e2a\u533a\u5757\u6709_^_0_^_\u5206\u8fa8\u7387
\n", "Noise resolution starts from _^_0_^_
\n": "\u566a\u58f0\u5206\u8fa8\u7387\u4ece_^_0_^_
\n", "Number of blocks in the generator (calculated based on image resolution)
\n": "\u751f\u6210\u5668\u4e2d\u7684\u5757\u6570\uff08\u6839\u636e\u56fe\u50cf\u5206\u8fa8\u7387\u8ba1\u7b97\uff09
\n", "Number of images
\n": "\u56fe\u50cf\u6570\u91cf
\n", "Number of layers in the mapping network
\n": "\u5236\u56fe\u7f51\u7edc\u4e2d\u7684\u56fe\u5c42\u6570
\n", "Number of steps to accumulate gradients on. Use this to increase the effective batch size.
\n": "\u7d2f\u79ef\u68af\u5ea6\u7684\u6b65\u6570\u3002\u4f7f\u7528\u5b83\u53ef\u4ee5\u589e\u52a0\u6709\u6548\u6279\u6b21\u5927\u5c0f\u3002
\n", "Optimizers
\n": "\u4f18\u5316\u5668
\n", "Path length penalty calculation interval
\n": "\u8def\u5f84\u957f\u5ea6\u60e9\u7f5a\u8ba1\u7b97\u95f4\u9694
\n", "Probability of mixing styles
\n": "\u6df7\u5408\u6837\u5f0f\u7684\u6982\u7387
\n", "Random cross-over point
\n": "\u968f\u673a\u4ea4\u53c9\u70b9
\n", "Reset gradients
\n": "\u91cd\u7f6e\u6e10\u53d8
\n", "Resize the image
\n": "\u8c03\u6574\u56fe\u50cf\u5927\u5c0f
\n", "Return images and _^_0_^_
\n": "\u8fd4\u56de\u56fe\u50cf\u548c_^_0_^_
\n", "Return noise tensors
\n": "\u8fd4\u56de\u566a\u58f0\u5f20\u91cf
\n", "Run the training loop
\n": "\u8fd0\u884c\u8bad\u7ec3\u5faa\u73af
\n", "Sample _^_0_^_ and _^_1_^_
\n": "\u6837\u672c_^_0_^_\u548c_^_1_^_
\n", "Sample images from generator
\n": "\u6765\u81ea\u751f\u6210\u5668\u7684\u6837\u672c\u56fe\u50cf
\n", "Save model checkpoints
\n": "\u4fdd\u5b58\u6a21\u578b\u68c0\u67e5\u70b9
\n", "Set configurations and override some
\n": "\u8bbe\u7f6e\u914d\u7f6e\u5e76\u8986\u76d6\u4e00\u4e9b
\n", "Set models for saving and loading
\n": "\u8bbe\u7f6e\u7528\u4e8e\u4fdd\u5b58\u548c\u52a0\u8f7d\u7684\u6a21\u578b
\n", "Set tracker configurations
\n": "\u8bbe\u7f6e\u8ddf\u8e2a\u5668\u914d\u7f6e
\n", "Skip calculating path length penalty during the initial phase of training
\n": "\u5728\u8bad\u7ec3\u7684\u521d\u59cb\u9636\u6bb5\u8df3\u8fc7\u8ba1\u7b97\u8def\u5f84\u957f\u5ea6\u635f\u5931
\n", "Start the experiment
\n": "\u5f00\u59cb\u5b9e\u9a8c
\n", "Take a training step
\n": "\u8fc8\u51fa\u8bad\u7ec3\u4e00\u6b65
\n", "Take optimizer step
\n": "\u91c7\u53d6\u4f18\u5316\u5668\u6b65\u9aa4
\n", "The first block has only one _^_0_^_ convolution
\n": "\u7b2c\u4e00\u4e2a\u65b9\u5757\u53ea\u6709\u4e00\u4e2a_^_0_^_\u5377\u79ef
\n", "The interval at which to compute gradient penalty
\n": "\u8ba1\u7b97\u68af\u5ea6\u60e9\u7f5a\u7684\u95f4\u9694
\n", "Total number of training steps
\n": "\u8bad\u7ec3\u6b65\u6570\u603b\u6570
\n", "Train the discriminator
\n": "\u8bad\u7ec3\u9274\u522b\u5668
\n", "Train the generator
\n": "\u8bad\u7ec3\u53d1\u7535\u673a
\n", "Training mode state for logging activations
\n": "\u65e5\u5fd7\u8bb0\u5f55\u6fc0\u6d3b\u7684\u8bad\u7ec3\u6a21\u5f0f\u72b6\u6001
\n", "Transformation
\n": "\u8f6c\u578b
\n", "Update _^_0_^_. Set whether to log activation
\n": "\u66f4\u65b0_^_0_^_\u3002\u8bbe\u7f6e\u662f\u5426\u8bb0\u5f55\u6fc0\u6d3b
\n", "We need to calculate gradients w.r.t. real images for gradient penalty
\n": "\u6211\u4eec\u9700\u8981\u7528\u771f\u5b9e\u56fe\u50cf\u8ba1\u7b97\u68af\u5ea6\u4ee5\u83b7\u5f97\u68af\u5ea6\u60e9\u7f5a
\n", "Whether to log model layer outputs
\n": "\u662f\u5426\u8bb0\u5f55\u6a21\u578b\u5c42\u8f93\u51fa
\n", "Without mixing
\n": "\u4e0d\u6df7\u5408
\n", "