CGAN-Jittor

A Conditional Generative Adversarial Network (CGAN) implemented with Jittor, designed for conditional image generation on the MNIST dataset.

This is the warm-up problem for “第五届计图人工智能挑战赛“.

Requirements

• Python 3.8+
• Jittor
• Cuda, CuDNN (Optional)

Install dependencies:

pip install jittor

Usage

To train and generate images:

python ./CGAN.py

Optional arguments (all have default values and can be set via command line):

• --n_epochs: Number of training epochs (default: 100)
• --batch_size: Batch size (default: 64)
• --lr: Learning rate (default: 0.0002)
• --latent_dim: Dimension of the noise vector (default: 100)
• --n_classes: Number of classes (default: 10)
• --img_size: Image size (default: 32)
• --channels: Number of image channels (default: 1)
• --sample_interval: Interval for saving generated images (default: 1000)

Example:

python ./CGAN.py --n_epochs 200 --batch_size 128

Output

• During training, generated images are saved as 0.png, 1000.png, … etc.
• Model checkpoints are saved every 10 epochs: generator_last.pkl, discriminator_last.pkl
• At the end of the script, a row of images is generated and saved as result.png

Dataset

• Uses the built-in MNIST dataset from Jittor by default.

Example Results

• result.png: A concatenated image of digits generated according to a preset sequence in the script.
• Training images like xxxx.png show samples generated at different training stages.

References

• Jittor Documentation
• Conditional GAN Paper