stylegan2-pytorch

To utilize all available GPUs on a machine, add the --multi-gpus flag. The batch size will be divided evenly across the GPUs. Use the CUDA_VISIBLE_DEVICES environment variable to restrict training to specific GPUs.

$ stylegan2_pytorch --data ./data --multi-gpus --batch-size 32 --gradient-accumulate-every 1

After training, use the --generate flag to create images from the latest checkpoint.

Options:

• --generate-interpolation --interpolation-num-steps <int>: Generate a video of interpolation through two random points in latent space.
• --generate-interpolation --save-frames: Save each individual frame of the interpolation.
• --generate --load-from <checkpoint_number>: Load a specific previous checkpoint (useful if a previous generator was better).
• --generate --trunc-psi <float>: Control truncation (typically 0.5 to 1.0, default is 0.75). Lower values increase sample quality/fidelity but decrease variety.

Run training by pointing to a directory containing your images using the --data flag.

By default, sample images are saved to results/default and model checkpoints are saved to models/default.

$ stylegan2_pytorch --data /path/to/images

When training with limited data (e.g., 1k-2k images), use differentiable augmentation to prevent the discriminator from overfitting.

Configuration:

• --aug-prob <float>: Set the augmentation probability (recommended between 0.0 and 0.7).
• --aug-types [type1,type2,...]: Specify augmentation types. Available types are translation, cutout, and color. Note: Do not include spaces between items in the list.

To monitor training quality using FID scores, first install pytorch-fid, then use the --calculate-fid-every <int> flag. Results are logged to ./results/{name}/fid_scores.txt.

$ pip install pytorch-fid
$ stylegan2_pytorch --data ./data --calculate-fid-every 5000

Install the package via pip. A machine with a GPU and CUDA installed is required.

Windows Installation: If you are on Windows, it is recommended to install PyTorch and TorchVision via Conda first.

$ pip install stylegan2_pytorch

To log losses to the Aim experiment tracker, add the --log flag to your command. This requires Docker to be installed and running aim up in your terminal.

$ stylegan2_pytorch --data ./data --log

Improve results by adding self-attention to designated layers using the --attn-layers flag.

Note: Do not include spaces after commas in the list.

Use the following flags to customize your training run:

• --name <name>: Specify a project name (defaults to default).
• --results_dir <path>: Specify the directory for intermediate results.
• --models_dir <path>: Specify the directory for model checkpoints.
• --network-capacity <int>: Increase network capacity (default is 16) to improve generation at the cost of more memory.
• --image-size <int>: Set the image resolution.
• --batch-size <int>: Set the training batch size.
• --gradient-accumulate-every <int>: Number of steps to accumulate gradients before updating weights.
• --num-train-steps <int>: Total number of training steps.
• --new: Restart training with new settings instead of resuming from the last checkpoint.

You can use the CLI to perform inference without running a full training loop by setting the generate or generate_interpolation flags.

• To generate static samples: Set --generate True and provide a path to a checkpoint via --load_from.
• To generate interpolations: Set --generate_interpolation True and provide a path to a checkpoint via --load_from.

Results will be saved to the directory specified by --results_dir using a timestamped filename.

Use the ModelLoader class to load checkpoints and generate images in Python code.

import torch
from torchvision.utils import save_image
from stylegan2_pytorch import ModelLoader

loader = ModelLoader(
    base_dir = '/path/to/directory',   # path where you invoked the CLI
    name = 'default'                   # project name
)

noise   = torch.randn(1, 512).cuda()
styles  = loader.noise_to_styles(noise, trunc_psi = 0.7)  # pass through mapping network
images  = loader.styles_to_images(styles) # call generator on style vectors

save_image(images, './sample.jpg')

You can apply contrastive learning to the discriminator during standard GAN training to potentially improve stability and final result quality using the --cl-reg flag.

stylegan2_pytorch --data ./data --cl-reg