- 1Python入门基础教程（非常详细）
- 2NLPIR数据挖掘是AI技术在NLP上的应用
- 3Spring Cloud Alibaba AI：快速体验通义千问大模型问答_通义千问 spring ai
- 4Python中的random模块及相关模块详解
- 5固定Linux的ip地址，通过图形化界面操作，简单易上手_linux图形化界面配置ip地址
- 6AI 绘画｜Midjourney设计Logo提示词
- 7大模型（LLMs）算法工程师的面试题_大模型开发工程师 面试问题(1)_面试题 lora微调 vs
- 8Redis宣布商用后，Redis国产化替代方案有那些？_redis 国产替代
- 9【栅格地图路径规划】基于A星算法实现带电机器人取送货路径规划附matlab代码_栅格图 机器人路径规划
- 10云原生网络之微隔离
经典神经网络(11)VQ-VAE模型及其在MNIST数据集上的应用_vqvae
赞
踩
经典神经网络(11)VQ-VAE模型及其在MNIST数据集上的应用
-
我们之前已经了解了PixelCNN模型。
-
今天，我们了解下DeepMind在2017年提出的一种基于离散隐变量（Discrete Latent variables）的生成模型：VQ-VAE。
-
VQ-VAE采用离散隐变量，而不是像VAE那样采用连续的隐变量。其实VQ-VAE本质上是一种AE，只能很好地完成图像压缩，把图像变成一个短得多的向量，而不支持随机图像生成。
-
那么，VQ-VAE会被归类到图像生成模型中呢？这是因为VQ-VAE单独训练了一个基于自回归的模型如PixelCNN来学习先验（prior），对VQ-VAE的离散编码空间采样。而不是像VAE那样采用一个固定的先验（标准正态分布）。 -
此外，VQ-VAE还是一个强大的无监督表征学习模型，它学习的离散编码具有很强的表征能力：
- OpenAI在2021年发布的文本转图像模型DALL-E就是基于VQ-VAE。
- 另外，在BEiT中也用VQ-VAE得到的离散编码作为训练目标。
注：推荐下EleutherAI团队的lucidrains（Phil Wang）的github，他开源复现了ViT、AlphaFold 2、DALLE、 DALLE2、imagen等项目
https://github.com/lucidrains
1 VQ-VAE
1.1 从AE到VQ-VAE
-
AE是一类能够把图片压缩成较短的向量的神经网络模型。
- AE的编码器编码出来的向量空间是不规整的。也就是说，解码器只认识经编码器编出来的向量，而不认识其他的向量。
- 如下图，我们在code空间上，两张图片的编码点中间处取一点，然后将这一点交给解码器，我们希望新的生成图片是一张清晰的图片（类似3/4全月的样子）。但是，实际的结果是，生成图片是模糊且无法辨认的乱码图。
-
只要AE的编码空间比较规整，符合某个简单的数学分布（比如最常见的标准正态分布，如下图所示），那我们就可以从这个分布里随机采样向量，再让解码器根据这个向量来完成随机图片生成了。
- VAE就是这样一种改进版的AE，它用一些巧妙的方法约束了编码向量z，使得z满足标准正态分布。
- 训练完成后，我们就可以扔掉编码器，用来自标准正态分布的随机向量和解码器来实现随机图像生成了。
-
VQ-VAE的作者认为，VAE的生成图片之所以质量不高，是因为图片被编码成了连续向量。而实际上，把图片编码成离散向量会更加自然。
-
至于离散编码的原因，作者解释如下：https://avdnoord.github.io/homepage/slides/SANE2017.pdf
1.2 VQVAE概述
把图像编码成离散向量后，会带来两个问题：
-
第一个问题是，神经网络会默认输入满足一个连续的分布，而不善于处理离散的输入。
- 如果你直接输入0, 1, 2这些数字，神经网络会默认1是一个处于0, 2中间的一种状态。为了解决这一问题，我们可以借鉴NLP中对于离散单词的处理方法。
- 我们可以把嵌入层加到VQ-VAE的解码器前，这个嵌入层就是
embedding space（嵌入空间），也称codebook。 注意：其实Encoder编码出来的是二维离散编码，下图画的是一维。
-
另一个问题是离散向量不好采样。
- VAE之所以把图片编码成符合正态分布的连续向量，就是为了能在图像生成时把编码器扔掉，让随机采样出的向量也能通过解码器变成图片。现在，VQ-VAE把图片编码了一个离散向量，这个离散向量构成的空间是不好采样的。
- VQ-VAE的作者之前设计了一种图像生成网络，叫做PixelCNN。可以用PixelCNN生成离散编码，再利用VQ-VAE的解码器把离散编码变成图像。
-
VQ-VAE的架构图，如下图所示：
- 训练VQ-VAE的编码器和解码器，使得VQ-VAE能把图像变成
latent image(下图zq)，也能把latent image(下图zq)变回图像。 - 训练PixelCNN，让它学习怎么生成
latent image(下图zq)。 - 生成(采样)时，先用PixelCNN采样出
latent image(下图zq)，再用VQ-VAE把latent image(下图zq)翻译成最终的生成图像。
- 训练VQ-VAE的编码器和解码器，使得VQ-VAE能把图像变成
1.3 VQ-VAE设计细节
1.3.1 关联编码器的输出与解码器的输入
如何关联编码器的输出与解码器的输入呢？
假设嵌入空间codebook已经训练完毕，那么对于编码器的每个输出向量 z e ( x ) ze(x) ze(x)，我们需要找出它在嵌入空间里的最近邻 z q ( x ) zq(x) zq(x)，把 z e ( x ) ze(x) ze(x)替换成 z q ( x ) zq(x) zq(x)作为解码器的输入。- 方式是：求最近邻，即先计算向量与嵌入空间K个向量每个向量的距离，再对距离数组取一个
argmin，求出最近的下标(如上图中的shape为[1,7,7])，最后用下标去嵌入空间里取向量，就得到了 z q zq zq(如上图中的shape为[1,32,7,7])。下标构成的多维数组，也正是VQ-VAE的离散编码。
1.3.2 梯度复制
- 我们现在能把编码器和解码器拼接到一起，但怎么让梯度从解码器的输入 z q ( x ) zq(x) zq(x)传到 z e ( x ) ze(x) ze(x)？从 z e ( x ) ze(x) ze(x)到 z q ( x ) zq(x) zq(x)的变换是一个从数组里取值，这个操作无法求导。
- VQ-VAE使用了一种叫做"straight-through estimator"的技术【即前向传播和反向传播的计算可以不对应】来完成梯度复制。VQ-VAE使用了一种叫做sg(stop gradient，停止梯度)的运算：
s
g
(
x
)
=
{
x
,
前向传播
0
,
反向传播
前向传播时，
s
g
里的值不变；反向传播时，
s
g
按值为
0
求导，即此次计算无梯度。
sg(x)=
由于VQ-VAE其实是一个AE，误差函数里应该只有原图像和目标图像的重建误差：
L
r
e
c
o
n
s
t
r
u
c
t
=
∣
∣
x
−
d
e
c
o
d
e
r
(
z
q
(
x
)
)
∣
∣
2
2
L_{reconstruct}=||x-decoder(z_q(x))||_2^2
Lreconstruct=∣∣x−decoder(zq(x))∣∣22
我们现在利用sg运算，设计新的重建误差：
L
r
e
c
o
n
s
t
r
u
c
t
=
∣
∣
x
−
d
e
c
o
d
e
r
(
z
e
(
x
)
+
s
g
[
z
q
(
x
)
−
z
e
(
x
)
]
)
∣
∣
2
2
前向传播时，就是拿解码器的输入
z
q
(
x
)
来算误差：
L
r
e
c
o
n
s
t
r
u
c
t
=
∣
∣
x
−
d
e
c
o
d
e
r
(
z
e
(
x
)
+
z
q
(
x
)
−
z
e
(
x
)
)
∣
∣
2
2
=
∣
∣
x
−
d
e
c
o
d
e
r
(
z
q
(
x
)
)
∣
∣
2
2
反向传播时，等价于把解码器的梯度全部传给
z
e
(
x
)
：
L
r
e
c
o
n
s
t
r
u
c
t
=
∣
∣
x
−
d
e
c
o
d
e
r
(
z
e
(
x
)
+
s
g
[
z
q
(
x
)
−
z
e
(
x
)
]
)
∣
∣
2
2
=
∣
∣
x
−
d
e
c
o
d
e
r
(
z
e
(
x
)
)
∣
∣
2
2
L_{reconstruct}=||x-decoder(z_e(x)+sg[z_q(x)-z_e(x)])||_2^2\\ 前向传播时，就是拿解码器的输入z_q(x)来算误差：\\ L_{reconstruct}=||x-decoder(z_e(x)+z_q(x)-z_e(x))||_2^2\\ =||x-decoder(z_q(x))||_2^2\\ 反向传播时，等价于把解码器的梯度全部传给z_e(x)：\\ L_{reconstruct}=||x-decoder(z_e(x)+sg[z_q(x)-z_e(x)])||_2^2\\ =||x-decoder(z_e(x))||_2^2
Lreconstruct=∣∣x−decoder(ze(x)+sg[zq(x)−ze(x)])∣∣22前向传播时，就是拿解码器的输入zq(x)来算误差：Lreconstruct=∣∣x−decoder(ze(x)+zq(x)−ze(x))∣∣22=∣∣x−decoder(zq(x))∣∣22反向传播时，等价于把解码器的梯度全部传给ze(x)：Lreconstruct=∣∣x−decoder(ze(x)+sg[zq(x)−ze(x)])∣∣22=∣∣x−decoder(ze(x))∣∣22
在PyTorch里，(x).detach()就是sg(x)，它的值在前向传播时取x，反向传播时取0。
# stop gradient
decoder_input = ze + (zq - ze).detach()
# decode
x_hat = decoder(decoder_input)
# l_reconstruct
l_reconstruct = mse_loss(x, x_hat)
- 1
- 2
- 3
- 4
- 5
- 6
1.3.3 优化嵌入空间codebook
嵌入空间的优化目标是什么呢？嵌入空间的每一个向量应该能概括一类编码器输出的向量。因此，嵌入空间的向量应该和其对应编码器输出尽可能接近。
L
e
=
∣
∣
z
e
(
x
)
−
z
q
(
x
)
∣
∣
2
2
z
e
(
x
)
是编码器的输出向量，
z
q
(
x
)
是其在嵌入空间的最近邻向量
L_e=||z_e(x)-z_q(x)||_2^2\\ z_e(x)是编码器的输出向量，z_q(x)是其在嵌入空间的最近邻向量
Le=∣∣ze(x)−zq(x)∣∣22ze(x)是编码器的输出向量，zq(x)是其在嵌入空间的最近邻向量
作者认为，编码器和嵌入向量的学习速度应该不一样快。
于是，他们再次使用了停止梯度的技巧，把上面那个误差函数拆成了两部分。其中，β控制了编码器的相对学习速度。作者发现，算法对β的变化不敏感，β取0.1~2.0都差不多。
L
e
=
∣
∣
s
g
[
z
e
(
x
)
]
−
z
q
(
x
)
∣
∣
2
2
+
β
∣
∣
z
e
(
x
)
−
s
g
[
z
q
(
x
)
]
∣
∣
2
2
L_e=||sg[z_e(x)]-z_q(x)||_2^2+\beta||z_e(x)-sg[z_q(x)]||_2^2\\
Le=∣∣sg[ze(x)]−zq(x)∣∣22+β∣∣ze(x)−sg[zq(x)]∣∣22
# vq loss
l_embedding = mse_loss(ze.detach(), zq)
# commitment loss
l_commitment = mse_loss(ze, zq.detach())
- 1
- 2
- 3
- 4
VQ-VAE总体的损失函数可以写成:
L
t
o
t
a
l
=
L
r
e
c
o
n
s
t
r
u
c
t
+
L
e
=
∣
∣
x
−
d
e
c
o
d
e
r
(
z
e
(
x
)
+
s
g
[
z
q
(
x
)
−
z
e
(
x
)
]
)
∣
∣
2
2
+
α
∣
∣
s
g
[
z
e
(
x
)
]
−
z
q
(
x
)
∣
∣
2
2
+
β
∣
∣
z
e
(
x
)
−
s
g
[
z
q
(
x
)
]
∣
∣
2
2
L_{total}=L_{reconstruct} + L_e \\ =||x-decoder(z_e(x)+sg[z_q(x)-z_e(x)])||_2^2 +\alpha||sg[z_e(x)]-z_q(x)||_2^2\\+\beta||z_e(x)-sg[z_q(x)]||_2^2
Ltotal=Lreconstruct+Le=∣∣x−decoder(ze(x)+sg[zq(x)−ze(x)])∣∣22+α∣∣sg[ze(x)]−zq(x)∣∣22+β∣∣ze(x)−sg[zq(x)]∣∣22
# reconstruct loss
l_reconstruct = mse_loss(x, x_hat)
# vq loss
l_embedding = mse_loss(ze.detach(), zq)
# commitment loss
l_commitment = mse_loss(ze, zq.detach())
# total loss
loss = l_reconstruct + \
l_w_embedding * l_embedding + l_w_commitment * l_commitment
- 1
- 2
- 3
- 4
- 5
- 6
- 7
- 8
- 9
- 10
1.3.4 先验模型PixelCNN
- 训练好VQ-VAE后，还需要训练一个先验模型来完成数据生成，论文中采用PixelCNN模型。
- 这里我们不再是学习生成原始的pixels，而是学习生成离散编码：
- 首先，我们需要用已经训练好的VQ-VAE模型对训练图像推理，得到每张图像对应的离散编码；
- 然后用一个PixelCNN来对离散编码进行建模
- 最后的预测层采用基于softmax的多分类，类别数为embedding空间的大小K。
- 那么，生成图像的过程就比较简单了，首先用训练好的PixelCNN模型来采样一个离散编码样本(上图中shape为[1, 32, 7, 7])，然后送入VQ-VAE的decoder中，得到生成的图像。
- 实际上，PixelCNN不是唯一可用的拟合离散分布的模型。
我们可以把它换成Transformer，甚至是diffusion模型。
2 VQ-VAE模型在MNIST数据集上的应用
这里使用的模型为Gated PixelCNN模型，具体可参考：
经典神经网络(10)PixelCNN模型、Gated PixelCNN模型及其在MNIST数据集上的应用
网络结构图如下所示：
2.1 VQ-VAE模型
VQVAE的编码器和解码器的结构很简单，仅由普通的上/下采样层和残差块组成。
- 编码器先是有两个3x3卷积+2倍下采样卷积的模块，再有两个残差块(ReLU, 3x3卷积, ReLU, 1x1卷积)；
- 解码器则反过来，先有两个残差块，再有两个3x3卷积+2倍上采样反卷积的模块。
# Reference: https://github.com/SingleZombie/DL-Demos/tree/master/dldemos/VQVAE import os import time import cv2 import einops import numpy as np import torch import torch.nn as nn import torch.nn.functional as F import torchvision from torch.utils.data import DataLoader, Dataset from torch.utils.data.distributed import DistributedSampler from torchvision import transforms from GatedPixelCNNDemo import GatedPixelCNN, GatedBlock class ResidualBlock(nn.Module): def __init__(self, dim): super().__init__() self.relu = nn.ReLU() self.conv1 = nn.Conv2d(dim, dim, kernel_size=3, stride=1, padding=1) self.conv2 = nn.Conv2d(dim, dim, kernel_size=1) def forward(self, x): tmp = self.relu(x) tmp = self.conv1(tmp) tmp = self.relu(tmp) tmp = self.conv2(tmp) return x + tmp class VQVAE(nn.Module): def __init__(self, input_dim, dim, n_embedding): super().__init__() # 1、编码器 self.encoder = nn.Sequential(nn.Conv2d(input_dim, dim, kernel_size=4, stride=2, padding=1), nn.ReLU(), nn.Conv2d(dim, dim, kernel_size=4, stride=2, padding=1), nn.ReLU(), nn.Conv2d(dim, dim, kernel_size=3, stride=1, padding=1), ResidualBlock(dim), ResidualBlock(dim) ) self.vq_embedding = nn.Embedding(n_embedding, dim) # 初始化为均匀分布 self.vq_embedding.weight.data.uniform_(-1.0 / n_embedding, 1.0 / n_embedding) # 2、解码器 self.decoder = nn.Sequential( nn.Conv2d(dim, dim, 3, 1, 1), ResidualBlock(dim), ResidualBlock(dim), nn.ConvTranspose2d(dim, dim, 4, 2, 1), nn.ReLU(), nn.ConvTranspose2d(dim, input_dim, 4, 2, 1) ) self.n_downsample = 2 def forward(self, x): # encode [N, 1, 28, 28] -> [N, 32, 7, 7] ze = self.encoder(x) # ze: [N, C, H, W] # embedding [K, C] [32, 32] embedding = self.vq_embedding.weight.data N, C, H, W = ze.shape K, _ = embedding.shape # 求解最近邻 embedding_broadcast = embedding.reshape(1, K, C, 1, 1) ze_broadcast = ze.reshape(N, 1, C, H, W) distance = torch.sum((embedding_broadcast - ze_broadcast)**2, 2) nearest_neighbor = torch.argmin(distance, 1) # make C to the second dim zq = self.vq_embedding(nearest_neighbor).permute(0, 3, 1, 2) # stop gradient decoder_input = ze + (zq - ze).detach() # decode x_hat = self.decoder(decoder_input) return x_hat, ze, zq @torch.no_grad() def encode(self, x): ze = self.encoder(x) embedding = self.vq_embedding.weight.data # ze: [N, C, H, W] # embedding [K, C] N, C, H, W = ze.shape K, _ = embedding.shape embedding_broadcast = embedding.reshape(1, K, C, 1, 1) ze_broadcast = ze.reshape(N, 1, C, H, W) distance = torch.sum((embedding_broadcast - ze_broadcast)**2, 2) nearest_neighbor = torch.argmin(distance, 1) return nearest_neighbor @torch.no_grad() def decode(self, discrete_latent): zq = self.vq_embedding(discrete_latent).permute(0, 3, 1, 2) x_hat = self.decoder(zq) return x_hat # Shape: [C, H, W] def get_latent_HW(self, input_shape): C, H, W = input_shape return (H // 2**self.n_downsample, W // 2**self.n_downsample)
- 1
- 2
- 3
- 4
- 5
- 6
- 7
- 8
- 9
- 10
- 11
- 12
- 13
- 14
- 15
- 16
- 17
- 18
- 19
- 20
- 21
- 22
- 23
- 24
- 25
- 26
- 27
- 28
- 29
- 30
- 31
- 32
- 33
- 34
- 35
- 36
- 37
- 38
- 39
- 40
- 41
- 42
- 43
- 44
- 45
- 46
- 47
- 48
- 49
- 50
- 51
- 52
- 53
- 54
- 55
- 56
- 57
- 58
- 59
- 60
- 61
- 62
- 63
- 64
- 65
- 66
- 67
- 68
- 69
- 70
- 71
- 72
- 73
- 74
- 75
- 76
- 77
- 78
- 79
- 80
- 81
- 82
- 83
- 84
- 85
- 86
- 87
- 88
- 89
- 90
- 91
- 92
- 93
- 94
- 95
- 96
- 97
- 98
- 99
- 100
- 101
- 102
- 103
- 104
- 105
- 106
- 107
- 108
- 109
- 110
2.2 先验模型
我们已经有了一个普通的PixelCNN模型GatedPixelCNN
- 需要在整个模型的最前面套一个嵌入层，嵌入层的嵌入个数等于离散编码的个数(
color_level)，嵌入长度等于模型的特征长度(p)。 - 由于嵌入层会直接输出一个长度为
p的向量，我们还需要把第一个模块的输入通道数改成p。
# 继承自我们之前实现的模型GatedPixelCNN
class PixelCNNWithEmbedding(GatedPixelCNN):
def __init__(self, n_blocks, p, linear_dim, bn=True, color_level=256):
super().__init__(n_blocks, p, linear_dim, bn, color_level)
self.embedding = nn.Embedding(color_level, p)
self.block1 = GatedBlock('A', p, p, bn)
def forward(self, x):
x = self.embedding(x)
x = x.permute(0, 3, 1, 2).contiguous()
return super().forward(x)
- 1
- 2
- 3
- 4
- 5
- 6
- 7
- 8
- 9
- 10
- 11
- 12
2.3 两种模型的训练
- 下面就是常规的训练代码
- 先训练VQVAE、再训练PixelCNN
def train_vqvae(model: VQVAE, img_shape=None, device='cuda', ckpt_path='./model.pth', batch_size=64, dataset_type='MNIST', lr=1e-3, n_epochs=100, l_w_embedding=1, l_w_commitment=0.25): print('batch size:', batch_size) dataloader = get_dataloader(dataset_type, batch_size, img_shape=img_shape) model.to(device) model.train() optimizer = torch.optim.Adam(model.parameters(), lr) mse_loss = nn.MSELoss() tic = time.time() for e in range(n_epochs): total_loss = 0 for x in dataloader: current_batch_size = x.shape[0] x = x.to(device) x_hat, ze, zq = model(x) # 1、reconstruct loss l_reconstruct = mse_loss(x, x_hat) # 2、vq loss + commitment loss l_embedding = mse_loss(ze.detach(), zq) l_commitment = mse_loss(ze, zq.detach()) # total loss loss = l_reconstruct + \ l_w_embedding * l_embedding + l_w_commitment * l_commitment optimizer.zero_grad() loss.backward() optimizer.step() total_loss += loss.item() * current_batch_size total_loss /= len(dataloader.dataset) toc = time.time() torch.save(model.state_dict(), ckpt_path) print(f'epoch {e} loss: {total_loss} elapsed {(toc - tic):.2f}s') print('Done') def train_generative_model(vqvae: VQVAE, model, img_shape=None, device='cuda', ckpt_path='./gen_model.pth', dataset_type='MNIST', batch_size=64, n_epochs=50): print('batch size:', batch_size) dataloader = get_dataloader(dataset_type, batch_size, img_shape=img_shape) vqvae.to(device) vqvae.eval() model.to(device) model.train() optimizer = torch.optim.Adam(model.parameters(), 1e-3) # 交叉熵损失 loss_fn = nn.CrossEntropyLoss() tic = time.time() for e in range(n_epochs): total_loss = 0 for x in dataloader: current_batch_size = x.shape[0] with torch.no_grad(): x = x.to(device) # 1、训练好的VQ-VAE模型对训练图像推理，得到每张图像对应的离散编码 x = vqvae.encode(x) # 2、用一个PixelCNN来对离散编码进行建模 predict_x = model(x) # 3、预测层采用基于softmax的多分类 loss = loss_fn(predict_x, x) optimizer.zero_grad() loss.backward() optimizer.step() total_loss += loss.item() * current_batch_size total_loss /= len(dataloader.dataset) toc = time.time() torch.save(model.state_dict(), ckpt_path) print(f'epoch {e} loss: {total_loss} elapsed {(toc - tic):.2f}s') print('Done') def reconstruct(model, x, device, dataset_type='MNIST'): model.to(device) model.eval() with torch.no_grad(): x_hat, _, _ = model(x) n = x.shape[0] n1 = int(n**0.5) x_cat = torch.concat((x, x_hat), 3) x_cat = einops.rearrange(x_cat, '(n1 n2) c h w -> (n1 h) (n2 w) c', n1=n1) x_cat = (x_cat.clip(0, 1) * 255).cpu().numpy().astype(np.uint8) cv2.imwrite(f'work_dirs/vqvae_reconstruct_{dataset_type}.jpg', x_cat)
- 1
- 2
- 3
- 4
- 5
- 6
- 7
- 8
- 9
- 10
- 11
- 12
- 13
- 14
- 15
- 16
- 17
- 18
- 19
- 20
- 21
- 22
- 23
- 24
- 25
- 26
- 27
- 28
- 29
- 30
- 31
- 32
- 33
- 34
- 35
- 36
- 37
- 38
- 39
- 40
- 41
- 42
- 43
- 44
- 45
- 46
- 47
- 48
- 49
- 50
- 51
- 52
- 53
- 54
- 55
- 56
- 57
- 58
- 59
- 60
- 61
- 62
- 63
- 64
- 65
- 66
- 67
- 68
- 69
- 70
- 71
- 72
- 73
- 74
- 75
- 76
- 77
- 78
- 79
- 80
- 81
- 82
- 83
- 84
- 85
- 86
- 87
- 88
- 89
- 90
- 91
- 92
- 93
- 94
- 95
- 96
- 97
- 98
- 99
- 100
class MNISTImageDataset(Dataset): def __init__(self, img_shape=(28, 28)): super().__init__() self.img_shape = img_shape self.mnist = torchvision.datasets.MNIST(root='/root/autodl-fs/data/minist') def __len__(self): return len(self.mnist) def __getitem__(self, index: int): img = self.mnist[index][0] pipeline = transforms.Compose( [transforms.Resize(self.img_shape), transforms.ToTensor()]) return pipeline(img) def get_dataloader(type, batch_size, img_shape=None, dist_train=False, num_workers=0, **kwargs): if type == 'MNIST': if img_shape is not None: dataset = MNISTImageDataset(img_shape) else: dataset = MNISTImageDataset() if dist_train: sampler = DistributedSampler(dataset) dataloader = DataLoader(dataset, batch_size=batch_size, sampler=sampler, num_workers=num_workers) return dataloader, sampler else: dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers) return dataloader cfg = dict(dataset_type='MNIST', img_shape=(1, 28, 28), dim=32, n_embedding=32, batch_size=32, n_epochs=20, l_w_embedding=1, l_w_commitment=0.25, lr=2e-4, n_epochs_2=50, batch_size_2=32, pixelcnn_n_blocks=15, pixelcnn_dim=128, pixelcnn_linear_dim=32, vqvae_path='./model_mnist.pth', gen_model_path='./gen_model_mnist.pth') if __name__ == '__main__': os.makedirs('work_dirs', exist_ok=True) device = 'cuda' if torch.cuda.is_available() else 'cpu' img_shape = cfg['img_shape'] # 初始化模型 vqvae = VQVAE(img_shape[0], cfg['dim'], cfg['n_embedding']) gen_model = PixelCNNWithEmbedding(cfg['pixelcnn_n_blocks'], cfg['pixelcnn_dim'], cfg['pixelcnn_linear_dim'], True, cfg['n_embedding']) # 1. Train VQVAE train_vqvae(vqvae, img_shape=(img_shape[1], img_shape[2]), device=device, ckpt_path=cfg['vqvae_path'], batch_size=cfg['batch_size'], dataset_type=cfg['dataset_type'], lr=cfg['lr'], n_epochs=cfg['n_epochs'], l_w_embedding=cfg['l_w_embedding'], l_w_commitment=cfg['l_w_commitment']) # 2. Test VQVAE by visualizaing reconstruction result vqvae.load_state_dict(torch.load(cfg['vqvae_path'])) dataloader = get_dataloader(cfg['dataset_type'], 16, img_shape=(img_shape[1], img_shape[2])) img = next(iter(dataloader)).to(device) reconstruct(vqvae, img, device, cfg['dataset_type']) # 3. Train Generative model (Gated PixelCNN) vqvae.load_state_dict(torch.load(cfg['vqvae_path'])) train_generative_model(vqvae, gen_model, img_shape=(img_shape[1], img_shape[2]), device=device, ckpt_path=cfg['gen_model_path'], dataset_type=cfg['dataset_type'], batch_size=cfg['batch_size_2'], n_epochs=cfg['n_epochs_2']) # 4. Sample VQVAE vqvae.load_state_dict(torch.load(cfg['vqvae_path'])) gen_model.load_state_dict(torch.load(cfg['gen_model_path'])) sample_imgs(vqvae, gen_model, cfg['img_shape'], device=device, n_sample=1, dataset_type=cfg['dataset_type'])
- 1
- 2
- 3
- 4
- 5
- 6
- 7
- 8
- 9
- 10
- 11
- 12
- 13
- 14
- 15
- 16
- 17
- 18
- 19
- 20
- 21
- 22
- 23
- 24
- 25
- 26
- 27
- 28
- 29
- 30
- 31
- 32
- 33
- 34
- 35
- 36
- 37
- 38
- 39
- 40
- 41
- 42
- 43
- 44
- 45
- 46
- 47
- 48
- 49
- 50
- 51
- 52
- 53
- 54
- 55
- 56
- 57
- 58
- 59
- 60
- 61
- 62
- 63
- 64
- 65
- 66
- 67
- 68
- 69
- 70
- 71
- 72
- 73
- 74
- 75
- 76
- 77
- 78
- 79
- 80
- 81
- 82
- 83
- 84
- 85
- 86
- 87
- 88
- 89
- 90
- 91
- 92
- 93
- 94
- 95
- 96
- 97
- 98
- 99
- 100
- 101
- 102
- 103
- 104
- 105
- 106
- 107
- 108
- 109
- 110
- 111
- 112
- 113
- 114
2.4 图像生成(采样)
def sample_imgs(vqvae: VQVAE, gen_model, img_shape, n_sample=81, device='cuda', dataset_type='MNIST'): vqvae = vqvae.to(device) vqvae.eval() gen_model = gen_model.to(device) gen_model.eval() C, H, W = img_shape H, W = vqvae.get_latent_HW((C, H, W)) input_shape = (n_sample, H, W) # 初始化为0 x = torch.zeros(input_shape).to(device).to(torch.long) with torch.no_grad(): # 逐像素预测 for i in range(H): for j in range(W): output = gen_model(x) prob_dist = F.softmax(output[:, :, i, j], -1) # 从概率分布中采样 pixel = torch.multinomial(prob_dist, 1) x[:, i, j] = pixel[:, 0] # 解码 imgs = vqvae.decode(x) imgs = imgs * 255 imgs = imgs.clip(0, 255) imgs = einops.rearrange(imgs, '(n1 n2) c h w -> (n1 h) (n2 w) c', n1=int(n_sample**0.5)) imgs = imgs.detach().cpu().numpy().astype(np.uint8) cv2.imwrite(f'work_dirs/vqvae_sample_{dataset_type}.jpg', imgs)
- 1
- 2
- 3
- 4
- 5
- 6
- 7
- 8
- 9
- 10
- 11
- 12
- 13
- 14
- 15
- 16
- 17
- 18
- 19
- 20
- 21
- 22
- 23
- 24
- 25
- 26
- 27
- 28
- 29
- 30
- 31
- 32
- 33
- 34
- 35
- 36
