时间戳

首页 > 解决方案 > 使用分组卷积进行重复卷积的反向传播

python - 使用分组卷积进行重复卷积的反向传播

问题描述

我有一个 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：我需要优化速度

标签： pythonpytorchcomputer-visionconv-neural-networkbackpropagation

解决方案

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

对于使用组时给定的卷积层，您可以将其视为通过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你就不必担心每次都用内核梯度的总和来更新梯度。通过保持对self.weight和self.biaswith的引用，自动进行累积Tensor.expand。

推荐阅读