构建常规等变 CNN 的原则-Python教程-PHP中文网

其中一个原则简单地表述为“让内核旋转”，我们将在本文中重点介绍如何将其应用到您的架构中。等变架构使我们能够训练对某些群体行为无关的模型。为了理解这到底意味着什么，让我们在 MNIST 数据集（0-9 的手写数字数据集）上训练这个简单的 CNN 模型。

class SimpleCNN(nn.Module):

    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.cl1 = nn.Conv2d(in_channels=1, out_channels=8, kernel_size=3, padding=1)
        self.max_1 = nn.MaxPool2d(kernel_size=2)
        self.cl2 = nn.Conv2d(in_channels=8, out_channels=16, kernel_size=3, padding=1)
        self.max_2 = nn.MaxPool2d(kernel_size=2)
        self.cl3 = nn.Conv2d(in_channels=16, out_channels=16, kernel_size=7)
        self.dense = nn.Linear(in_features=16, out_features=10)

    def forward(self, x: torch.Tensor):
        x = nn.functional.silu(self.cl1(x))
        x = self.max_1(x)
        x = nn.functional.silu(self.cl2(x))
        x = self.max_2(x)
        x = nn.functional.silu(self.cl3(x))
        x = x.view(len(x), -1)
        logits = self.dense(x)
        return logits

登录后复制

Accuracy on test	Accuracy on 90-degree rotated test
97.3%	15.1%

_{表1：SimpleCNN模型的测试精度} 正如预期的那样，我们在测试数据集上获得了超过 95% 的准确率，但是如果我们将图像旋转 90 度呢？如果不采取任何对策，结果会下降到仅比猜测稍好一些。这个模型对于一般应用来说是没有用的。相比之下，让我们训练一个具有相同数量参数的类似等变架构，其中组动作恰好是 90 度旋转。

Accuracy on test	Accuracy on 90-degree rotated test
96.5%	96.5%

_{表2：使用与SimpleCNN模型相同数量的参数来测试EqCNN模型的准确性} 准确性保持不变，我们甚至没有选择数据增强。这些模型在 3D 数据的帮助下变得更加令人印象深刻，但我们将继续使用这个示例来探索核心思想。如果您想亲自测试一下，您可以在 Github-Repo 下免费访问用 PyTorch 和 JAX 编写的所有代码，并且只需两个命令即可使用 Docker 或 Podman 进行训练。玩得开心！

那么什么是等方差？

等变架构保证了某些群体行为下特征的稳定性。组是简单的结构，其中组元素可以组合、反转或不执行任何操作。有兴趣的话可以去维基百科查一下正式的定义。出于我们的目的，您可以想象一组作用于方形图像的 90 度旋转。我们可以将图像旋转 90、180、270 或 360 度。为了反转该动作，我们分别应用 270、180、90 或 0 度旋转。很容易看出，我们可以对表示为的组进行组合、反转或不执行任何操作

C_{4} C_4

。该图像将图像上的所有操作可视化。

Figure 1: Rotated MNIST image by 90°, 180°, 270°, 360°, respectively
_{图 1：将 MNIST 图像分别旋转 90°、180°、270°、360°}

Now, given an input image $x x$ , our CNN model classifier $f_{θ} f_\theta$ , and an arbitrary 90-degree rotation $g g$ , the equivariant property can be expressed as
$f_{θ} (rotate x by g) = f_{θ} (x) f_\theta(\text{rotate } x \text{ by } g) = f_\theta(x)$

Generally speaking, we want our image-based model to have the same outputs when rotated.

As such, equivariant models promise us architectures with baked-in symmetries. In the following section, we will see how our principle can be applied to achieve this property.

How to Make Our CNN Equivariant

The problem is the following: When the image rotates, the features rotate too. But as already hinted, we could also compute the features for each rotation upfront by rotating the kernel.
We could actually rotate the kernel, but it is much easier to rotate the feature map itself, thus avoiding interference with PyTorch's autodifferentiation algorithm altogether.

So, in code, our CNN kernel

x = nn.functional.silu(self.cl1(x))

登录后复制

now acts on all four rotated images:

x_0 = x
x_90 = torch.rot90(x, k=1, dims=(2, 3))
x_180 = torch.rot90(x, k=2, dims=(2, 3))
x_270 = torch.rot90(x, k=3, dims=(2, 3))

x_0 = nn.functional.silu(self.cl1(x_0))
x_90 = nn.functional.silu(self.cl1(x_90))
x_180 = nn.functional.silu(self.cl1(x_180))
x_270 = nn.functional.silu(self.cl1(x_270))

登录后复制

Or more compactly written as a 3D convolution:

self.cl1 = nn.Conv3d(in_channels=1, out_channels=8, kernel_size=(1, 3, 3))
...
x = torch.stack([x_0, x_90, x_180, x_270], dim=-3)
x = nn.functional.silu(self.cl1(x))

登录后复制

The resulting equivariant model has just a few lines more compared to the version above:

class EqCNN(nn.Module):

    def __init__(self):
        super(EqCNN, self).__init__()
        self.cl1 = nn.Conv3d(in_channels=1, out_channels=8, kernel_size=(1, 3, 3))
        self.max_1 = nn.MaxPool3d(kernel_size=(1, 2, 2))
        self.cl2 = nn.Conv3d(in_channels=8, out_channels=16, kernel_size=(1, 3, 3))
        self.max_2 = nn.MaxPool3d(kernel_size=(1, 2, 2))
        self.cl3 = nn.Conv3d(in_channels=16, out_channels=16, kernel_size=(1, 5, 5))
        self.dense = nn.Linear(in_features=16, out_features=10)

    def forward(self, x: torch.Tensor):
        x_0 = x
        x_90 = torch.rot90(x, k=1, dims=(2, 3))
        x_180 = torch.rot90(x, k=2, dims=(2, 3))
        x_270 = torch.rot90(x, k=3, dims=(2, 3))

        x = torch.stack([x_0, x_90, x_180, x_270], dim=-3)
        x = nn.functional.silu(self.cl1(x))
        x = self.max_1(x)

        x = nn.functional.silu(self.cl2(x))
        x = self.max_2(x)

        x = nn.functional.silu(self.cl3(x))

        x = x.squeeze()
        x = torch.max(x, dim=-1).values
        logits = self.dense(x)
        return logits

登录后复制

But why is this equivariant to rotations?
First, observe that we get four copies of each feature map at each stage. At the end of the pipeline, we combine all of them with a max operation.

This is key, the max operation is indifferent to which place the rotated version of the feature ends up in.

To understand what is happening, let us plot the feature maps after the first convolution stage.

_{Figure 2: Feature maps for all four rotations}