Rust Candle 框架与 Pytorch nn 模块网络层转换(1)

之前大体上了解完 pytorch 与 candle 之间张量的一些操作，现在来到了图的构建块: nn 模块的一些相同或者不同的地方。

torch.nn 模块是构建和训练神经网络的核心工具包，里面有大量的网络层、损失函数、优化器等的定义。

由于 candle 重推理而不是训练，因此这里重点讲讲 candle 和 pytorch 网络层之间的相互转换。

表格预览

✅ 表示有对应实现
🚫 表示无对应实现
☢️ 表示有代替实现

功能	pytorch	candle	是否实现
顺序容器	Sequential	Sequential	✅
1D 卷积	nn.Conv1d	conv1d/conv1d_no_bias	✅
2D 卷积	nn.Conv2d	conv2d/conv2d_no_bias	✅
3D 卷积	nn.Conv3d	未实现	🚫
1D 转置卷积	nn.ConvTranspose1d	conv_transpose1d conv_transpose1d_no_bias	✅
2D 转置卷积	nn.ConvTranspose2d	conv_transpose2d conv_transpose2d_no_bias	✅
3D 转置卷积	nn.ConvTranspose3d	未实现	🚫
1D 最大池化	nn.MaxPool1d	未实现	🚫
2D 最大池化	nn.MaxPool2d	max_pool2d max_pool2d_with_stride	✅
3D 最大池化	nn.MaxPool3d	未实现	🚫
1D 平均池化	nn.AvgPool1d	未实现	🚫
2D 平均池化	nn.AvgPool2d	avg_pool2d avg_pool2d_with_stride	✅
3D 平均池化	nn.AvgPool3d	未实现	🚫
逐元素应用修正线性单元函数	nn.ReLU	relu	✅
逐元素应用 ReLU6 函数	nn.ReLU6	relu6	✅
逐元素应用随机泄漏修正线性单元函数	nn.RReLU	未实现	✅
应用高斯误差线性单元函数	nn.GELU	gelu	✅
Sigmoid 函数	nn.Sigmoid	sigmoid	✅
Sigmoid 线性单元 (SiLU) 函数	nn.SiLU	silu	✅
逐元素应用双曲正切 (Tanh) 函数	nn.Tanh	tanh	✅
逐元素应用指数线性单元 (ELU) 函数。	nn.ELU	elu	✅

开始之前

为了对比两边的操作是否等效，我们需要找到一个方法固定并共享 pytorch 和 candle 两个框架的测试权重。因为有一些参数如果初始化时不设定权重，会默认随机分配权重值。比如下面这个代码：

m = nn.Conv1d(3, 16, 3)
input = torch.ones(1, 3, 224,dtype=torch.float32)
output = m(input)

上面的代码由于没有初始化权重，事实上 output 的结果每次都不一样，因此为了方便复现，我们需要保存权重，把随机的参数都固定下来。

最好的办法是把python 的权重保存到文件，然后在 candle 中复现。 pytorch 保存权重，这里使用 safetensors 格式：

from safetensors.torch import save_model
m = nn.Conv1d(3, 16, 3, stride=2,bias=False)
save_model(m, "model.safetensors")

candle 中，直接读取model.safetensors 权重即可复现，同时这也是最接近使用 candle 实现 pytorch 模型的方式：

let vb = unsafe {
        VarBuilder::from_mmaped_safetensors(&["./model.safetensors"], DType::F32, &Device::Cpu)?
    };

Sequential

Sequential 只是容器，本身并无任何功能，而且也不会影响权重的加载。因此我们可以使用 candle 官方的，也可以自己定义。 pytorch

seq = torch.nn.Sequential(
        nn.Conv2d(3, 64, 3, 1, 1),
        nn.ReLU(),
        nn.Conv2d()
    )

candle 官方的定义本质上也是内部维护一个 vec：

    let mut seq = candle_nn::seq();
    seq.add(line);

因此如果官方的不满足（为什么不满足? rust 由于类型系统存在中没有那么方便，比如官方的实现类型是：Vec<Box<dyn Module>> ），可以自己实现：

use candle_core::{ Module, Tensor, Result };

#[derive(Debug, Clone)]
pub struct Sequential<T: Module> {
    layers: Vec<T>,
}

pub fn seq<T: Module>(cnt: usize) -> Sequential<T> {
    let v = if cnt == 0 { vec![] } else { Vec::with_capacity(cnt) };
    Sequential { layers: v }
}

impl<T: Module> Sequential<T> {
    pub fn len(&self) -> usize {
        self.layers.len()
    }

    pub fn is_empty(&self) -> bool {
        self.layers.is_empty()
    }

    pub fn push(&mut self, layer: T) {
        self.layers.push(layer);
    }

    pub fn add(&mut self, layer: T) {
        self.layers.push(layer);
    }
}
impl<T: Module> Module for Sequential<T> {
    fn forward(&self, xs: &candle_core::Tensor) -> Result<Tensor> {
        let mut xs = xs.clone();
        for layer in self.layers.iter() {
            xs = xs.apply(layer)?;
        }
        Ok(xs)
    }
}

1D 卷积

conv1d 主要用于处理 序列数据 ，其核心是通过滑动窗口在单个维度（通常是时间或序列维度）上提取局部特征。比如传感器数据（如温度、股票价格）、医疗信号（如 ECG）、工业监控等。

Conv1d 的公式大概如下：

out(N_i​,C_{out_j}​​)=bias(C_{out_j})+ \sum_{k=0}^{C_{in}​−1}​weight(C{out_j}​​​​,k)⋆input(N_i​,k)

输入(批次，通道，长度)：

(N,C_{in},L_{in})

输出(批次，通道，长度)

(N,C_{out},L_{out})

其中 $L_{out}$ 公式如下：

L_{out} = \frac{L_{in}+2*padding-dilation*(kernel\_size - 1) - 1}{stride} + 1

pytorch

    # 输入通道为3，输出通道为16，卷积核大小为3，步长为2
    # Conv1d(3, 16, kernel_size=(3,), stride=(2,))
    m = nn.Conv1d(3, 16, 3, stride=2,bias=False)
    input = torch.ones(1, 3, 224,dtype=torch.float32)
    output = m(input)
    print(output)
    print(output.size()) #Tensor[dims 1, 16, 111; f32]

candle

    // Conv1d(3, 16, kernel_size=(3,), stride=(2,))
    let cfg = candle_nn::Conv1dConfig{
        stride: 2,
        padding:0,
        dilation:1,
        groups:1,
    };
    let conv1d = candle_nn::conv1d_no_bias(3, 16, 3, cfg, vb)?;
    let x = Tensor::ones((1,3,224), DType::F32, &Device::Cpu)?;
    let y = conv1d.forward(&x)?;
    println!("{y}"); //Tensor[[1, 16, 111], f32]

这里需要注意的是，pytorch 里 bias =False 时，对应的是 candle 中的 conv1d_no_bias，当 bias = True 时，才对应 candle 中的 conv1d。

2D 卷积

conv2d 主要用于主处理 图像数据 ，通过在二维空间（高度和宽度）上滑动滤波器（kernel）提取局部特征。

公式如下：

out(N_i​,C_{out_j}​​) = bias(C_{out_j}) + \sum^{C_{in}-1}_{k=0} weight(C_{out_j},k)⋆input(N_i,k)

输入和输出都是（N,C,H,W）。其中输出的尺寸 $H_{out}$ 和 $W_{out}$ 分别如下:

H_{out} = \frac{H_{in}+2*padding[0]-dilation[0]*(kernel\_size[0]-1)-1}{stride[0]} + 1

W_{out} = \frac{W_{in}+2*padding[1]-dilation[1]*(kernel\_size[1]-1)-1}{stride[1]} + 1

pytorch

    # 输入通道为3，输出通道为16，卷积核大小为3，步长为2
    # Conv1d(3, 16, kernel_size=(3,), stride=(2,))
    m = nn.Conv2d(3, 16, 3, stride=2,bias=False)
    # 输入参数： (N,C,H,W)
    input = torch.ones(1, 3, 224,224,dtype=torch.float32)
    output = m(input)
    print(output)
    print(output.size()) #torch.Size([1, 16, 111, 111])

candle:

    // Conv1d(3, 16, kernel_size=(3,), stride=(2,))
    let cfg = candle_nn::Conv2dConfig{
        stride: 2,
        padding:0,
        dilation:1,
        groups:1,
    };
    let conv1d = candle_nn::conv2d_no_bias(3, 16, 3, cfg, vb)?;
    let x = Tensor::ones((1,3,224,224), DType::F32, &Device::Cpu)?;
    let y = conv1d.forward(&x)?;
    println!("{y}"); // Tensor[[1, 16, 111, 111], f32]

bias 上和 conv1d 同理。

3D 卷积

conv3d 主要用于处理具有 4 维的 时空数据 或 体积数据，比如视频。很遗憾的是，candle 不支持3D 卷积，甚至不支持任何3D 的操作[2025年 3 月前]。

1D 转置卷积/反卷积

一维转置卷积层（也称为反卷积层），主要用于上采样或特征图尺寸恢复。

输入输出同 1D 卷积。输入(批次，通道，长度)：

(N,C_{in},L_{in})

输出(批次，通道，长度)

(N,C_{out},L_{out})

其中长度 $L_{out}$ 的计算公式如下：

L_{out​}=(L_{in}​−1)×stride−2×padding+dilation×(kernel\_size−1)+output\_padding+1

pytorch:

    m = nn.ConvTranspose1d(3, 16, 3, stride=2,bias=False)
    input = torch.ones(1, 3, 224,dtype=torch.float32)
    output = m(input)
    print(output)
    print(output.size()) #torch.Size([1, 16, 449])

candle:

    let cfg = candle_nn::ConvTranspose1dConfig{
        stride: 2,
        padding:0,
        dilation:1,
        groups:1,
        ..Default::default()
    };
    let conv1d = candle_nn::conv_transpose1d_no_bias(3, 16, 3, cfg, vb)?;
    let x = Tensor::ones((1,3,224), DType::F32, &Device::Cpu)?;
    let y = conv1d.forward(&x)?;
    println!("{y}"); // Tensor[[1, 16, 449], f32]

2D 转置卷积/反卷积

二维转置卷积层（也称为反卷积层），主要用于 上采样 或 恢复图像分辨率 。它通过转置卷积操作将低分辨率特征图映射到高分辨率空间，常用于图像生成（如 GAN）、语义分割等任务。和2D 卷积一样，输入和输出都是（N,C,H,W）。

其中输出的高H 和宽 W 公式如下

H_{out​}=(H_{in}​−1)×stride[0]−2×padding[0]+dilation[0]×(kernel\_size[0]−1)+output\_padding[0]+1

W_{out}​=(W_{in}​−1)×stride[1]−2×padding[1]+dilation[1]×(kernel\_size[1]−1)+output\_padding[1]+1

pytorch:

    m = nn.ConvTranspose2d(3, 16, 3, stride=2,bias=False)
    input = torch.ones(1, 3, 224,224,dtype=torch.float32)
    output = m(input)
    print(output)
    print(output.size()) #torch.Size([1, 16, 449, 449])

candle:

    let cfg = candle_nn::ConvTranspose2dConfig{
        stride: 2,
        padding:0,
        dilation:1,
        ..Default::default()
    };
    let conv1d = candle_nn::conv_transpose2d_no_bias(3, 16, 3, cfg, vb)?;
    let x = Tensor::ones((1,3,224,224), DType::F32, &Device::Cpu)?;
    let y = conv1d.forward(&x)?;
    println!("{y}"); // Tensor[[1, 16, 449, 449], f32]

2D 最大池化

最大池化在 pytorch 可以做为层存在，但是 candle 是直接对 tensor 进行操作的。 pytorch

    p = nn.MaxPool2d(3, stride=2)
    input = torch.ones(1, 3, 224,224,dtype=torch.float32)
    output = p(input)

candle

    let x = Tensor::ones((1,3,224,224), DType::F32, &Device::Cpu)?;
    let y = x.max_pool2d_with_stride(3,2)?;

2D 平均池化

2D 平均池化操作和上面 2D 最大池化一样，只是 max_pool2d 变成了avg_pool2d,max_pool2d_with_stride变成了avg_pool2d_with_stride。

激活层

激活层主要是引入非线性变换，使模型能够学习复杂的模式。

pytorch 也有函数的调用方式，这里不展开。

pytorch

    p = nn.ReLU()
    p = nn.GELU()
    p = nn.SiLU()
    p = nn.Tanh()
    p = nn.ELU()
    p = nn.Sigmoid()
    p = nn.ReLU6()
    p = nn.LeakyReLU(0.01)

candle 两种方法都可用，本质上，candle_nn::Activation 实现了 Module，是对前一种方法的封装。

    let x = Tensor::ones((1,3,224,224), DType::F32, &Device::Cpu)?;
    let y = x.relu()?; //relu
    let y = x.gelu()?;//gelu
    let y = x.silu()?;//gelu
    let y = x.tanh()?;//tanh
    let y = x.elu(1f64)?;//elu
    let y = candle_nn::ops::sigmoid(&x)?;//sigmoid
    let y = x.clamp(0f32, 6f32)?; //relu6
    let y = x.relu()?.sqr()?;//Relu2
    let y = candle_nn::ops::leaky_relu(&x, 0.01)?;
	// 或者下面这种方式。
    let activation = candle_nn::Activation::Relu;
    let activation = candle_nn::Activation::Gelu;
    let activation = candle_nn::Activation::Silu;
    let activation = candle_nn::Activation::Elu(0.01);
    let activation = candle_nn::Activation::Sigmoid;
    let activation = candle_nn::Activation::Relu6;
    //...

    let y = activation.forward(&x)?;

更多层见下篇内容。