lightweight-gan

Augmentation is critical for low-data settings. By default, translation and cutout are used (color is omitted). You can specify custom augmentation types and probability using --aug-types and --aug-prob.

$ lightweight_gan --data ./path/to/images --aug-prob 0.25 --aug-types [translation,cutout,color]

Run the training process using the lightweight_gan CLI. By default, models are saved to ./models/{name} every 1000 iterations, and samples are saved to ./results/{name}. If --name is not provided, it defaults to default.

$ lightweight_gan --data ./path/to/images --image-size 512

Generate images using a specific checkpoint. If --load-from is omitted, the latest checkpoint is used. Results are saved to a folder with the postfix -generated-{checkpoint_num}.

$ lightweight_gan \
  --name {name of run} \
  --load-from {checkpoint num} \
  --generate \
  --generate-types {types of result, default: [default,ema]} \
  --num-image-tiles {count of image result}

Generate a sequence of images showing the model's progress across checkpoints. Results are saved in a folder with the postfix -progress. You can convert these to a video using ffmpeg.

$ lightweight_gan \
  --name {name of run} \
  --show-progress \
  --generate-types {types of result, default: [default,ema]} \
  --num-image-tiles {count of image result}

# Convert to video:
$ ffmpeg -framerate 10 -pattern_type glob -i '*-ema.jpg' out.mp4

Create smooth transitions between latent vectors using the --generate-interpolation flag.

$ lightweight_gan --name {name of run} --generate-interpolation

Use the --aug-test flag to see how an image will be augmented. This creates a file named {original_name}_augs.jpg containing a grid of augmented tiles.

lightweight_gan \
    --aug-test \
    --data ./path/to/lena.jpg \
    --batch-size 16 \
    --num-image-tiles 4 \
    --aug-types [color,translation]

Install the package via pip to use the Lightweight GAN implementation in PyTorch.

$ pip install lightweight-gan

Customize the training run using various CLI flags for batch size, gradient accumulation, and total steps.

$ lightweight_gan \
    --data ./path/to/images \
    --name {name of run} \
    --batch-size 16 \
    --gradient-accumulate-every 4 \
    --num-train-steps 200000

Several advanced flags are available to tune the architecture and loss:

• --disc-output-size {size}: Changes discriminator output size (default is 1). Use 5 for art datasets.
• --attn-res-layers [res1,res2]: Adds linear + axial attention to specific resolution layers (e.g., [32,64]).
• --dual-contrast-loss: Uses a novel contrastive loss between real and fake logits instead of the default hinge loss.
• --transparent: Enables training with transparent images.
• --greyscale: Enables training with greyscale images.

The following augmentation types can be passed to --aug-types:

• color: Randomly changes brightness, saturation, and contrast.
• cutout: Creates random black boxes on the image.
• offset: Randomly moves the image by x and y-axis with repeating image.
  • offset_h: Only x-axis movement.
  • offset_v: Only y-axis movement.
• translation: Randomly moves the image on the canvas with a black background.

Note: Horizontal flip is applied by default and is not controllable via --aug-types.