这几年做三维视觉和几何深度学习相关项目我反复读过好几遍 Spherical CNN 的参考实现s2cnn 仓库。说实话论文里那些公式推导第一次看得人头晕但代码反而把思路讲得很清楚——球面谐波变换怎么变成矩阵乘法、旋转等变性怎么在张量运算里落地、频域卷积为什么能替代空间域滑窗这些在源码里都有非常直观的答案。这篇文章就是把我读源码时的笔记整理出来顺着代码的执行逻辑从球面卷积的数学动机一路拆到最终的算子实现希望能帮你把这条线串起来。如果你是拿了全景图、球形点云、分子结构这类球面数据想做卷积或者单纯想理解“旋转等变的网络结构到底怎么写”这篇解析会覆盖从原理到排坑的完整过程。1. 先搞懂球面卷积到底在卷积什么1.1 平面卷积为什么拿球面数据没办法先聊一个基础问题普通 CNN 的卷积核为什么在球面上不能直接用。平面卷积的成立依赖“平移等价性”——物体在图像里挪一个位置卷积核跟着挪就能产生相同模式的响应。平面像一个无限大的棋盘每个格子之间只有“上下左右挪几格”这一种关系所以卷积核滑动天然合理。球面不一样。球面上没有一个全局性的“平移”操作。你想象一个罗盘图标从赤道移到北极它的朝向、它和周围像素的相对关系都会发生变化。在球面上把一个局部结构从一个位置带到另一个位置唯一自然的操作是旋转。也就是说球面卷积要对齐的是旋转等变性而不是平移等变性输入球面信号转一个角度网络输出的特征也相应旋转同样的角度而不是“换一个位置出现了同一套响应”。还有个实际问题是参数化畸变。全景图常用等矩形投影展开成一张 2:1 的矩形图看起来很方便但纬度越高同一行像素覆盖的实际球面面积越小。到了两级附近一整行像素几乎挤在一个点上。直接把展开图丢给 2D CNN赤道附近学到的卷积模式到了高纬区域就完全失真了。网络要么浪费大量参数去强行适应这种畸变要么干脆学不到有效特征。这正是 Spherical CNN 要解决的核心痛点。1.2 球面卷积定理把卷积换成频域乘法源码没有尝试在球面上设计“可滑动的卷积核”而是走了一条更巧妙的路线——频域乘法。你回想一下普通信号处理两个信号做卷积等于各自做傅里叶变换后逐点相乘再逆变换回来。球面卷积也有一模一样的定理。球面上的基函数不再是指数函数 e^{ikx}而是一族名为球面谐波Spherical Harmonics的正交基函数 Y_l^m。任意球面信号 f(theta, phi) 都可以展开成这些基函数的线性组合f(theta, phi) sum_l sum_m f_hat(l, m) * Y_l^m(theta, phi)这里的 f_hat(l, m) 就是球面谐波系数等价于“球面上的傅里叶系数”。球面卷积定理说的是球面信号 f 和球面滤波器 h 做卷积先各自做球面谐波变换在频域里把系数相乘最后再做逆变换就能得到卷积结果。更关键的一点是旋转等变性在这里天然成立。旋转一个球面信号等价于在谐波系数上施加一个旋转矩阵后面会讲到具体的 Wigner-D 矩阵。而频域的逐点乘法对每个 (l, m) 分量是独立的所以先旋转再做卷积和先做卷积再旋转在频域上看是完全一致的操作顺序。等变性就这样被嵌入到了计算流程里不需要数据增强来硬学。这也是为什么 s2cnn 这套源码里到处都在处理“系数矩阵”而不是处理“滑动窗口”。网络的大部分计算实际上是在做“球面网格采样值”和“谐波系数”之间的来回转换。1.3 为什么源码里会同时出现 S2 和 SO3读源码时你一定会注意到两个出现频率极高的名词S2 和 SO3。S2 是球面本身也就是我们输入数据所在的空间SO3 是三维旋转群所有可能的旋转操作组成的空间。为什么卷积过程中需要从 S2 跑到 SO3 上去我的理解是这样第一层卷积在球面上扫描各种特征模式但这些特征模式出现时滤波器本身也是带方向的。一个球面信号经过一次“提升”卷积之后输出不能再只是球面上每个位置的一个标量而需要编码“这个特征出现在哪个位置、朝向什么方向”。旋转群 SO3 上的一个点恰好能同时表达“位置”和“朝向”于是第一层的输出被定义成了 SO3 上的函数。这种从 S2 到 SO3 的转变是整个架构的核心设计。后边的每一层都在 SO3 上继续做卷积因为特征已经不再是简单的球面值而是“携带旋转属性”的分布。源码里对应的就是两个模块s2_conv.py 处理从球面到旋转群的第一次卷积so3_conv.py 处理旋转群上的后续卷积。两者在实现上有区别但数学骨架是连贯的——都是在频域做乘法。2. 源码仓库结构一个模块负责一件事2.1 主仓库的基本文件布局s2cnn 这个 repo 的文件不算多但每个文件对应一个独立的数学概念。按代码依赖关系整理如下文件职责我读代码时的关键观察s2cnn/polar.py用 numpy 预计算球面谐波变换需要的所有矩阵核心是生成等矩形网格再逐点求 Y_l^m 的值s2cnn/s2_ft.pyS2FT 层把球面网格信号变换成谐波系数作为可训练的 TensorFlow 层本质就是 tf.matmul 乘一个预计算矩阵代码非常短s2cnn/so3_ft.pySO3FT 层把 SO3 网格上的信号变换成 Wigner-D 域系数和 S2FT 逻辑对称但涉及三重索引 (j, k, l)s2cnn/s2_conv.pyS2 卷积层定义先用 S2FT 变换再在频域乘可训练滤波器最后逆变换s2cnn/so3_conv.pySO3 卷积层定义同理但在 Wigner-D 域做乘法s2cnn/rotate.py对特征做旋转操作的辅助模块用于验证等变性也用于构造旋转不变符号examples/MNIST 球面投影等示例模型适合作为读代码的入口一个直观感受是整套源码的核心不是“卷积”本身而是两套傅里叶变换。卷积在代码里反而显得很“薄”——只不过是在变换域里做一次逐元素乘法。2.2 一套数据从进网络到出分类结果的完整链路我建议读代码时按照一条数据的前向路径来读而不是按文件名字典序读。以仓库里 MNIST 球面化的例子来说完整链路是这样的输入是一张平面 MNIST 手写数字先把它投影到球面上得到一个等矩形采样的球面信号形状大概是 (B, 2B, 3)这里的 2B 是经度方向采样数B 是纬度方向采样数。这个球面信号进入第一个 S2Conv 层。层内部先做 S2FT把网格上的 N 个像素值变成 N 个谐波系数实际按带宽截断后数量会略少。谐波系数和可训练滤波器在频域相乘。做逆变换输出变成 SO3 群上的特征图。后面接若干层 SO3Conv特征维度逐渐增大但分辨率因为频带截断逐渐减小。最后是一个“旋转不变池化”层把所有旋转方向上的响应汇总起来接全连接层出分类结果。代码里这一整套流程并不是写在一个大模型文件里的而是通过把 S2Conv、SO3Conv 当成积木组装。我第一次看的时候最大的收获是每一个层的输入输出张量形状都对应了数学推导中的一个空间。搞清楚每个张量“生活在哪个空间”之后代码就通了。2.3 TF1.x 版本遗留问题和新框架适配有一个实际情况必须提醒原版 s2cnn 是基于 TensorFlow 1.x 写的大量使用 tf.get_variable、tf.name_scope 这些老接口项目里还有对 tf.contrib 的依赖。如果你直接 pip install 然后跑大概率会在环境阶段就卡住。我自己是把核心模块拆出来改成 PyTorch 复现的s2ft 和 so3ft 的预计算矩阵逻辑可以直接复用剩下的网络骨架重写并不算难。如果你没时间自己复现也可以找社区基于这套源码的思路做的 PyTorch 版本。关键要保留的是 polar.py 和 rotate.py 里的数学预计算逻辑这两个文件才是整套源码的灵魂。3. 球面谐波变换的矩阵化实现源码中最精巧的部分3.1 预计算变换矩阵的基本思路S2FT 在源码里实现得极其简洁就是一次矩阵乘法。能做到这一点的原因在于球面谐波 Y_l^m(theta, phi) 是固定函数只要采样网格固定每个网格采样点上的 Y_l^m 值就是固定的。于是我们可以提前把“从网格像素到谐波系数”这个线性变换的矩阵算出来。具体来说设带宽为 B那么 l 取 0 到 B-1m 取 -l 到 l一共有 B^2 个谐波基函数。等矩形网格是一组 (theta_i, phi_j) 的笛卡尔积总共有约 2B^2 个采样点。源码在 polar.py 中做这么几件事# 逻辑示意预计算 S2FT 变换矩阵 def compute_s2ft_matrix(B): # 1. 生成等矩形网格 thetas np.linspace(0, np.pi, B, endpointFalse) phis np.linspace(0, 2 * np.pi, 2 * B, endpointFalse) # 2. 对每个网格点、每个 (l, m) 计算 Y_l^m(theta, phi) # 3. 组装成矩阵 S形状为 (2B * B, B * B) return S等矩形网格看起来简单但它有个好处phi 方向上做离散傅里叶变换时天然周期theta 方向上配合勒让德多项式积分能做到数值精度可控。源码没有直接用 scipy 的现成函数去逐点暴力计算每个网格而是按 l 和 m 的结构做了分块计算把变换矩阵拆成若干小块每一块对应一个固定的 m 频率。这一点在读代码时可以留意它能帮你理解矩阵里的稀疏结构。3.2 谐波系数在张量里的排布方式球面谐波是复值函数所以谐波系数也是复数。但 TensorFlow 的矩阵乘法本身只支持实数张量源码里处理复数的方式不是用 float complex 类型而是把实部和虚部分别存成张量的两个通道。具体做法一个含有 C 个复数通道的特征会被存储成 2C 个实数通道前 C 个是实部后 C 个是虚部。S2Conv 层里对权重矩阵采取同样的拆分方式。这样 tf.matmul 就能直接跑不需要自定义复数乘法算子。代码里那个看起来有点绕的通道维度本质上就是在做复数的逐项乘法和加法。这里有一个容易踩的坑常规 2D 卷积的权重张量是“输出通道 x 输入通道 x 核高 x 核宽”但在 s2cnn 的频域滤波里“核”的概念变成了对每个频率条目 (l, m) 的一组复数权重。也就是说频域卷积的权重更像是一个逐通道的复数线性变换而不是一个空间上滑动的窗口。源码中权重的形状通常和输入频率的数量相关初始化方式也因此必须特殊对待否则实部虚部之间很容易出现方差失衡。3.3 bandwidth 的选择对变换精度的直接影响带宽 B也叫 band-limit是这套源码里最关键的参数。它决定了网络能表达多少频率细节也决定了所有张量的大小。从数学角度看一个带宽为 B 的球面信号可以用 B^2 个谐波系数完整表示。采样网格如果是 (B, 2B)刚好能无损地表达到这个带宽的信号。听起来很完美但问题在于实际数据不可能是严格带限的。图像投影到球面上会有高频细节采样前应该先做低通滤波否则混叠就出来了。源码对 MNIST 这类比较干净的数据影响不大但如果你处理真实全景图混叠会表现为训练时的高频噪声放大特征质量明显下降。调试时我习惯从小的带宽开始比如 B16 或 B32先验证模型能收敛再逐步调大。B64 时张量占用通常已经很可观了。带宽每翻一倍S2FT 变换矩阵的维度大约变成原来的 4 倍两个方向同时翻倍显存消耗也是类似的比例。这个增长规律在做显存规划时一定要心里有数。4. S2Conv 与 SO3Conv频域乘法如何落到代码4.1 S2Conv 的前向计算流程S2Conv 层是整个模型的第一层它的输入是球面网格图输出是 SO3 群上的特征。整个前向过程在源码里可以拆成四步S2FT调用预计算好的变换矩阵把球面网格信号变换为谐波系数。输入形状类似 (batch, B, 2B, in_channels)先在最后的通道维度上把球面像素值排列成向量然后做矩阵乘法得到 (batch, B^2, in_channels) 的频域系数。频域滤波对每个 (l, m) 条目把输入系数和训练权重做复数乘法。这一步是网络真正学习的内容它决定了哪些频率成分的球面模式被保留、被放大或抑制。逆变换把滤波后的系数通过逆 S2FT 变换回实空间——这里实际上是变换回 SO3 空间。非线性激活和 reshape得到本层输出。从代码结构看S2Conv 的构造函数主要做了两件事一是初始化频域权重二是调用 s2ft 层拿到变换矩阵。前向函数就是一次简单的乘法序列。真正隐藏的复杂度全在预计算里。4.2 SO3FT 和 Wigner-D 矩阵在卷积中的角色S2Conv 输出的是 SO3 上的函数要在这个空间继续卷积就需要 SO3 上的傅里叶变换。普通傅里叶变换的基函数是一维指数SO3 的“傅里叶基”是 Wigner-D 矩阵它是三维旋转群的不可约表示。这是一个比较高阶的数学对象但源码的处理方式仍然是同一个套路预计算变换矩阵前向时做矩阵乘法。SO3FT 和 S2FT 最大的差别在于索引结构。SO3 上的函数需要三个索引来标记两个用于编码“位置”的旋转角一个用于编码“朝向”。对应的 Wigner-D 矩阵条目是 D^l_{mn}(alpha, beta, gamma)其中 l 是频率级别m 和 n 都是从 -l 到 l 的整数。在频域做卷积时源码要对 m 和 n 两组索引分别做乘法所以参数量比 S2Conv 高一个量级。我最初读 so3_conv.py 时最大的障碍就是理不清这个张量维度输入是 (batch, B^3, channels)权重是 (B^3, in_channels, out_channels)但 B^3 内部是有结构的不能简单地当成一个平铺向量去理解。建议你把 (l, m, n) 三个索引的循环展开画出来对照源码里的 reshape 和 transpose 一步一对应很快就能理顺。4.3 参数初始化和正则化的代码陷阱频域卷积的权重不像普通卷积那样是空间核它实际上是每个频率分量上的复数线性映射。源码里对初始化的处理是这样对每个 l权重的方差按 1/(2l1)^2 的量级设置其中 (2l1) 是这一频率级别里独立系数的数量。这个设置的逻辑是保持不同频率级别的信号在经过卷积后能量不至于失衡让高频和低频分量在深层网络里有机会被公平地学习。如果直接把 torch.nn.Linear 的默认初始化套到这类权重上训练初期经常会出现 loss 剧烈震荡。原因是高频分量方差过大被激活函数放大或者低频分量方差过小被梯度更新忽略。复现时我建议先打印每一层输出的标准差观察各频率成分的尺度是否在同一个量级不一致就优先检查初始化公式。另外对频域权重做正则化时要小心。源码默认不会对权重做 L2 惩罚因为复数权重的实部和虚部如果被同一力度惩罚等价于对系数的模做惩罚影响的是频率能量而不是稀疏性。如果你在复现时加 weight decay我建议加在信号的模上而不是实部和虚部分别加效果更可控。5. 旋转等变性在代码里是怎么保证的5.1 旋转等变性的根源每个算子都等变整条网络具备旋转等变性的根本原因是组成它的每一个算子都各自具备旋转等变性。这不是“整个网络整体上近似等变”而是在数学上严格成立的性质。反过来说如果你在网络里随意插入一个不具备等变性的算子比如把特征图展开成一维再接全连接那么整个网络的等变性就立刻被打破。源码在最后的分类头前使用了一种“旋转不变池化”操作——对 SO3 特征的所有旋转维度取某种对称的聚合函数比如求最大值或均值把旋转信息压缩掉从而把等变特征转化成分类所再也不需要旋转索引。读源码时你会看到层与层之间很少出现随机 reshape因为每个变换都对应一个数学上等变的算子。这是一条很好的编码准则如果你自己设计新的球面网络层先用数学验证新算子的等变性再动手写前向代码能省掉大量调参时间。5.2 rotate.py 如何用 Wigner-D 矩阵转动特征rotate.py 这个模块不直接参与前向推理但它是理解网络性质和验证等变性的钥匙。它的作用是对已提取的谐波系数施加一个旋转模拟“输入球面信号旋转了一定角度”。代码逻辑大致是这样假设要对 S2 的谐波系数做旋转则在 (l, m) 索引上乘以一个由欧拉角决定的 Wigner-D 矩阵。这个矩阵的作用相当于把信号沿着某个旋转轴转动一个角度。源码实现了这一类旋转操作也实现了反向旋转用于对齐特征。我已经用过一个很实际的场景在训练完成一个分类模型后用 rotate.py 把测试集球面输入旋转若干个角度送入网络查看输出特征的变化。如果网络是严格等变的输出特征应当随着输入旋转而发生相应的置换或旋转但整体语义不变分类层的输出则应该几乎不变。源码作者在仓库里其实也写了一个等变性数值测试你可以直接调用验证自己复现的模型。5.3 等变性验证测试的源码逻辑等变性验证在代码里做起来很直接随机生成一个球面信号记作 f从一组随机旋转 R 中采样一个得到 Rf。然后分别走网络得到输出 Net(f) 和 Net(Rf)。由于 Net 是等变的应当有 Net(Rf) RNet(f)。在谐波域中这对应的就是输出系数之间相差一个 Wigner-D 矩阵变换。代码里直接计算两者之间的误差均值并和一个随机扰动产生的误差对比。这个测试值得每次改动网络结构之后都跑一遍。我遇到过的情况是有一次为了实现加速我在频域权重上做了一种类近似的稀疏化处理结果等变测试的误差突然从 1e-5 跳到了 1e-2。排查后发现是稀疏掩码破坏了 (l, m) 通道之间的对称结构。所以等变测试不只是学术验证也帮我抓住过实际实现里的 bug。5.4 等变性带来的实际收益网络上经常讨论等变性有什么实际价值。从源码能看出的最直接收益是网络不需要通过数据增强来学习“旋转前后是同一个东西”。对全景图或者三维物体识别这类任务来说旋转是天然会出现的干扰因素。普通 CNN 必须翻几十倍数据量去覆盖各种旋转姿态Spherical CNN 则用结构保证了旋转不变性数据效率高出不少。我之前在球面 MNIST 上做过对比实验同样的训练样本量下Spherical CNN 在旋转测试集上的准确率比平面 CNN 高几个点而且即使不做任何旋转增强测试准确率也不会因为输入旋转而波动。这个特性对样本量小、旋转变化大的实际项目非常关键。6. 复现过程中踩过的坑与调试建议6.1 环境依赖和 TF2 兼容原版仓库在 TF1.x 下运行如果你用的是 TF2 环境安装依赖时大概率会报错。主要的坑集中在tf.contrib模块被移除以及若干老 API 改名。我的处理方式是不追求在 TF2 里硬跑原版代码而是把 polar.py 和 rotate.py 这两个核心文件保留下来只用它们预计算变换矩阵网络主体用 PyTorch 重写。如果遇到必须用原版跑通的情况可以考虑安装一个单独的 Python 3.6 TensorFlow 1.15 虚拟环境这样依赖冲突最小。另外注意 numpy 版本不要太高老代码里一些np.float之类的写法在新版 numpy 里会直接报错。6.2 网格密度与显存的取舍球面网格的采样密度直接决定所有张量的大小。一个带宽为 B 的网格S2 采样点大约是 2B^2 个SO3 采样点大约 2B^3 个。随着网络层数加深SO3 域的特征图张量增长非常快。我的建议是先用 B16 把整个 pipeline 跑通确认损失函数能下降再逐步调到 B32 或 48。如果在训练初期就发现显存占用异常先查是不是某个 SO3Conv 层的输出分辨率设置过高了。另外批量大小不要一味追求大Spherical CNN 在 batch size 较小时反而更容易稳定收敛这大概和频域权重更新的噪声特性有关。6.3 常见报错和修复建议报错或症状可能原因修复建议初始化后 loss 为 NaN频域权重初始化方差过大检查初始化方差公式按 1/(2l1)^2 级别缩小等变测试误差过大某个层破坏了 (l, m) 索引结构逐层打印输出定位破环等变性的层训练正常但测试集掉点严重带宽太小导致信息丢失在数据输入前做低通滤波或调大 B显存不足SO3 特征图分辨率过高降低采样密度或减少 SO3Conv 层数梯度消失靠近输入的网络层权重几乎不变网络太深且激活函数选择不佳减少层数或改用残差连接连接这里要特别解释一下第一个问题。复数的实部和虚部如果在初始化时独立采样各自方差不一致就会出现某个频带能量特别大而另一个特别小。能量大的频带会主导梯度让优化过程退化。我建议初始化后先做一次前向小批量推理检查各频率分量的平均幅度是否相近再继续训练。6.4 从源码出发的扩展思路读完这套源码之后你会发现它其实抽象出了一整套“几何频域卷积”的框架。同样的设计思路可以迁移到很多场景球面上的气象数据预测、三维分子结构的性质回归、全景图场景理解甚至 3D 网格上的等变特征提取。我自己做的一个扩展是把 S2FT 预计算矩阵用在轻量化模型上——因为变换矩阵是固定的可以预先离线算好推理时只需要一个小矩阵乘法这个思路和模型压缩天然兼容。还有一次把旋转等变特征和注意力机制结合在球面分类上拿到了比原模型更好的精度同时仍然保持等变性质。这些扩展的基石都是这套源码里的变换模块所以把 polar.py 和 s2_ft.py 吃透收益率其实比一般想象的高很多。最后再分享一个小经验读这套源码别急着看卷积层先把“变换矩阵怎么来的”搞清楚。只要理解了等矩形网格到谐波系数的线性变换后面所有层的设计都能顺藤摸瓜理出来。数学公式看不懂的部分就回到代码里打印张量 shape对照每个索引的含义过一遍基本就能落地了。