自回归模型中的 PixelCNN





0/5 (0投票)
本文节选自书籍《PyTorch 深度学习实战》(PyTorch Deep Learning Hands-On),作者为 Sherin Thomas 和 Sudhanshu Passi。本书包含大量示例和动态 AI 应用,展示了 PyTorch 在机器学习和深度学习方面的简洁高效。
引言
PixelCNN 由 DeepMind 提出,是 DeepMind 推出的三种自回归模型之一。自首次推出以来,PixelCNN 已有数次迭代,以提高速度和效率。在本文中,我们将探讨基本的 PixelCNN。
本文节选自书籍《PyTorch 深度学习实战》(PyTorch Deep Learning Hands-On),作者为 Sherin Thomas 和 Sudhanshu Passi。这本书包含大量示例和动态 AI 应用,展示了 PyTorch 在机器学习和深度学习方面的简洁高效。
PixelCNN 一次生成一个像素,然后用它来生成下一个像素,接着用前两个像素来生成下一个。
在 PixelCNN 中,存在一个概率密度模型,它可以学习所有图像的密度分布并从中生成图像。但在这里,我们试图通过考虑所有先前预测的联合概率来条件化生成的每个像素,使其依赖于所有先前生成的像素。
PixelCNN 使用卷积层作为感受野,这提高了输入的读取速度。考虑一个被物体部分遮挡的图像;假设我们只有图像的一半。那么,我们有一半图像,我们的算法需要生成另一半。PixelCNN 通过卷积层一次性获取图像。然而,PixelCNN 的生成过程无论如何都必须是顺序的。您可能想知道,为什么只有一半图像被送入卷积层;答案是掩码卷积(masked convolution),我们将在稍后解释。
图 1.2 显示了如何在像素集上应用卷积操作来预测中心像素。自回归模型相较于其他模型的主要优势在于,其联合概率学习技术是可行的,并且可以使用梯度下降进行学习。没有近似,也没有变通方法;我们只是尝试在给定所有先前像素值的情况下预测每个像素值,并且训练完全由反向传播支持。然而,自回归模型在可扩展性方面存在挑战,因为生成总是顺序的。PixelCNN 是一个精心设计的模型,用于在生成新像素时,将单个概率的乘积作为所有先前像素的联合概率。在 RNN 模型中,这是默认行为,但 CNN 模型通过使用巧妙设计的掩码来实现这一点。
PixelCNN 在参数中捕获像素之间依赖关系的分布,这与其他方法不同。VAE 通过生成隐藏的潜在向量来学习这种分布,这引入了独立的假设。在 PixelCNN 中,学习到的依赖关系不仅存在于之前的像素之间,还存在于不同的通道之间;在普通彩色图像中,它们是红、绿、蓝(RGB)。
存在一个基本问题:如果 CNN 尝试使用当前像素或未来像素来学习当前像素怎么办?这也由掩码处理,该掩码将自我(self)的粒度进一步扩展到通道级别。例如,当前像素的红色通道不会从当前像素本身学习,而是会从之前的像素中学习。但绿色通道现在可以使用当前像素的红色通道和所有之前的像素。同样,蓝色通道可以从当前像素的绿色和红色通道以及所有之前的像素中学习。
在整个网络中使用了两种类型的掩码,但后面的层不需要这种安全机制,尽管它们仍然需要模拟顺序学习,同时进行并行卷积操作。因此,PixelCNN 论文(https://arxiv.org/pdf/1606.05328.pdf)引入了两种类型的掩码:A 型和 B 型。
PixelCNN 与其他传统 CNN 模型的一个主要架构区别在于没有池化层。由于 PixelCNN 的目标不是将图像的本质捕获为降维形式,并且我们不能通过池化损失上下文,因此作者故意去掉了池化层。
fm = 64
net = nn.Sequential(
MaskedConv2d('A', 1, fm, 7, 1, 3, bias=False),
nn.BatchNorm2d(fm), nn.ReLU(True),
MaskedConv2d('B', fm, fm, 7, 1, 3, bias=False),
nn.BatchNorm2d(fm), nn.ReLU(True),
MaskedConv2d('B', fm, fm, 7, 1, 3, bias=False),
nn.BatchNorm2d(fm), nn.ReLU(True),
MaskedConv2d('B', fm, fm, 7, 1, 3, bias=False),
nn.BatchNorm2d(fm), nn.ReLU(True),
MaskedConv2d('B', fm, fm, 7, 1, 3, bias=False),
nn.BatchNorm2d(fm), nn.ReLU(True),
MaskedConv2d('B', fm, fm, 7, 1, 3, bias=False),
nn.BatchNorm2d(fm), nn.ReLU(True),
MaskedConv2d('B', fm, fm, 7, 1, 3, bias=False),
nn.BatchNorm2d(fm), nn.ReLU(True),
MaskedConv2d('B', fm, fm, 7, 1, 3, bias=False),
nn.BatchNorm2d(fm), nn.ReLU(True),
nn.Conv2d(fm, 256, 1))
前面的代码片段是完整的 PixelCNN 模型,它被封装在一个顺序单元中。它由一系列 `MaskedConv2d` 实例组成,该实例继承自 `torch.nn.Conv2d`,并使用 `torch.nn` 中 `Conv2d` 的所有 `*args` 和 `**kwargs`。每个卷积单元后面都跟着一个批归一化层和一个 ReLU 层,这被认为是卷积层成功的组合。作者没有在最后一层使用线性层,而是决定使用一个标准的二维卷积,这被证明比线性层效果更好。
掩码卷积
PixelCNN 中使用掩码卷积来防止在训练网络时,信息从未来像素和当前像素流向生成任务。这一点至关重要,因为在生成像素时,我们无法访问未来像素或当前像素。然而,有一个例外,正如前面所述。当前绿色通道值的生成可以使用红色通道的预测,而当前蓝色通道的生成可以使用绿色和红色通道的预测。
掩码是通过将所有不需要的像素归零来实现的。会创建一个与卷积核大小相等的掩码张量,该张量包含 1 和 0。对于所有不必要的像素,其值为 0。然后,在进行卷积操作之前,将此掩码张量与权重张量相乘。
由于 PixelCNN 不使用池化层和反卷积层,因此随着数据流的进展,通道大小应保持不变。虽然掩码 A 仅负责阻止网络从当前像素学习值,但掩码 B 保持通道大小为三个(RGB),并通过允许当前像素值依赖于其自身值来为网络提供更大的灵活性。
class MaskedConv2d(nn.Conv2d):
def __init__(self, mask_type, *args, **kwargs):
super().__init__(*args, **kwargs)
assert mask_type in ('A', 'B')
self.register_buffer('mask', self.weight.data.clone())
_, _, kH, kW = self.weight.size()
self.mask.fill_(1)
self.mask[:, :, kH // 2, kW // 2 + (mask_type == 'B'):] =
0
self.mask[:, :, kH // 2 + 1:] = 0
def forward(self, x):
self.weight.data *= self.mask
return super(MaskedConv2d, self).forward(x)
前面的 `MaskedConv2d` 类继承自 `torch.nn.Conv2d`,而不是继承自 `torch.nn.Module`。虽然我们通常继承 `torch.nn.Module` 来创建自定义模型类,但由于我们要增强 `Conv2d` 的掩码操作,因此我们继承自 `torch.nn.Conv2D`,而 `torch.nn.Conv2D` 又继承自 `torch.nn.Module`。类方法 `register_buffer` 是 PyTorch 提供的一个方便的 API,用于将任何张量添加到 `state_dict` 字典对象中,当您尝试将模型保存到磁盘时,该对象会与模型一起被保存到磁盘。
添加一个状态变量并使其能在前向函数中重用的显而易见的方法是将其添加为对象属性。
self.mask = self.weight.data.clone()
但这永远不会成为 `state_dict` 的一部分,也永远不会被保存到磁盘。使用 `register_buffer`,我们可以确保我们创建的新张量将成为 `state_dict` 的一部分。然后,掩码张量使用就地操作 `fill_` 填充 1,然后添加 0 以获得一个类似于图 1.3 的张量,尽管图显示的是一个二维张量,而实际的权重张量是三维的。前向函数只是通过与掩码张量相乘来掩码权重张量。乘法保留了掩码为 1 的索引处的所有值,同时删除了掩码为 0 的索引处的所有值。然后,通过对父 `Conv2d` 层进行常规调用,使用权重张量执行二维卷积。
网络的最后一层是一个 softmax 层,它预测像素值在 256 个可能值中的概率,从而离散化了网络的输出生成,而之前的最先进的自回归模型在最后一层使用了连续值生成。
optimizer = optim.Adam(net.parameters())
for epoch in range(25):
net.train(True)
for input, _ in tr:
target = (input[:,0] * 255).long()
out = net(input)
loss = F.cross_entropy(out, target)
optimizer.zero_grad()
loss.backward()
optimizer.step()
训练使用带有默认动量率的 Adam 优化器。此外,损失函数是从 PyTorch 的 Functional 模块创建的。除了目标变量的创建之外,所有其他内容都与正常训练操作相同。
到目前为止,我们一直在进行监督学习,其中标签是明确给出的,但在这种情况下,目标与输入相同,因为我们试图重现相同的输出。torchvision 包应用了转换和归一化来处理像素,并将像素值范围从 0 到 255 转换为 -1 到 1。我们需要转换回 0 到 255 的范围,因为我们在最后一层使用了 softmax,它会在 0 到 255 的范围内生成概率分布。
在本文中,我们实现了 WaveNet 的基础块 PixelCNN,它建立在自回归卷积神经网络 (CNN) 之上。PixelCNN 使用卷积层作为感受野,这提高了输入的读取速度。
关于作者
Sherin Thomas 的职业生涯始于一名信息安全专家,随后将重心转移到基于深度学习的安全系统。他曾帮助全球多家公司建立人工智能管道,最近还为总部位于班加罗尔的一家快速发展的初创公司 CoWrks 工作。Sherin 正在参与多个开源项目,包括 PyTorch、RedisAI 等,并领导 TuringNetwork.ai 的开发。
Sudhanshu Passi 是 CoWrks 的一名技术专家。除其他职责外,他一直是 CoWrks 所有与机器学习相关工作的推动者。他擅长简化复杂概念,使他的作品成为初学者和专家都能阅读的理想读物。这一点可以从他的许多博客和首次出版的书籍中得到证实。业余时间,他会在当地的游泳池里一边游泳一边在水下计算梯度下降。