E(n)等变可导向CNNs (escnn)

escnn 是一个用于等变深度学习的 PyTorch 扩展。 escnn 是 e2cnn 库的后继者，后者仅支持平面等距变换。 相比之下，escnn 支持对 2D 和 3D 等距变换都具有等变性的可导向 CNN，以及等变的 MLP。

如果您更喜欢使用 Jax，请查看我们库的这个分支 escnn_jax！

---

等变神经网络 保证其特征空间在输入发生变换时表现出特定的变换行为。 例如，经典的卷积神经网络（CNN）在设计上对其输入的平移具有等变性。 这意味着图像的平移会导致网络特征图的相应平移。 本包提供了在图像平面 (\mathbb{R}^2) 的所有 等距变换 (\mathrm{E}(2)) 以及 3D 空间 (\mathbb{R}^3) 的所有 等距变换 (\mathrm{E}(3)) 下——即在 平移、旋转 和 反射 下——具有等变性的神经网络模块实现（并且理论上可以扩展到 (\mathbb{R}^n) 的所有等距变换 (\mathrm{E}(n))）。 与传统 CNN 不同，(\mathrm{E}(n)) 等变模型能够保证在这些变换上的泛化能力，因此数据效率更高。

(\mathrm{E}(n)) 等变可导向 CNN 的特征空间被定义为 特征场 空间，其特点是由它们在旋转和反射下的变换规律所决定。 典型的例子包括标量场（如灰度图像或温度场）或矢量场（如光流或电磁场）。

用户无需指定通道数，而是需要通过指定场的 类型 和它们的 重数 来定义特征空间。 给定指定的输入和输出特征空间，我们的 R2conv 和 R3conv 模块会实例化两者之间最一般的卷积映射。 我们的库还提供了许多其他等变操作来处理特征场，包括非线性层、生成不变特征的映射、批归一化和丢弃层。

理论上，特征场定义在连续空间 (\mathbb{R}^n) 上。 实际上，它们要么被采样在 像素网格 上，要么以 点云 的形式给出。 escnn 使用 GeometricTensor 对象来表示特征场，这些对象将一个 torch.Tensor 包装起来，并附加相应的变换规律。 所有等变操作都会进行动态类型检查，以确保对特征场的几何处理是合理的。

为了参数化对任意紧致群 (G) 具有等变性的可导向核空间， 我们在 论文 中，在 Group Equivariant Convolution Kernels 的 Wigner-Eckart 定理 的基础上，将该定理从 (G)-齐次空间推广到了承载 (G)-作用的一般空间 (X)。 简而言之，我们的方法利用一个 (G)-可导向基底，针对整个欧几里得空间 (\mathbb{R}^n) 上的无约束标量滤波器，生成具有任意输入和输出场 类型 的可导向核空间。 例如，下图左侧展示了两个 (\mathrm{SO}(2))-可导向基底元素，它们用于定义 (\mathbb{R}^2) 上的函数，进而生成右侧两个 (\mathrm{SO}(2))-等变可导向核的基底元素。 具体来说，这些可导向核将一个频率 (l=1) 的矢量场（2 个通道）映射到一个频率 (J=2) 的矢量场（2 个通道）。

(\mathrm{E}(n)) 等变可导向 CNN 将一系列等距变换等变的 CNN 统一并推广到一个单一框架中。 示例包括：

• Group Equivariant Convolutional Networks
• Harmonic Networks: Deep Translation and Rotation Equivariance
• Steerable CNNs
• Rotation equivariant vector field networks
• Learning Steerable Filters for Rotation Equivariant CNNs
• HexaConv
• Roto-Translation Covariant Convolutional Networks for Medical Image Analysis
• 3D Steerable CNNs
• Tensor Field Networks
• Cormorant: Covariant Molecular Neural Networks
• 3D GCNNs for Pulmonary Nodule Detection

更多细节请参阅我们在 ICLR 2022 发表的论文 A Program to Build E(N)-Equivariant Steerable CNNs 以及我们在 NeurIPS 2019 发表的论文 General E(2)-Equivariant Steerable CNNs。

---

该库由四个子包组成，每个子包具有不同的高级功能：

组件	描述
escnn.group	实现了群和表示论的基本概念
escnn.kernels	求解等变卷积核空间
escnn.gspaces	定义欧几里得空间及其对称性
escnn.nn	包含用于构建深度神经网络的等变模块

---

警告： escnn.kernels 在 1.0.0 版本中进行了重大重构，与库的旧版本不兼容。这些更改不会影响库其他部分提供的接口，但有时使用旧版本训练的网络权重可能无法正确加载到新实例化的模型中。 为了向后兼容，我们建议使用 v0.1.9 版本。

演示

由于 $\mathrm{E}(2)$-可导向 CNN 对旋转和反射具有等变性，因此其推理结果与图像方向的选择无关。 下面的可视化演示了这一特性：我们将旋转后的图像输入到一个随机初始化的 $\mathrm{E}(2)$-可导向 CNN 中（左图）。 中间的图表展示了经过几层网络后，由一个标量场（颜色编码）和一个矢量场（箭头）组成的特征空间的等变变换。 在右图中，我们通过将响应场反向旋转，将其转换为随动参考系（稳定视图）。

在随动参考系中特征的不变性从实验上验证了 $\mathrm{E}(2)$-可导向 CNN 的旋转等变性。 请注意，响应的波动是由于图像在像素网格上的采样导致的离散化误差，这使得精确的连续旋转无法实现。

作为对比，我们展示了传统 CNN 在不同图像方向下的特征图响应：

由于传统 CNN 不具备旋转等变性，其响应会随着图像方向的变化而随机变化。这阻碍了 CNN 自动地在不同参考系之间泛化已学习的模式。

实验结果

$\mathrm{E}(n)$-可导向卷积可以作为 CNN 中传统卷积的直接替代品。 在使用相同基础架构（内存和计算成本相近）的情况下， 与 CNN 基线相比，性能得到了显著提升（数值为测试准确率，单位：%）。

模型	旋转后的 ModelNet10
CNN 基线	82.5 ± 1.4
SO(2)-CNN	86.9 ± 1.9
Octa-CNN	89.7 ± 0.6
Ico-CNN	90.0 ± 0.6
SO(3)-CNN	89.5 ± 1.0

所有模型的架构和宽度大致相同。 更多细节请参阅我们的论文。

本库支持我们在先前的 e2cnn 库中实现的 $\mathrm{E}(2)$-可导向 CNN 作为特例； 我们在此列出了一些来自那里的二维场景中的代表性结果：

模型	CIFAR-10	CIFAR-100	STL-10
CNN 基线	2.6   ± 0.1	17.1   ± 0.3	12.74 ± 0.23
E(2)-CNN *	2.39 ± 0.11	15.55 ± 0.13	10.57 ± 0.70
E(2)-CNN	2.05 ± 0.03	14.30 ± 0.09	9.80 ± 0.40

在采用与 CNN 基线相同的训练设置（未进行额外的超参数调优）的情况下，等变模型取得了显著更好的结果（数值为测试错误率，单位：%）。 为了公平比较，未带 * 号的模型设计时尽量保持与基线模型参数数量相近，而带 * 号的模型则保留了通道数，从而保证计算量相当。 更多详情请参阅我们之前的 e2cnn 论文。

入门

escnn 使用起来非常方便，因为它提供了一个高层次的用户界面，将群论和表示论中的大部分复杂细节都抽象掉了。 下面的代码片段展示了如何从一张 RGB 图像执行到 10 个 正则 特征场的等变卷积（对应于 群卷积）。

from escnn import gspaces                                          #  1
from escnn import nn                                               #  2
import torch                                                       #  3
                                                                   #  4
r2_act = gspaces.rot2dOnR2(N=8)                                    #  5
feat_type_in  = nn.FieldType(r2_act,  3*[r2_act.trivial_repr])     #  6
feat_type_out = nn.FieldType(r2_act, 10*[r2_act.regular_repr])     #  7
                                                                   #  8
conv = nn.R2Conv(feat_type_in, feat_type_out, kernel_size=5)       #  9
relu = nn.ReLU(feat_type_out)                                      # 10
                                                                   # 11
x = torch.randn(16, 3, 32, 32)                                     # 12
x = feat_type_in(x)                                                # 13
                                                                   # 14
y = relu(conv(x))                                                  # 15

第 5 行指定了网络应保持等变性的图像平面 $\mathbb{R}^2$ 上的对称群作用。 我们选择了 循环群 $\mathrm{C}_8$，它描述了以 $2\pi/8$ 的倍数进行的离散旋转。 第 6 行指定了输入特征场的类型。 RGB 图像的三个颜色通道被识别为三个独立的标量场，它们按照 $\mathrm{C}_8$ 的 平凡表示 进行变换（当输入图像旋转时，RGB 值不会改变；可参考上方第一张图中的标量场和向量场）。 类似地，第 7 行中指定的输出特征空间由 10 个按照 $\mathrm{C}_8$ 的 正则表示 进行变换的特征场组成。 然后通过将输入和输出类型以及卷积核大小传递给构造函数来实例化 $\mathrm{C}_8$-等变卷积（第 9 行）。 第 10 行实例化了一个作用于输出场的等变 ReLU 非线性层，因此传入了输出场的类型。

第 12 和 13 行生成了一个随机的 RGB 图像小批量，并将其包装成一个 nn.GeometricTensor，以便将其与正确的场类型 feat_type_in 关联起来。 等变模块在第 15 行处理该几何张量。 每个模块都会检查传入的几何张量是否满足预期的变换规律。

由于在测试时不再需要更新参数，训练完成后，任何等变网络都可以转换为纯 PyTorch 模型，且与传统 CNN 相比没有任何额外的计算开销。 目前，代码支持通过 .export() 方法自动转换一些常用的模块；更多详细信息请参阅 文档。

为了帮助您入门，我们提供了一些示例和教程：

• 入门教程 介绍了该库的基本功能。
• 第二个 教程 讲述了如何在旋转后的 MNIST 数据集上构建并训练一个等变模型。
• 请注意，escnn 也支持等变 MLP；请参阅 这些示例。
• 此外，请查看阿姆斯特丹大学 深度学习 2 课程中使用本库的关于可导向 CNN 的 教程。

更复杂的 2D 等变 Wide Resnet 模型已在 e2wrn.py 中实现。 要尝试一个在反射下等变的模型，请运行：

cd examples
python e2wrn.py

而同时在反射和 90 度整数倍旋转下等变的同一模型版本，则可通过以下命令运行：

python e2wrn.py --rot90

您可以在 examples 文件夹中找到更多示例。 例如，se3_3Dcnn.py 实现了一个在 3D 空间中对旋转和平移等变的 3D CNN。您可以尝试运行：

cd examples
python se3_3Dcnn.py

学习等变性和可导向 CNN 的有用资料

如果您想更好地理解等变和可导向神经网络背后的理论，可以查阅以下参考资料：

• Erik Bekkers 在阿姆斯特丹大学 深度学习 2 课程中关于 几何深度学习 的 讲座。
• 该课程的教材还包括一个关于 群卷积 的 教程，以及另一个使用 本库 讲解可导向 CNN 的 教程。
• Gabriele 的 硕士论文 简要概述了理解可导向 CNN 所需的关键数学知识。
• Maurice 的 博士论文 发展了可导向 CNN 的表示理论，推导出其中最主要的几层，并解释了规范场论的观点。
• Gabriele 演讲的 幻灯片 和 录像 对理解可导向卷积及其滤波器的推导过程非常有帮助，并包含了该研究领域的最新进展。
• Lars Veefkind 的 硕士论文 以非常清晰的方式涵盖了群论、表示论和可导向卷积；特别是，其中还包含了一种新颖的 可导向滤波器的构造性和直观推导方法。

依赖项

该库基于 Python 3.7。

torch>=1.3
numpy
scipy
lie_learn
joblib
py3nj

可选依赖：

torch-geometric
pymanopt>=1.0.0
autograd

警告：py3nj 提供了快速计算 Clebsch-Gordan 系数的功能。如果未安装此包，我们的库将依赖数值方法来估计这些系数。这种数值方法不能保证返回与 py3nj 计算得到的相同系数（它们可能仅相差一个符号）。因此，使用和不使用 py3nj 构建的模型可能不兼容。

要成功安装 py3nj，您可能需要在环境中安装 Fortran 编译器。

安装

您可以安装最新的 release：

pip install escnn

或者您可以克隆此仓库并手动安装：

pip install git+https://github.com/QUVA-Lab/escnn

贡献

您想为 escnn 做出贡献吗？太好了！

请查看 CONTRIBUTING.md 中的说明，帮助我们改进这个库！

您有任何疑问吗？或者有一些想法想要讨论吗？欢迎在 Discussions 中开启新话题！

引用

本库的开发是我们论文工作的一部分，包括 A Program to Build E(N)-Equivariant Steerable CNNs 和 General E(2)-Equivariant Steerable CNNs。 如果您使用我们的代码，请引用以下文献：

   @inproceedings{cesa2022a,
        title={A Program to Build {E(N)}-Equivariant Steerable {CNN}s },
        author={Gabriele Cesa and Leon Lang and Maurice Weiler},
        booktitle={International Conference on Learning Representations},
        year={2022},
        url={https://openreview.net/forum?id=WE4qe9xlnQw}
    }
    
   @inproceedings{e2cnn,
       title={{General E(2)-Equivariant Steerable CNNs}},
       author={Weiler, Maurice and Cesa, Gabriele},
       booktitle={Conference on Neural Information Processing Systems (NeurIPS)},
       year={2019},
       url={https://arxiv.org/abs/1911.08251}
   }

欢迎随时 联系我们。

许可证

escnn 采用 BSD Clear 许可证进行分发。详情请参阅 LICENSE 文件。

escnn 快速上手指南

escnn 是一个基于 PyTorch 的扩展库，用于构建等变深度学习模型。它支持 2D 和 3D 空间中的等距变换（平移、旋转、反射），是 e2cnn 库的后继版本。通过使用 escnn，开发者可以构建在几何变换下具有严格数学保证的神经网络，从而显著提高数据效率和泛化能力。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux, macOS 或 Windows
• Python: 3.7 或更高版本
• 核心依赖:
  • PyTorch (推荐最新稳定版)
  • numpy
  • scipy

注意：escnn 强依赖于 PyTorch。请先根据您的 CUDA 版本安装合适的 PyTorch。

安装步骤

1. 安装 PyTorch (如未安装)

访问 PyTorch 官网 获取适合您环境的安装命令。 国内用户推荐使用清华或中科大镜像源加速安装：

pip install torch torchvision torchaudio --index-url https://pypi.tuna.tsinghua.edu.cn/simple

2. 安装 escnn

您可以直接通过 pip 从 PyPI 安装最新版本：

pip install escnn

如果下载速度较慢，可指定国内镜像源：

pip install escnn -i https://pypi.tuna.tsinghua.edu.cn/simple

或者，如果您需要最新的功能（包含尚未发布的修复），可以从 GitHub 源码安装：

pip install git+https://github.com/QUVA-Lab/escnn.git

基本使用

escnn 提供了高级接口，抽象了群论和表示理论的复杂性。核心概念是使用 GeometricTensor 来包装数据，并指定特征场的类型（Type）和重数（Multiplicity）。

以下是最简单的示例，展示如何创建一个支持 2D 旋转和平移等变性的卷积层，并对一个 RGB 图像进行前向传播。

代码示例：构建 E(2) 等变卷积网络

import torch
import escnn
from escnn import gspaces, nn

# 1. 定义几何空间
# 这里我们定义一个 2D 平面空间，具有旋转和平移对称性 (E(2))
r2_space = gspaces.rot2dOnR2(N=8)  # N=8 表示离散化为 8 个旋转角度

# 2. 定义输入和输出特征场类型
# 输入：RGB 图像，视为标量场 (标量在旋转下不变)，通道数为 3
# triv 表示平凡表示 (标量)
input_type = r2_space.trivial_repr * 3 

# 输出：定义一些更复杂的特征场，例如 4 个标量场和 4 个向量场 (频率为 1)
# r2_space.irreps(1) 表示频率为 1 的不可约表示 (类似向量)
output_type = (r2_space.trivial_repr * 4) + (r2_space.irreps(1) * 4)

# 3. 创建等变卷积层
# in_type 和 out_type 必须基于同一个几何空间定义
conv = nn.R2Conv(in_type=input_type, out_type=output_type, kernel_size=3, padding=1)

# 4. 准备输入数据
# 假设输入是一个 batch size 为 2 的 RGB 图像 (2, 3, 32, 32)
input_tensor = torch.randn(2, 3, 32, 32)

# 5. 将普通 Tensor 包装为 GeometricTensor
# 这一步告诉网络输入数据的几何属性（这里是标量场）
input_geotensor = escnn.GeometricTensor(input_tensor, input_type)

# 6. 前向传播
output_geotensor = conv(input_geotensor)

# 7. 获取结果
# 输出仍然是一个 GeometricTensor，包含变换后的特征场
print(f"输入形状：{input_geotensor.tensor.shape}")
print(f"输出形状：{output_geotensor.tensor.shape}")
print(f"输出类型：{output_geotensor.type}")

# 如果需要将其转换回普通 Tensor 用于后续非等变操作或损失计算
output_tensor = output_geotensor.tensor

关键点说明

• gspaces: 用于定义空间对称性（如 rot2dOnR2 用于 2D 旋转，rot3dOnR3 用于 3D 旋转）。
• GeometricTensor: 核心数据结构，它将 torch.Tensor 与其变换规则（Type）绑定，确保运算的几何正确性。
• R2Conv / R3Conv: 自动根据输入和输出的类型，计算出最通用的等变卷积核。
• 类型系统: 不再仅仅指定“通道数”，而是指定“场类型”（如标量、向量、高阶张量场）及其数量。

通过以上步骤，您即可快速构建具有旋转和平移等变性的神经网络模块。更多复杂的非线性激活函数、池化层和批归一化操作可在 escnn.nn 模块中找到。