如何矢量化我的自定义 pytorch conv2d 操作？

Question

我有一个自定义 conv2d 层的向后传递，我认为可以对其进行矢量化。

for s in range(num_samples):
    for c in range(num_channels):
        temp_input = input[s, c, :, :].unsqueeze(0).unsqueeze(0)
        temp_doutput = doutput[s, :, :, :].unsqueeze(1)
        temp_conv2d = torch.nn.functional.conv2d(temp_input, temp_doutput, stride=dilation, padding=padding, dilation=stride, groups=groups).squeeze(0).squeeze(0)
        cut_conv2d = temp_conv2d[:, :kernel_size, :kernel_size]
        grad_w[:, c, :, :] += cut_conv2d

我认为它可以矢量化的原因是，如果它没有矢量化，输入 50000 个大小为 512x512 的图像，在当前配置下，它必须在一个 epoch 中执行数十亿个 conv2d 函数。我猜这是不对的，但我想不出一种方法来对其进行矢量化。

Answer 1

想通了。不确定是否可以删除最后一个 for 循环，但这对我来说已经加速了 100 倍，足够好了。

for s in range(num_samples):
            temp_input = input[s, :, :, :].unsqueeze(1)
            temp_doutput = doutput[s, :, :, :].unsqueeze(1)
            
            temp_conv2d = torch.nn.functional.conv2d(temp_input, temp_doutput, stride=dilation, padding=padding, dilation=stride, groups=groups).squeeze()
            cut_conv2d = torch.permute((temp_conv2d[:, :, :kernel_size, :kernel_size]), (1, 0, 2, 3))
            grad_w[:, :, :, :] += cut_conv2d