Jittor 热身赛 基于CGAN的手写数字生成方法

简介

本项目包含了第四届计图人工智能挑战赛 - 热身赛的代码实现。本项目的特点是： 使用 MNIST 数据集训练Conditional GAN模型，实现了对生成图像内容的控制的目的，并在比赛中取得了0.9999的得分。

安装

本项目可在 1 张 4090 上运行，训练时间约为 0.5 小时。

运行环境

• ubuntu 20.04 LTS
• python >= 3.8
• jittor >= 1.3.0

安装依赖

执行以下命令安装 python 依赖

pip install -r requirements.txt

训练

单卡训练可运行以下命令：

python train.py --n_epochs=1000 --batch_size=32

–n_epochs 指定训练的轮次，–batch_size指定每个批次的数据数目

推理

生成指定数字的图像可以运行以下命令：

python test.py --model_path="./generator_last.pkl" --save_img_path="./result.png"

–model_path指定生成模型权重文件的路径，–save_img_path指定生成图像的保存路径

致谢

此项目基于论文 Conditional Generative Adversarial Nets 实现，部分代码参考了 jittor-gan。