Rust Candle 的 3D 卷积层

Candle 框架一直以来都没有推进 3D 算子的实现。

比较常见的 3D 算子比如： BatchNorm3d, Conv3d, Dropout3d, MaxPool3d, ConvTranspose3d 通通没有实现。

没有 3D 算子 就意味着 Candle 框架无法处理有关视频数据的任务，例如视频分类、视频对象检测、VLM 等。这也是为什么 Candle 框架一直没有 Qwen2_VL,GLM 等视觉大模型入库的原因。

不过，最近发现有人使用 Candle 在 Qwen2_VL 上实现了一个在时间维度的卷积核大小固定为2的 Conv3d 层。

2D卷积模拟3D卷积

思路大概如下，将3D卷积分解为两个并行的2D卷积：

// 原始3D卷积权重形状：[输出通道, 输入通道/组数, 时间=2, 高度, 宽度]
// 权重形状示例：[64, 32, 2, 3, 3]  ← 时间维度为2

// 分解为两个2D卷积权重：
let w1 = ws.i((.., .., 0, .., ..))?;  // 形状：[64, 32, 3, 3] ← 第1帧的卷积核
let w2 = ws.i((.., .., 1, .., ..))?;  // 形状：[64, 32, 3, 3] ← 第2帧的卷积核

通过对拆分输入张量的时间维度，对每一帧分别进行2D卷积，再把结果合并起来。

向前传播大概如下：

fn forward(&self, xs: &Tensor) -> Result<Tensor> {
    // 1. 拆分输入张量的时间维度
    let xs1 = xs.i((.., .., 0, .., ..))?;  // 第1帧的特征图
    let xs2 = xs.i((.., .., 1, .., ..))?;  // 第2帧的特征图
    
    // 2. 对每一帧分别进行2D卷积
    let out1 = self.conv2d_1.forward(&xs1)?;
    let out2 = self.conv2d_2.forward(&xs2)?;
    
    // 3. 合并结果（逐元素相加）
    // 4. 恢复时间维度（从 [B, C, H, W] → [B, C, 1, H, W]）
    (out1 + out2)?.unsqueeze(2)
}

局限性

这种实现是注定没有灵活性的，因为它只能处理时间维度为2的情况。目前，Candle框架在核心库中对3D卷积算子的支持进展缓慢，尚未有通用解决方案。