深度学习系列31：Dalle生成模型

Posted 2022-06-15 IE06

1. Dalle模型

前面介绍过VAVQE模型，它本质上是一个encoder-decoder模型，只是中间加了一个codebook。这样我们就可以把尺寸大大缩小。
得到codebook后，图片可以用其进行编码，然后使用自回归模型（比如transformer）来进行序列生成。Taming Transformer就是这样的一个模型。与之相对应的，是早起的PixelCNN、PixelRNN等直接在像素级别进行序列预测的模型，只能处理32*32这样的尺寸。
Dalle模型和Taming Transformer基本相同，只是把输入把文字tokens拼接到了图片tokens前面。

2. 模型训练代码

先安装：pip install dalle-pytorch
伪代码如下：
1）训练VAE的codebook

import torch
from dalle_pytorch import DiscreteVAE
vae = DiscreteVAE()
loss = vae(images, return_loss = True)
loss.backward()

这步可以跳过，直接使用OpenAI现成的VAE模型：

from dalle_pytorch import OpenAIDiscreteVAE
vae = OpenAIDiscreteVAE() 

或者用Taming Transformer中预训练的VQGAN VAE：

from dalle_pytorch import VQGanVAE
vae = VQGanVAE()

2）训练dalle模型

import torch
from dalle_pytorch import DALLE
dalle = DALLE(vae = vae)
loss = dalle(text, images, return_loss = True)
loss.backward()

3）生成图片

images = dalle.generate_images(text)
# or images = dalle.generate_images(
    text, img = img_prime,num_init_img_tokens = (14 * 32) )

3. 预测部分代码

网上有训练好的模型：https://github.com/robvanvolt/DALLE-models
然后执行：

python generate.py --dalle_path=模型路径 --taming --text=文本内容 --num_images=1 --batch_size=1 --outputs_dir=输出地址

参考这篇https://github.com/rom1504/dalle-service可以部署网页服务，或者在jupyter中执行：

开发者涨薪指南 48位大咖的思考法则、工作方式、逻辑体系