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

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 是由重复 3 次的 kernel 组成的。最后，我只需要更新的kernel。我该怎么做？

PS：我需要优化速度

Answer 1

要回答您自己的答案，平均权重实际上并不是一种准确的方法。您可以对梯度求和（见下文），但不能对权重进行操作。

---

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

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

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

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

• 在 nn.Module 初始化程序中定义您的单个内核。
• 在前向定义中，查看内核以匹配数字组。使用Tensor.expand 而不是Tensor.repeat（后者制作副本）。您不应该制作副本，它们必须保留对相同基础数据的引用即您的单个​​内核。然后，您可以使用论文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 的引用，自动进行累积。

Answer 2

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

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

这是最好的方法吗？还是一开始就对了？