Conditional GAN (CGAN) with MNIST - Jittor实现

项目简介

条件生成对抗网络(CGAN)是 GAN 的一种变体，通过在生成器和判别器中同时引入标签信息，使网络能够生成指定类别的图像。本项目中，我们基于 Jittor 框架，使用 CGAN 来生成 MNIST 数据集中的手写数字图像。

运行环境

Python 3.7+
Jittor
PIL (Pillow)
NumPy

由于 Jittor with CUDA 的安装比较不便，因此建议使用 Docker 安装预制镜像进行运行。Jittor 官方提供 jittor-cuda-11-1 镜像，然而对于现在常见的 RTX40 系列显卡对应的 Architecture 89，该镜像使用的 CUDA 版本过低，无法正常编译运行。因此我们提供了一个自定义 Dockerfile，在 jittor-cuda-11-1 的基础上安装 CUDA 12.2，支持 RTX40 系列显卡的 GPU 加速。

使用方法

# 如果没有创建自定义镜像
# 请自行修改 /path/to/project/ 为您的目录
cd /path/to/project/
docker build -t jittor-cuda-12-2 ./
# 创建自定义镜像之后
docker run -it --gpus all -v /path/to/project:/workspace jittor-cuda-12-2 bash
cd /workspace
python3.8 -m CGAN.py

部分可配置参数如下：

--n_epochs: #训练轮数(默认100)
--batch_size: #批大小(默认64)
--lr: #学习率(默认0.0002)
--latent_dim: #潜在空间维度(默认100)
--img_size: #图像大小(默认32)
--sample_interval: #采样间隔(默认1000)

模型简介

生成器

输入：随机噪声向量(latent_dim=100)和类别标签的嵌入 经过多层全连接网络处理

输出：生成的数字图像(1x32x32)

判别器

输入：图像和类别标签 经过多层全连接网络处理

输出：真假判断分数

项目结构

.
├── CGAN.py # 主代码文件
├── Dockerfile # 自定义 Dockerfile
├── README.md # README 文件
└── result.png # 生成的数字序列图像（示例）

参考资料

Jittor文档