GAN Compression (MIT HAN Lab, CVPR 2020), supporting Windows 11 and Ubuntu Linux 24.04: compresses conditional GAN generators through a four-stage pipeline — train the mobile teacher, distill it into a smaller student, fine-tune the student into a once-for-all supernet, then export one sub-network as the compressed generator. Two complete pipelines ship with all stage weights: CycleGAN on horse2zebra (compressed generator 2.5G MACs / 0.36M params vs the teacher's 18.3G / 1.99M) and Pix2pix on edges2shoes-r. The engine is vendored in src/gan_compression; every test stage profiles the model and computes FID.

Folder Structure

Path Description
CLAUDE.md / AGENT.md Program specification (identical content).
CHANGELOG.md Change log (including all engine compatibility fixes).
2003.08936.pdf The GAN Compression paper.
data/horse2zebra/, data/edges2shoes-r/ The two datasets, with the FID real-statistics .npz files.
model/<gan>/<dataset>/{mobile,distill,supernet,compressed}/ Shipped stage checkpoints (training also logs here).
model/pt_inception-*.pth Inception weights used by the FID metric.
model/result/ Test/inference outputs: translated images + index.html galleries.
src/gan_compression/ Vendored engine (mit-han-lab/gan-compression).
src/*.py The 8 pipeline wrappers (each takes --model cyclegan|pix2pix).
venv/ Python virtual environment, created by step 01.

Launcher Scripts

Launcher (.bat / .sh) Purpose
01_venv_create Fresh venv + CPU requirements. Aborts loudly if pip fails.
01_venv_create_gpu Same, then swaps torch for the CUDA 12.8 build. Training needs this.
02_venv_activate Activate the venv (Linux: source src/02_venv_activate.sh).
03_train Step 1 — train the mobile teacher. Overwrites the shipped mobile weights.
04_test_mobile Step 2 — test the teacher (profile + FID + gallery).
05_train_distill Step 3 — distillation. Overwrites the shipped distill weights.
06_test_distill Step 4 — test the student.
07_train_supernet Step 5 — once-for-all supernet training. Overwrites the shipped supernet weights.
08_test_supernet Step 6 — test one supernet sub-network.
09_export_compressed Step 7 — export the compressed generator. Overwrites the shipped compressed weights.
10_inference_compressed Step 8 — run the compressed generator on the test set.
11_tensorboard TensorBoard on model/ (port 6066).
99_venv_clear Delete the venv/ folder.

Usage Notes

  • Every wrapper takes --model cyclegan|pix2pix (default cyclegan) and forwards extra arguments to the engine; the defaults reproduce the original notebooks.
  • GPU is auto-detected (CUDA device 0 when available, CPU otherwise). Test/inference/export run fine on CPU; training needs a GPU.
  • Test stages accept --num_test N to limit the image count; open model/result/<gan>/<dataset>/<stage>/index.html for the gallery.
  • Channel configurations: CycleGAN 16_16_32_16_32_32_16_16 (ngf 32), Pix2pix 32_32_48_32_48_48_16_16 (ngf 48).

Quick Start

Windows 11 (GPU)

src\01_venv_create_gpu.bat
src\10_inference_compressed.bat

Ubuntu 24.04

bash src/01_venv_create_gpu.sh
bash src/10_inference_compressed.sh

Pix2pix pipeline, or a quick CPU smoke test:

bash src/10_inference_compressed.sh --model pix2pix
bash src/10_inference_compressed.sh --num_test 4

Full retraining order: 03 → 04 → 05 → 06 → 07 → 08 → 09 → 10 (back up model/ first — the training steps overwrite the shipped weights).

Results

CycleGAN on horse2zebra — real_A is the input horse image, fake_B is the horse translated into a zebra.

Gan compress cyclegan.png

Pix2pix on edges2shoes-r — real_A is the shoe contour image, fake_B is the contour translated into a shoe photo, and real_B is the real shoe photo.

GAN Compression pix2pix.png

References

 

This SDK is built in AppForAI - AI Dev Tools.

 

Purchase license separately: USD 600, permanent authorization, single APP authorization, single machine authorization, one-year activation, one-year download, one-year update, one-year email technical support.