使用分组卷积的重复卷积反向传播

Question

我有一个 3D 张量，每个通道都与一个内核进行卷积。通过快速搜索，最快方法是使用分组卷积，其中组数为通道数。

这是一个可重现的小例子：

import torch
import torch.nn as nn
torch.manual_seed(0)


x = torch.rand(1, 3, 3, 3)
first  = x[:, 0:1, ...]
second = x[:, 1:2, ...]
third  = x[:, 2:3, ...]

kernel = nn.Conv2d(1, 1, 3)
conv = nn.Conv2d(3, 3, 3, groups=3)
conv.weight.data = kernel.weight.data.repeat(3, 1, 1, 1)
conv.bias.data = kernel.bias.data.repeat(3)

>>> conv(x)
tensor([[[[-1.0085]],

         [[-1.0068]],

         [[-1.0451]]]], grad_fn=<MkldnnConvolutionBackward>)

>>> kernel(first), kernel(second), kernel(third)
(tensor([[[[-1.0085]]]], grad_fn=<ThnnConv2DBackward>),
 tensor([[[[-1.0068]]]], grad_fn=<ThnnConv2DBackward>),
 tensor([[[[-1.0451]]]], grad_fn=<ThnnConv2DBackward>))

你可以看到完美的作品。

现在来回答我的问题。我需要对此（kernel 对象）进行反向传播。这样做时，conv 的每个权重都会得到自己的更新。但实际上，conv是由kernel重复3次组成的。最后我只需要一个更新的 kernel。我该怎么做？

PS：我需要优化速度

Answer 1

一个可能的答案是像这样在梯度更新后取平均值

kernel.weight.data = conv.weight.data.mean(0).unsqueeze(0)

这是最好的方法吗？或者这在一开始就正确吗？

Answer 2

给自己的答案一个回复，平均权重其实不是一个准确的方法。您可以通过对梯度求和（见下文）来对梯度进行操作，但不能对权重进行求和。

对于使用组的给定卷积层，您可以将其视为通过内核传递多个 groups 元素。因此，梯度是累积的，而不是平均的。结果梯度实际上是梯度的总和：

kernel.weight.grad = conv.weight.grad.sum(0, keepdim=True)

你可以用笔和纸来验证这一点，如果你平均权重，你最终会平均前一步的权重和每个内核的梯度。对于更高级的优化器来说甚至不是这样，它们不会仅仅依赖像 θ_t = θ _t-1 - lr*grad 这样的简单更新方案。因此，您应该直接使用梯度，而不是生成的权重。

解决此问题的另一种方法是实现您自己的共享内核卷积模块。这可以通过以下两个步骤完成：

在 nn.Module 初始化器中定义你的单一内核。
在前向定义中，查看您的内核以匹配数字组。使用 Tensor.expand instead of Tensor.repeat (the latter makes a copy). You should not make copies, they must remain references of the same underlying data i.e. your single kernel. Then, you can apply the grouped convolution with more flexibility using the functional variant of the paper torch.nn.functional.conv2d.

从那里您可以随时反向传播，梯度将在单个基础权重（和偏差）参数上累积。

让我们在实践中看看：

class SharedKernelConv2d(nn.Module):
   def __init__(self, kernel_size, groups, **kwargs):
      super().__init__()
      self.kwargs = kwargs
      self.groups = groups
      self.weight = nn.Parameter(torch.rand(1, 1, kernel_size, kernel_size))
      self.bias = nn.Parameter(torch.rand(1))

   def forward(self, x):
      return F.conv2d(x, 
         weight=self.weight.expand(self.groups, -1, -1, -1), 
         bias=self.bias.expand(self.groups), 
         groups=self.groups, 
         **self.kwargs)

这是一个非常简单但有效的实现。让我们比较一下两者：

>>> sharedconv = SharedKernelConv2d(3, groups=3):

使用另一种方法：

>>> conv = nn.Conv2d(3, 3, 3, groups=3)
>>> conv.weight.data = torch.clone(conv.weight).repeat(3, 1, 1, 1)
>>> conv.bias.data = torch.clone(conv.bias).repeat(3)

在 sharedconv 层上反向传播：

>>> sharedconv(x).mean().backward()

>>> sharedconv.weight.grad
tensor([[[[0.7920, 0.6585, 0.8721],
          [0.6257, 0.3358, 0.6995],
          [0.5230, 0.6542, 0.3852]]]])
>>> sharedconv.bias.grad
tensor([1.])

与重复张量上的梯度求和相比：

>>> conv(x).mean().backward()

>>> conv.weight.grad.sum(0, keepdim=True)
tensor([[[[0.7920, 0.6585, 0.8721],
          [0.6257, 0.3358, 0.6995],
          [0.5230, 0.6542, 0.3852]]]])
>>> conv.bias.grad.sum(0, keepdim=True)
tensor([1.])

使用SharedKernelConv2d，您不必担心每次都使用内核梯度的总和来更新梯度。通过使用 Tensor.expand.

保留对 self.weight 和 self.bias 的引用，自动进行累积。

使用分组卷积的重复卷积反向传播

Backprop for Repeated Convolution using Grouped Convolution

python

pytorch

computer-vision

conv-neural-network

backpropagation