Jittor 实现 Conditional GAN 手写数字生成

在数字图片数据集 MNIST 上训练 Conditional GAN（Conditional generative adversarial nets）模型，通过输入一个随机向量 z 和额外的辅助信息 y (如类别标签)，生成特定数字的图像。

简介

本项目为第三节计图挑战赛热身赛——手写数字生成挑战的代码实现。本项目使用Conditional GAN网络对输入向量（随机噪声）与一个类别标签映射为一个对应类别标签（数字）的二维图片。

安装

本项目可在 1 张 Nvidia RTX 2080Ti 显卡上运行，训练时间约为 3 小时。

使用计图框架官方 Docker 镜像 CUDA 版。见安装指南

运行环境

• ubuntu 18.04 LTS
• python >= 3.7
• jittor >= 1.3.0

安装依赖

官方 Docker 镜像默认环境即可完成该任务

预训练模型

预训练模型参数保存在文件discriminator_last.pkl、generator_last.pkl中，下载后放到项目根目录下。

训练

运行以下命令：

python CGAN.py