Image GPT
基于论文《从像素到生成预训练》(Generative Pretraining from Pixels)(Chen et al.)和伴随的代码的Image GPT的PyTorch实现。
模型生成的测试集半图片补全结果。第一列是输入;最后一列是原始图片
iGPT-S在CIFAR10上预训练。由于模型仅在CIFAR10上训练而不是在整个ImageNet上,因此补全结果相对较差。
工作进行中
- 在GPU上进行批量k-means以对更大数据集进行量化(目前使用的是
sklearn.cluster.MiniBatchKMeans
)。 - BERT风格的预训练(目前仅支持生成式预训练)。
- 从OpenAI加载预训练模型。
- 复现至少iGPT-S的结果。
根据他们的博客文章,最大的模型iGPT-L(1.4M参数)经过2500 V100天的训练。通过大大减少注意力头的数量、层数和输入大小(这会影响模型大小的平方),我们可以在单个NVIDIA 2070上,在不到2小时内在Fashion-MNIST上训练我们自己的模型(26K参数)。
使用
预训练模型
一些预训练模型位于models
目录中。运行./download.sh
以下载cifar10
预训练的iGPT-S模型。
计算质心
下载图像,并使用k-means和num_clusters
聚类来计算质心。这些质心在图像输入模型之前用于量化图像。
# 选项:mnist, fmnist, cifar10
python src/compute_centroids.py --dataset mnist --num_clusters=8
# 创建 data/<dataset>_centroids.npy
注意:使用与模型中的num_vocab
相同的num_clusters
。
训练
模型可以使用src/run.py
和train
子命令进行训练。
生成式预训练
通过指定数据集和模型配置文件,可以进行模型的预训练。configs/s_gen.yml
对应于论文中的iGPT-S,configs/xxs_gen.yml
是一个超小模型用于在计算资源有限的玩具数据集上进行尝试。
python src/run.py --dataset mnist train configs/xxs_gen.yml
分类微调
通过将预训练检查点的路径传递给--pretrained
,以及配置文件和数据集,可以对预训练模型进行微调。
python src/run.py --dataset mnist train configs/xxs_clf.yml --pretrained=models/mnist_gen.ckpt`
采样
可以使用测试集中的随机图像创建如上所示的图像:
# 输出为figure.png
python src/sample.py models/mnist_gen.ckpt
可以像我的推文中看到的那样制作动图:
# 输出为out.gif
python src/gif.py models/mnist_gen.ckpt