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

![主要结果](images/result1.png)

## 简介

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

## 安装

### 运行环境

- Jittor 1.3.4.12
- Numpy 1.22.4
- Pillow 9.1.1

### 安装依赖
执行以下命令安装 python 依赖

```bash
pip install -r requirements.txt
```

### 预训练模型

预训练模型[下载地址](https://drive.google.com/drive/folders/1qnS1SXwtmR-2H-i8vDxW1W4H5k1E2ENI?usp=sharing)。请将预训练模型放入 `saved/` 文件夹中供代码加载，并不要改变其文件名。

## 运行

### 训练

```bash
python CGAN.py train
```

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

### 推理

```bash
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