基于分组卷积的重复卷积背靠背

Question

我有一个三维张量，每个通道都要和一个核一起旋转。通过快速搜索，最快的方法是使用分组卷积和组数作为信道数。

下面是一个可复制的小例子：

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组成的。最后，我只需要一个更新的kernel。我该怎么做？

PS:我需要优化速度

Answer 1

要回答你自己的答案，平均权重实际上不是一个准确的方法。您可以通过对梯度的求和(见下文)来操作渐变，但不能对权重进行运算。

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

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

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

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

initializer.

• in
• 在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，您不必担心每次用内核梯度之和更新梯度。通过将self.weight和self.bias的引用与Tensor.expand保持在一起，积累将自动发生。

Answer 2

一个可能的答案是，在梯度更新之后采取如下方法：

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

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