PyTorch高维张量乘法与conv2d/conv3d维度语义解析

拿到一个(N, C, H, W)的张量，想跟另一个张量乘一下，结果报了个维度不匹配；换成conv2d又提示 expected 4-dimensional input，可你手里明明是五维的视频数据。这类报错我在带新人的时候见过太多次，问题往往不在 API 记不住，而在于没有把"维度语义"和"算子语义"对上号。PyTorch 里张量的数学运算这一块，二维矩阵乘法大家都很熟，一旦升到四维、五维，形状怎么读、该选matmul还是bmm、conv2d的四个维和conv3d的五个维分别代表什么、输出尺寸怎么手算、显存会不会炸，这些细节就全堆在一起了。

这篇内容就围绕高维张量的乘法与卷积展开，重点落在四维张量走conv2d、五维张量走conv3d这两条主线，顺带把高维乘法（matmul、bmm、einsum）的选型讲透。适合已经能跑通 MNIST、开始接触注意力机制、视频理解、医学影像体素分割这类任务的同学；也适合那些"代码能跑但不知道自己写了多少 FLOPs"的朋友。我会把形状推导、参数量手算、显存估算、验证卷积本质的写法都摆出来，你可以直接抄去改。

1. 先把维度的语义读明白，再谈运算

1.1 四维和五维分别在什么场景下出现

二维张量是矩阵，三维张量最常见的读法是(batch, seq_len, feature)，到四维就有两套主流排布：图像领域的(N, C, H, W)和注意力领域的(B, heads, L, d)。这两个虽然都是四维，但语义完全不同——前者后两维是空间坐标，后者后两维是"每个头内部的 Query/Key 维度"。你选什么算子，本质上取决于你把哪两维看成"要做矩阵乘的那两维"，以及你把哪些维看成"批次"。

五维张量在工程里出现得更集中一些，典型是(N, C, D, H, W)的体数据：视频片段、CT/MRI 体素、点云体素化、3D 检测的 BEV 特征堆叠。这里的D是深度（时间帧数或切片数），它和H、W的地位在conv3d里是平权的——卷积核在三个方向上同时滑动。这一点很关键，因为很多人下意识把D当成 batch 维去处理，结果核在时间轴上不共享，模型表达能力直接掉一截。

我习惯在写代码前先在手边记一行注释，比如# x: (N, C, D, H, W) N=8 C=3 D=16 H=W=64，这个习惯救过我很多次。因为 PyTorch 的报错信息只会告诉你"期望 5 维但拿到 4 维"，不会告诉你少的是哪一维。

1.2 乘法与卷积在工程里的分工边界

这两类运算经常被混着用，但职责其实很清楚。乘法负责"特征维度的线性组合"：把每个位置上的 C 维向量映射到 K 维，或者算注意力分数，或者做投影。它不关心空间邻域，每个位置独立计算。卷积负责"空间邻域的加权聚合"：同一个核在 H、W（甚至 D）上滑动，权重共享，因此参数量与输入尺寸无关，只与核大小和通道数有关。

这个分工带来一个实操上的判断标准：如果你发现自己在用一个大矩阵乘一个(C, C)的权重，而且这个权重在不同空间位置上是同一份，那你其实在做 1x1 卷积；反过来，如果你需要每个位置有独立的权重，那才是真正的矩阵乘。1x1 卷积在实现上就是一次matmul（或者更准确地说是bmm），这也是为什么 MobileNet 那类结构里 1x1 卷积被叫作 pointwise convolution——它只做通道混合，不动空间。

再往深一层，卷积在底层实现里通常会被展开成矩阵乘（im2col 那一套），所以你看到的conv2d和matmul在计算图上并不是泾渭分明。理解这一点，后面算 FLOPs 和估显存的时候会很顺。

2. 高维乘法：mm、bmm、matmul、einsum 怎么选

2.1 四个 API 的能力边界对照

先把这张表记住，能省掉大量试错：

API	输入维度要求	广播支持	典型用途
torch.mm	严格 2D × 2D	无	纯矩阵乘，最省心
torch.bmm	严格 3D × 3D，batch 维必须一致	无	批量矩阵乘，形状最明确
torch.matmul/@	任意维，最后两维做矩阵乘	有	通用场景，写起来最短
torch.einsum	任意维	无（靠下标表达）	多维收缩，可读性最好
torch.tensordot	任意维	无	指定两个维度集合做收缩

选型逻辑我一般是这么走的：形状简单、维度确定，直接用@；需要明确表达"这是 B 个独立矩阵"，用bmm，因为它会强制校验 batch 维，出错更早暴露；涉及三个以上维度同时收缩（比如bhld,bhsd->bhls这种），用einsum，可读性比手写transpose + matmul高一个量级。

有个细节值得说：torch.mm在输入不是 2D 时会直接抛错，而matmul会自动帮你处理。新手阶段用matmul舒服，但到了需要严格控制内存和性能的阶段，bmm反而更好排错——因为它不会偷偷广播，形状对不上就是不对上。

2.2 广播规则：高维乘法最容易翻车的地方

matmul的规则可以一句话概括：最后两维按矩阵乘处理，前面的所有维度按广播规则对齐。举个例子，(8, 32, 64) @ (64, 16)得到(8, 32, 16)，这里(64, 16)被当成"所有 batch 共享的权重"自动扩展。而(8, 32, 64) @ (8, 64, 16)走的是 batch 对齐，结果是(8, 32, 16)。

麻烦出在后面这种：(8, 1, 32, 64) @ (8, 5, 64, 16)。前面的(8, 1)和(8, 5)按广播规则对齐，1扩展成5，结果是(8, 5, 32, 16)。这个行为本身是对的，但它会悄悄复制数据，内存占用按倍数涨上去。我在处理多头注意力的时候踩过一次，heads维不小心写成了 1，结果一个(B, 1, L, d)和(B, heads, d, d)相乘，显存直接多吃了好几倍，还没报错。

注意：广播虽然方便，但它不会提示你"我复制了数据"。凡是看到某两维大小是 1，而另一个张量对应维不是 1，就要有意识地去估一下实际显存峰值。

2.3 用 einsum 表达多维收缩，顺便自检形状

einsum最大的价值不在性能，在于把形状约束写进了字符串里，等于自带一份文档。比如四维张量做投影：

import torch B, H, W, C, K = 2, 8, 8, 16, 32 x = torch.randn(B, H, W, C) w = torch.randn(C, K) y1 = torch.einsum('bhwc,ck->bhwk', x, w) y2 = x @ w # 直接 matmul 也能得到同样结果 y3 = torch.matmul(x, w) print(y1.shape, y2.shape, y3.shape) # 都是 (2, 8, 8, 32) print(torch.allclose(y1, y2), torch.allclose(y1, y3))

这段代码我经常拿来给新人演示，因为它同时讲清了两件事：第一，matmul对四维张量确实只动最后两维；第二，einsum的下标写法是"输入有哪些字母、输出保留哪些字母"的直接映射，c没出现在输出里就说明这一维被收缩掉了。

再看一个更贴近实战的注意力分数计算：

B, heads, L, d = 4, 8, 128, 64 q = torch.randn(B, heads, L, d) k = torch.randn(B, heads, L, d) scores_einsum = torch.einsum('bhld,bhsd->bhls', q, k) scores_matmul = q @ k.transpose(-2, -1) print(scores_einsum.shape) # (4, 8, 128, 128) print(torch.allclose(scores_einsum, scores_matmul, atol=1e-5))

transpose(-2, -1)把(B, heads, L, d)变成(B, heads, d, L)，再做matmul，逻辑上等价。这里transpose之后张量在内存上不连续，matmul内部会处理，但如果你后面还要接view，就得先.contiguous()。这是下一节要重点说的问题。

einsum在 PyTorch 里默认会尝试走opt_einsum的路径做收缩顺序优化，涉及三个以上张量的链式收缩时，性能差异会很明显。但两个张量的简单情形，它和matmul基本打平，甚至因为少了 kernel 特化而略慢一点点。所以我的建议是：两三个张量的简单乘法用matmul，多维收缩用einsum，别为了可读性牺牲不必要的性能。

3. 四维张量上的乘法实操与显存账

3.1 把四维拆成批处理矩阵的两种写法

假设你有一批特征图(N, C, H, W)，想对每个空间位置做通道变换，权重是(C, K)。有两种写法，语义一样，写法差别很大：

N, C, H, W, K = 8, 64, 32, 32, 128 x = torch.randn(N, C, H, W) w = torch.randn(C, K) # 写法一：直接 matmul，依赖最后两维是矩阵乘的规则 out_a = (x.transpose(1, -1) @ w).transpose(1, -1) # (N, K, H, W) print(out_a.shape) # 写法二：reshape 成三维，走 bmm x_flat = x.permute(0, 2, 3, 1).reshape(N, H * W, C) # (N, H*W, C) out_b = torch.bmm(x_flat, w.expand(N, C, K)) # (N, H*W, K) out_b = out_b.reshape(N, H, W, K).permute(0, 3, 1, 2) print(out_b.shape)

写法一的问题在于transpose之后张量不连续，matmul内部会做一次隐式的内存重排，代价不小。写法二里permute + reshape会触发一次真实的数据搬运（reshape在不连续时会自动contiguous），显存峰值会短暂翻倍。这两种写法都不是白给的。

实测下来，如果只是通道变换，直接换成nn.Conv2d(C, K, kernel_size=1)往往更快，因为 cuDNN 对 1x1 卷积有专门优化，而且不需要你自己处理布局。这点反直觉，但确实是很多人的经验：不要用纯matmul去手动实现 1x1 卷积，除非你有别的理由必须这么做（比如权重需要动态生成、或者要做一些核函数上的特殊处理）。

3.2 连续性、view 与 reshape 的恩怨

view要求张量在内存上连续，reshape不要求，必要时会自动拷贝。这个区别在处理四维、五维张量时会被无限放大，因为高维张量的permute、transpose几乎必然破坏连续性。

x = torch.randn(2, 3, 4, 5) y = x.permute(0, 2, 3, 1) # (2, 4, 5, 3) print(y.is_contiguous()) # False try: y.view(2, -1) except RuntimeError as e: print('view 失败:', e) print(y.reshape(2, -1).shape) # reshape 能过，内部帮你拷贝了一次 print(y.contiguous().view(2, -1).shape) # 显式拷贝后再 view，同样能过

reshape和contiguous().view()结果一样，但语义不同：前者是"我不在乎有没有拷贝，给我个形状对的就行"，后者是"我明确知道这里要付一次拷贝的代价"。在显存紧张的场景下，我倾向于写contiguous().view()，因为这样代码里每一次拷贝都是显式的，回头查显存峰值的时候有据可依。

提示：permute之后立刻reshape，是最常见的隐式拷贝来源。如果你的显存估算和实测差得比较多，先去代码里搜一遍这种组合。

3.3 参数量与 FLOPs 的手算方法

不管乘法还是卷积，估算的思路是一致的：乘法次数 = 输出元素个数 × 每个输出元素需要的乘加次数。

对于(N, C, H, W) @ (C, K)这种逐位置变换，输出是N*H*W*K个元素，每个元素需要C次乘加，总 FLOPs（按乘加算两次浮点运算）约2 * N * H * W * C * K。拿N=8, C=64, H=W=32, K=128代入：2 * 8 * 32 * 32 * 64 * 128 ≈ 1.34e8，也就是 1.34 亿次浮点运算。这个量级在现代 GPU 上属于"瞬间完成"，所以你很难从耗时上感知到写法差异，得靠 profiler 才看得出来。

卷积的算法稍有不同：输出元素个数是N * C_out * H_out * W_out，每个输出元素要扫过(C_in / groups) * kH * kW个输入值，所以

FLOPs ≈ 2 * N * C_out * H_out * W_out * (C_in / groups) * kH * kW

举例：N=8, C_in=64, C_out=128, kH=kW=3, H_out=W_out=32, groups=1，得到2 * 8 * 128 * 32 * 32 * 64 * 9 ≈ 1.2e9，约 12 亿次。比刚才那个矩阵乘大了一个数量级。这就是为什么 3x3 卷积是绝大多数 CNN 的计算主力，而 1x1 卷积常被用来做通道升降维。

4. conv2d：四维张量的每个位置都要对上号

4.1 NCHW 各维含义与权重形状

nn.Conv2d的输入形状是(N, C_in, H, W)，权重形状是(C_out, C_in / groups, kH, kW)，偏置是(C_out,)，输出(N, C_out, H_out, W_out)。这四个字母我建议你背下来，因为报错信息里经常只写 "expected input..."，不解释。

有一个很常见的误区：把(N, H, W, C)这种通道在最后的排布直接丢给Conv2d。它不会报错——因为维度数量对得上——但算出来的结果完全是乱的。通道维被当成了空间维，核在错误的方向上滑动。这类 bug 最恶心的地方在于不报错、loss 也能降一点，你很难第一时间发现。防御手段很简单：在模型forward的开头加一句assert x.shape[1] == self.in_channels，或者在__init__里把期望形状记下来。

import torch import torch.nn as nn conv = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, stride=1, padding=1) print(conv.weight.shape) # (16, 3, 3, 3) print(conv.bias.shape) # (16,) x = torch.randn(4, 3, 32, 32) y = conv(x) print(y.shape) # (4, 16, 32, 32)

padding=1、kernel_size=3、stride=1这个组合是"保持空间尺寸不变"，公式推一下就明白，下面单独说。

4.2 输出尺寸公式与有效核大小的换算

标准公式是：

H_out = floor((H_in + 2 * padding - dilation * (kH - 1) - 1) / stride + 1)

dilation * (kH - 1) - 1这一串不好记，我习惯先算有效核大小k_eff = dilation * (kH - 1) + 1，公式就变成floor((H_in + 2p - k_eff) / s + 1)，直观很多：相当于拿一个大小为k_eff的普通核，中间那些空洞只是不参与计算。

下面这张表是我平时手算时对照的，列举了几种常见配置：

输入 H	kernel	stride	padding	dilation	H_out
224	7	2	3	1	112
224	3	2	1	1	112
224	3	1	1	1	224
224	3	1	2	2	224
224	5	1	0	1	220
224	3	1	1	4	224

验证第一个：k_eff = 1*(7-1)+1 = 7，(224 + 6 - 7)/2 + 1 = 112.5，向下取整是 112。注意这个.5——224+6=230，230-7=223，223/2=111.5，111.5+1=112.5，取整 112。这意味着输入边缘有半个核的跨度被丢掉了，这就是为什么padding=3, kernel=7, stride=2之后输出是 112 而不是 113。

常用场景里还有一个padding='same'，它只支持stride=1，内部会自动算出需要的 padding（可能是非对称的）。用起来省事，但要注意：非对称 padding 在导出到某些推理引擎时可能不支持，我一般还是手写padding=k//2，可控性更好。

4.3 分组、空洞与 padding_mode 的实操细节

groups这个参数很多人不敢碰，其实逻辑很简单：把输入通道和输出通道各切成groups份，第 i 份输出只跟第 i 份输入卷。约束是C_in % groups == 0且C_out % groups == 0，权重的第二维变成C_in / groups。groups=C_in时就是深度可分离卷积的 depthwise 部分，参数量直接砍到C_in * kH * kW。

dw = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, groups=64) print(dw.weight.shape) # (64, 1, 3, 3) print(sum(p.numel() for p in dw.parameters())) # 64*9 + 64 = 640

对比普通 3x3 卷积的64*64*9 + 64 = 36928，差了两个数量级。这就是为什么移动端模型对 depthwise 情有独钟。但实话说，depthwise 卷积在 GPU 上的计算密度低，理论 FLOPs 降了，实际耗时不一定降那么多，选型时得实测，别光看算力数字。

dilation的坑在于和padding的配合。要让空洞卷积保持空间尺寸，padding 得按dilation * (k-1) / 2算，不是k // 2。比如kernel=3, dilation=2，有效核大小是 5，padding 要给 2 才能保持尺寸。写错了不会报错，但特征图会一路缩水，到最后view的时候才炸。

padding_mode支持'zeros'、'reflect'、'replicate'、'circular'，默认补零。补零在图像边缘会引入"硬边界"，做密集预测（分割、深度估计）时边缘质量差，换成'reflect'通常有小幅提升。这个点在论文里很少写，但确实是我做分割任务时的常规操作。

5. conv3d：五维张量的多一个维度意味着什么

5.1 维度排列与内存布局的实际影响

nn.Conv3d的输入是(N, C_in, D, H, W)，权重(C_out, C_in/groups, kD, kH, kW)，输出(N, C_out, D_out, H_out, W_out)。三个空间方向平权，核在 D、H、W 上同时滑动。

内存占用是这里最直接的痛点。一个N=8, C=3, D=16, H=W=224的输入，元素个数是8*3*16*224*224 ≈ 1.93e7，fp32 下约 77MB。第一层卷到 64 通道，输出就变成8*64*16*224*224 ≈ 4.1e8个元素，约 1.57GB。单层就吃掉一张卡的显存，这还没算反向传播存下来的中间激活。

所以做 3D 任务的第一个决策不是网络结构，而是输入的体素尺寸。常见做法是把 D 方向降采样（抽帧）、把 H/W 裁剪成 patch（比如 128x128）、或者干脆用(1, 3, 3)的核先把 D 当独立帧处理。我见过有人在 16GB 卡上直接怼(1, 1, 128, 512, 512)的 CT 体数据进去，结果 OOM 之后以为是框架的问题，其实是显存账从来没算过。

import torch import torch.nn as nn def nbytes(*shape, dtype=torch.float32): n = 1 for s in shape: n *= s return n * torch.finfo(dtype).bits // 8 x = torch.randn(2, 3, 16, 64, 64) print(f'输入: {nbytes(2,3,16,64,64)/1024/1024:.2f} MB') # 约 1.5 MB conv3d = nn.Conv3d(3, 16, kernel_size=3, stride=1, padding=1) y = conv3d(x) print(y.shape) # (2, 16, 16, 64, 64) print(f'输出: {nbytes(2,16,16,64,64)/1024/1024:.2f} MB') # 约 8 MB

单看一层不大，但真实的 3D 网络要堆十几层，中间还要存反向传播的激活，实际占用是逐层累加的。做显存规划的时候我习惯按"参数量 × 4 字节 + 激活值 × 4 字节 × 层数"粗估，宁可估高。

5.2 输出尺寸的逐维推导与常见误算

三个方向用同一个公式，但每一维的输入长度不同，结果也可能不同：

D_out = floor((D_in + 2*pD - dD*(kD-1) - 1) / sD + 1) H_out = floor((H_in + 2*pH - dH*(kH-1) - 1) / sH + 1) W_out = floor((W_in + 2*pW - dW*(kW-1) - 1) / sW + 1)

nn.Conv3d的stride、padding、dilation都支持元组，分别对应三个方向。这里最常见的错误是只传一个整数——它会被自动广播到三个方向，如果你的 D 和 H/W 尺寸差得多，输出尺寸就可能对不上预期。比如 D=16、H=W=64，kernel=3, stride=1, padding=1各方向都保持原尺寸，没问题；但如果D=16而你想要 D 方向做 2 倍下采样、H/W 保持不变，就得写stride=(2, 1, 1)，写成stride=2会把 H/W 也砍一半。

还有padding='same'在Conv3d里也是支持的，同样只限stride=1。D 方向如果是奇数长度、kernel 是偶数（比如kernel=(2,3,3)），'same'的 padding 在 D 上会是非对称的，个别推理后端处理起来有差异，我一般避免用偶数核。

5.3 用 (1,3,3) 核把 conv3d 退化成逐帧 conv2d

这是一个非常实用的技巧，尤其是在你还没想好要不要用真 3D 卷积、但代码里已经写死了Conv3d的时候。

conv3d_fake = nn.Conv3d(3, 16, kernel_size=(1, 3, 3), stride=1, padding=(0, 1, 1)) x = torch.randn(2, 3, 8, 32, 32) y = conv3d_fake(x) print(y.shape) # (2, 16, 8, 32, 32) # 等价写法：把 D 当 batch，走 conv2d conv2d_equiv = nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1) conv2d_equiv.weight.data = conv3d_fake.weight.data[:, :, 0].clone() conv2d_equiv.bias.data = conv3d_fake.bias.data.clone() x_2d = x.reshape(2 * 8, 3, 32, 32) y_2d = conv2d_equiv(x_2d).reshape(2, 8, 16, 32, 32).permute(0, 2, 1, 3, 4) print(torch.allclose(y, y_2d, atol=1e-6))

注意最后那个permute：conv2d的输出是(N*D, C_out, H, W)，reshape 回(N, D, C_out, H, W)之后要把 C 换到第二位。这一步忘记 axis 顺序是踩坑高发区，allclose一定要写。

这个等价关系还有一个好处：参数量对得上。(1,3,3)的 3D 核参数是16*3*1*3*3 = 432，2D 核是16*3*3*3 = 432，完全一致。如果你的模型前期用伪 3D 跑 baseline、后期想换真 3D，权重是可以直接迁移的（在 D 方向复制核值再归一化）。

6. 从构造到校验：一套完整的实操流程

6.1 手写 unfold 验证 conv2d 本质就是乘加

我一直觉得，理解卷积最直接的方式是手动把它拆成矩阵乘。F.unfold就是把滑动窗口按列展开的操作，配合一次矩阵乘就能复现conv2d的结果：

import torch import torch.nn.functional as F torch.manual_seed(0) x = torch.randn(1, 1, 4, 4) w = torch.randn(1, 1, 3, 3) b = torch.randn(1) conv_out = F.conv2d(x, w, b) patches = F.unfold(x, kernel_size=3) # (1, 1*3*3, 4) 4 个滑动位置 manual = w.view(1, -1) @ patches # (1, 1, 4) manual = manual + b.view(1, 1) manual = manual.view(1, 1, 2, 2) print(torch.allclose(conv_out, manual, atol=1e-6)) # True

走一遍这个流程，你会对两件事有实感：第一，unfold展开后的通道数是C * kH * kW，这就是 im2col 的显存放大倍数；第二，卷积的"权重共享"体现在w.view(1, -1)是所有滑动位置共用的同一份权重，如果每个位置用不同权重，那就退化成全连接了。

顺带算一下 im2col 的显存：N * C * kH * kW * L * 4字节，L = H_out * W_out。N=8, C=64, k=3, H_out=W_out=56时是8*64*9*3136*4 ≈ 57.8MB。这就是为什么大卷积层在训练时显存会突然冲高——cuDNN 内部不一定真的物理展开，但在某些算法下确实会。

6.2 conv3d 的完整跑通脚本

把前面几块拼起来，给一段可以直接跑的最小代码：

import torch import torch.nn as nn import time torch.manual_seed(42) device = 'cuda' if torch.cuda.is_available() else 'cpu' model = nn.Sequential( nn.Conv3d(3, 16, kernel_size=3, stride=1, padding=1), nn.BatchNorm3d(16), nn.ReLU(inplace=True), nn.Conv3d(16, 32, kernel_size=3, stride=2, padding=1), nn.BatchNorm3d(32), nn.ReLU(inplace=True), ).to(device) x = torch.randn(4, 3, 16, 64, 64, device=device) # 先算一遍理论输出形状 def out_size(inp, k, s, p, d=1): return (inp + 2 * p - d * (k - 1) - 1) // s + 1 print('第1层后 D/H/W:', out_size(16,3,1,1), out_size(64,3,1,1), out_size(64,3,1,1)) print('第2层后 D/H/W:', out_size(16,3,2,1), out_size(64,3,2,1), out_size(64,3,2,1)) with torch.no_grad(): y = model(x) print('实际输出:', tuple(y.shape)) # 参数量与显存 n_param = sum(p.numel() for p in model.parameters()) print(f'参数量: {n_param:,}') print(f'输入显存: {x.numel()*4/1024/1024:.2f} MB') print(f'输出显存: {y.numel()*4/1024/1024:.2f} MB') # 简单计时（第一次含 kernel 编译，取第二次） if device == 'cuda': torch.cuda.synchronize() t0 = time.time() with torch.no_grad(): model(x) torch.cuda.synchronize() print(f'前向耗时: {(time.time()-t0)*1000:.2f} ms')

第二层stride=2之后，D 从 16 变 8、H/W 从 64 变 32，和我手算的一致。这种"先写公式算、再拿 shape 对"的习惯，比出问题后 debug 快得多。

注意：计时一定要在torch.cuda.synchronize()之后取，否则拿到的是 kernel 入队时间，不是真实耗时。我第一次做性能对比时就栽在这上面，测出来的数字比实际小十倍。

6.3 布局与后端开关的两个小优化

如果你的卷积层通道数比较多（比如 128 以上）、核是 3x3，可以试试channels_last：

model = model.to(memory_format=torch.channels_last) x = x.to(memory_format=torch.channels_last_3d) # 3D 用 channels_last_3d

对conv2d/conv3d来说，NHWC 布局在某些 GPU 上能让 Tensor Core 的利用率更高，提速 10%~30% 是常见的。但它不是万能的，小通道数（比如 3 通道输入）反而可能变慢，得实测。

另一个开关是torch.backends.cudnn.benchmark = True。它会让 cuDNN 在第一次遇到某个形状时搜索一遍算法，后续同形状直接复用。适合输入尺寸固定的场景（比如训练时固定 crop 到 224x224），如果每个 batch 的输入尺寸都不一样，搜索开销会拖慢整体，这时候保持默认False更合适。

7. 常见报错与排查速查

这些问题我基本都亲手遇到过，整理成表方便对照：

报错 / 现象	根因	处理方式
expected 4-dimensional input	少了 batch 维，或误传五维	单个样本用x.unsqueeze(0)；五维数据检查是否用错算子
view size is not compatible	permute后不连续	改reshape或显式.contiguous()
bmm报 batch 不一致	两个张量的第 0 维不同	先expand对齐，或改用matmul走广播
输出尺寸比预期小	忘了dilation对有效核的放大	用k_eff = d*(k-1)+1重算
groups报错	通道数不能被groups整除	调整C_in/C_out或改groups
conv3d训练直接 OOM	体数据尺寸过大	降 D、裁 patch、减小 batch、用 checkpoint
结果全对但精度差很多	通道维放错位置，如 NHWC 喂给 Conv2d	检查shape[1]是否等于in_channels
推理比训练慢	没切eval()、没关梯度	model.eval()+torch.no_grad()
channels_last反而变慢	通道数太小，布局转换开销占主导	通道数 < 32 时别开

再补几个不那么"报错"但很烦人的问题：

第一个是padding='same'配stride>1。它会直接抛错，提示不支持。这时候要么改stride=1，要么自己算 padding。我一般后者，公式就是max((k_eff - 1) // 2, 0)加上stride带来的额外偏移，虽然麻烦，但行为完全可控。

第二个是输入尺寸不是 stride 的整数倍。比如H=225配stride=2，最后一行像素会被丢掉。做分割任务时这会导输出和标签对不齐，loss计算出 NaN 或者形状不匹配。解决办法是在DataLoader里统一做 padding 到 32 的倍数，别指望网络自己处理。

第三个是混合精度下的卷积。autocast会把conv的输入转成 fp16，但bias参与累加时如果被转成 fp16，大通道数下会有精度损失。PyTorch 的autocast对卷积有白名单处理（保持 fp32 累加），一般不用管；但如果你自己手写了matmul版的 1x1 卷积，就得手动with torch.autocast(...)或者显式float()一下，不然数值会飘。

8. 我在实际项目里的一些体会

踩过几次坑之后，我形成了一个固定动作：任何高维张量进模型之前，先把 shape、dtype、device 打印一遍，再用一行公式算出期望输出，最后才forward。多花十秒钟，能省掉半小时的 debug。尤其是conv3d，形状错一位可能不报错，但训练几轮之后 loss 不降，你会怀疑人生。

另一个体会是关于"选matmul还是选conv"的。早期我总想用纯矩阵运算去替代卷积，觉得这样"更底层、更可控"。后来发现，除非你有非常特殊的权重生成逻辑，否则交给 cuDNN 处理永远是更优解——它的 kernel 库针对不同形状做了大量特化，手写很难追上。真正需要你手写matmul的场景其实很窄：注意力机制里的 QK^T、动态生成的投影权重、以及一些图结构上的消息传递。

最后一个技巧，关于conv3d的调试。如果你的 3D 模型训练不收敛，可以先把它退化成(1,3,3)核，等价于逐帧 2D 卷积，跑通之后再逐步把 D 方向的核放大到 3、5。这样能把"数据管线问题"和"3D 建模问题"分开排查，效率高很多。我在一个体素分割任务上就是靠这个办法定位到问题出在数据增强的空间翻转上——3D 翻转把左右方向搞反了，标签没跟着翻，模型自然学不动。