python - torch 猫内存爆炸

Question

我尝试在 ResNet50 上用这个模块替换 Conv2d。

class SubtractedConv(nn.Module):
    def __init__(self, input_ch, output_ch, kernels, stride=1):
        super().__init__()
        self.point_wise = nn.Conv2d(input_ch, output_ch//2, 1, bias=False, stride=stride)
        self.depth_wise = nn.Conv2d(output_ch // 2, output_ch // 2, kernels, groups=output_ch // 2, bias=False, padding=kernels // 2)
        self.low_pass = nn.Conv2d(output_ch // 2, output_ch // 2, kernels, bias=False, padding=kernels // 2)
    def forward(self, x):
        p = self.point_wise(x)
        d = self.depth_wise(p)
        d -= p 
        l = self.low_pass(p)
        x = torch.cat((d, l), 1)
        return x

预期的输出应该与正常的 Conv2d 具有相同的 channel ，但我在 torch.cat() 处发现 cuda 内存不足。 我想知道为什么？又该如何处理呢？

我使用Pytorch 0.4.0，在Tesla P100上运行，图像大小224*224，批量大小16。

事实上，类似的东西适用于 Keras，并且参数较少(ResNet50 中为 16M，而普通 Conv2D 为 25M)。

def subtractedconv(input_tensor, kernel_size, filters, stride=1):
    p = kl.Conv2D(filters//2, (1, 1), use_bias=False, strides=stride, padding='same')(input_tensor)
    d = DepthwiseConv2D(kernel_size, use_bias=False, padding='same')(p)
    d = kl.subtract([d, p])
    l = kl.Conv2D(filters//2, kernel_size, use_bias=False, padding='same')(p)
    x = kl.Concatenate(axis=-1)([d, l])
    return x

Answer 1

PyTorch 的问题很可能是创建的中间张量，而不是 torch.cat 本身。为了通过 nn.Conv2d 进行反向传播，您需要将此操作的输入保留在内存中。当您遍历各层时，内存消耗会增加(所有中间结果都会保留)。现在在你的 forward 代码中你有三个

p = self.point_wise(x) # x is kept
d = self.depth_wise(p) # p is kept
d -= p # here we do not need to keep d, because of derivative formula for subtraction
l = self.low_pass(p)
x = torch.cat((d, l), 1) # but since this goes into further processing, we will need to keep d anyway

请注意，即使您的操作可能具有计算效率(例如具有较小的内核)，它们仍然需要相同数量的内存来保存输入特征图 - 换句话说，您为每个 nn.Conv2d 支付大量固定成本，无论其本身的复杂性如何。因此，显然，如果将 1 个 nn.Conv2d 替换为 3 个，则内存消耗预计会增加大约三倍。

不过，对于您的情况，有一个解决方法。由于您的整个操作是线性的(您只执行卷积，这是线性的，减法，这是线性的和串联，这在某种意义上是线性的)，您可以将整个计算归结为单个卷积，并使用精心准备的内核。如果我们将卷积视为线性运算符，并用 point_wise 表示 P 运算，用 depth_wise 表示 D ，用 low_pass 表示 L ，我们得到你的前向计算 concatenate(Dx - Px, LPx) ，它可以简化为[concatenate(D-P, LP)]x。因此，您可以根据三组权重(分别对应 point_wise 、 depth_wise 和 low_pass )预先计算内核，然后调用 nn.functional.conv2d 一次。不过，实现这种预计算很困难，因为它需要对参数张量的形状进行相当复杂的转换，以保留操作的确切语义(例如，从 1x1 内核 D 中减去空间内核 P )。我尝试在 10 分钟内完成此任务，但失败了，但如果这非常重要，请考虑在 PyTorch 论坛上提问或在评论中告诉我。

至于为什么 Keras 会处理它，我不确定，但强烈猜测这要归功于 TensorFlow。与 PyTorch 不同，TensorFlow(主要)使用静态计算图，可以提前分析并优化。我希望 TensorFlow 能够识别三个线性运算符的序列并将它们组合为一个，或者至少部分利用它们的线性来优化内存使用。