Jittor 武汉大学队 手写数字生成

简介

本项目是 华中科技大学“武汉大学”队 参加 第二届计图人工智能挑战赛 热身赛 所使用的代码。项目使用 Jittor 框架实现的 Conditional GAN 模型在 MNIST 数据集上训练，并生成一串手写数字的图像。

安装

运行环境

• Jittor 1.3.4.12
• Numpy 1.22.4
• Pillow 9.1.1

安装依赖

执行以下命令安装 python 依赖

pip install -r requirements.txt

预训练模型

预训练模型下载地址。请将预训练模型放入 saved/ 文件夹中供代码加载，并不要改变其文件名。

运行

训练

python CGAN.py train

运行上述指令即可开始训练，训练使用 Jittor 框架内置的 MNIST 数据集格式，训练过程中每一个 epoch 的模型系数和结果文件将保存在 saved/ 文件夹下。

推理

python CGAN.py eval result.png 87

将预训练模型或自行训练的系数 pkl 文件放入 saved/ 文件夹下，并使用上述指令进行推理。推理命令的格式为 python CGAN.py [output file] [epoch]，预训练模型 epoch 数为 87，故给出的上述指令使用的是预训练模型推理。

致谢

本项目参考了以下资料与项目，感谢作者的分享！

• https://cg.cs.tsinghua.edu.cn/jittor/tutorial/2020-5-13-22-47-cgan/
• https://github.com/Jittor/gan-jittor/blob/master/models/cgan/cgan.py
• https://github.com/arturml/mnist-cgan/blob/master/mnist-cgan.ipynb