Rust AI 进阶（下） PyTorch 模型转换为 Rust candle 模型教程

限于篇幅太长，因此分为了上下两部分进行讲解。

Sequential

Sequential 在pytorch 是一个容器模型，它本身的功能只是按照添加的顺序将多个模块（层）封装在一起，形成一个线性堆栈。使用 Sequential 可以简化模型定义的过程，特别是当你的网络结构是简单的线性序列时，即每一层的输出直接作为下一层的输入，而没有分支或更复杂的连接方式。

Candle 里没有这个模块，需要我们实现。下面是我自己的一个实现

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

#[derive(Debug, Clone)]
// 定义一个泛型结构体 Sequential，其中 T 必须实现了 Module 特性
pub struct Sequential<T: Module> {
    // 包含多个 T 类型层的向量
    layers: Vec<T>,
}

// 定义一个构造函数 seq，用于创建 Sequential 实例
pub fn seq<T: Module>(cnt: usize) -> Sequential<T> {
    // 根据 cnt 的值初始化一个空向量或具有指定容量的向量
    let v = if cnt == 0 { vec![] } else { Vec::with_capacity(cnt) };
    // 返回一个 Sequential 实例，其 layers 字段初始化为 v
    Sequential { layers: v }
}

// 为 Sequential 结构体实现方法
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);
    }

    // 向序列中添加一个层（与 push 方法功能相同）
    pub fn add(&mut self, layer: T) {
        self.layers.push(layer);
    }
}

// 为 Sequential 结构体实现 Module 特性
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)
    }
}

这是一个通用模块，事实上很多模型都可以通用此模块。

1x1 卷积(2d卷积)

1x1 卷积 对应的是以下结构，这个卷积层有三个输入通道，64个输出通道，卷积核为7x7。

(conv1): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)

对此，Candle 有对应的API：candle_nn::conv2d_no_bias。

/// 创建一个二维卷积层（Conv2d）。
///
/// 该函数使用指定的输入通道数、输出通道数、卷积核大小、填充、步幅和变量构建器来创建一个没有偏置项的二维卷积层。
///
/// # 参数
/// - `in_planes` (usize): 输入图像的通道数。对于RGB图像，通常是3；对于灰度图像，通常是1。
/// - `out_planes` (usize): 输出特征图的通道数。决定了卷积层生成多少个特征图。
/// - `ksize` (usize): 卷积核的大小。例如，7表示7x7的卷积核。
/// - `padding` (usize): 在输入图像边界上添加的零填充量。这有助于保持输出特征图的尺寸与输入相似。
/// - `stride` (usize): 卷积核在输入图像上的移动步幅。较大的步幅会减少输出特征图的尺寸。
/// - `vb` (VarBuilder): 变量构建器，用于初始化卷积层中的权重。`VarBuilder` 通常提供了从预训练模型加载权重或随机初始化的功能。
///
/// # 返回
/// - `Result<Conv2d>`: 如果成功创建了卷积层，则返回一个包含`Conv2d`实例的结果；如果创建过程中发生错误，则返回错误信息。
///
/// # 示例
/// ```rust
/// let vb = VarBuilder::new(); // 假设这是初始化变量构建器的方式
/// let conv_layer = conv2d(3, 64, 7, 3, 2, vb).unwrap();
/// ```
fn conv2d(
    in_planes: usize,
    out_planes: usize,
    ksize: usize,
    padding: usize,
    stride: usize,
    vb: VarBuilder
) -> Result<candle_nn::Conv2d> {
    // 创建一个 Conv2dConfig 实例，并设置步幅和填充。其他配置项使用默认值。
    let conv2d_cfg = candle_nn::Conv2dConfig {
        stride,
        padding,
        ..Default::default()
    };

    // 使用 candle_nn 库提供的 conv2d_no_bias 函数创建一个没有偏置项的二维卷积层。
    // 这个函数接收输入通道数、输出通道数、卷积核大小、配置项和变量构建器作为参数。
    candle_nn::conv2d_no_bias(in_planes, out_planes, ksize, conv2d_cfg, vb)
}

二维批量归一化(BatchNorm2d)

BatchNorm2d 通常放置在卷积层之后和激活函数之前，以改善训练过程中的梯度传播问题，并加速网络的收敛。

(bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)

这个在 candle 中有等同的操作API：

nn:batch_norm(out_planes, 1e-5, vb.pp("bn1"))?;//out_planes=64

relu 激活函数

relu 公式如下： *f*(*x*)=max(0,*x*)

(relu): ReLU(inplace=True)

在candle 中，relu 并不需要去定义结构，只需要在 forward 中调用即可。

一个简单的示例：

impl Module for Block {
    fn forward(&self, xs: &candle_core::Tensor) -> Result<candle_core::Tensor> {
        let ys = xs.relu()?;
        Ok(ys)
    }
}

MaxPool2d 最大池化

最大池化主要是用来减小计算量，增加感受野，并强化特征。

(maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)

同样的，在candle 中，最大池化也不需要定义结构，只需要在 forward 中直接调用即可。

let xs = xs.max_pool2d_with_stride(3, 2)?; // kernel_size=3 stride=2

前面讲过， ResNet18 和ResNet 50 最大的区别是模型网络深度的不同，主要差异在残差块（residual block）的设计和通道数的变化。

残差块设计

ResNet18 用的 BasicBlock 结构，每个残差块包含两个3x3的卷积层，后接BN层和ReLU激活函数。

BasicBlock(
      (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (relu): ReLU(inplace=True)
      (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)

我们使用candle 定义 BasicBlock 结构：

/// 基本残差块结构
/// 包含两个卷积层和两个批量归一化层
/// 如果输入和输出通道数不同或步长不为1，则包含一个下采样层
#[derive(Debug, Clone)]
pub struct BasicBlock {
    /// 第一个卷积层
    conv1: nn::Conv2d,
    /// 第一个卷积层后的批量归一化层
    bn1: nn::BatchNorm,
    /// 第二个卷积层
    conv2: nn::Conv2d,
    /// 第二个卷积层后的批量归一化层
    bn2: nn::BatchNorm,
    /// 可选的下采样层，用于调整特征图的尺寸和通道数
    downsample: Option<Downsample>,
}

Downsample

downsample 这里是可选的，因为有一些 layer 层是有下采样层。这取决于输入和输出通道数以及步长。

/// 对输入特征图进行下采样的函数
/// 当输入输出通道数不一致或stride不为1时，使用1x1卷积调整通道数并进行下采样
/// 
/// # 参数
/// - `in_planes`: 输入特征图的通道数
/// - `out_planes`: 输出特征图的通道数
/// - `stride`: 卷积的步长
/// - `vb`: 用于初始化卷积和批归一化层的变量构建器
/// 
/// # 返回
/// - 如果stride不为1或输入输出通道数不一致，则返回Some(Downsample)实例
/// - 否则，返回None，表示不需要下采样
fn downsample(in_planes: usize, out_planes: usize, stride: usize, vb: VarBuilder) -> Result<Option<Downsample>> {
    // 检查是否需要进行降采样
    if stride != 1 || in_planes != out_planes {
        // 使用1x1卷积调整通道数并进行降采样
        let conv = conv2d(in_planes, out_planes, 1, 0, stride, vb.pp(0))?;
        // 使用批归一化稳定训练
        let bn = batch_norm(out_planes, 1e-5, vb.pp(1))?;
        // 返回下采样模块
        Ok(
            Some(Downsample{ conv2d: conv, bn2: bn, in_planes, out_planes, stride})
        )
    } else {
        // 不需要下采样，返回None
        Ok(None)
    }
}

实例化工厂函数 new 如下：

/// 创建一个新的ResidualBlock实例。
///
/// # 参数
///
/// - `vb`: 用于构建变量的变量构建器。
/// - `in_planes`: 输入通道数。
/// - `out_planes`: 输出通道数。
/// - `stride`: 卷积步长。
///
/// # 返回
///
/// - `Result<Self>`: 返回一个构建好的ResidualBlock实例。
///
/// # 说明
///
/// 该函数负责构建一个残差块，包括两个卷积层、两个批量归一化层和一个下采样层。
/// 每个卷积层后面都跟着一个批量归一化层，以加速网络的训练过程。
/// 下采样层用于匹配输入和输出的维度，以便进行残差相加操作。
pub fn new(vb: VarBuilder, in_planes: usize, out_planes: usize, stride: usize) -> Result<Self> {
    // 构建第一个卷积层，使用指定的输入通道数、输出通道数、卷积核大小、卷积步长等参数
    let conv1 = conv2d(in_planes, out_planes, 3, 1, stride, vb.pp("conv1"))?;

    // 构建第一个批量归一化层，用于对第一个卷积层的输出进行归一化处理
    let bn1 = batch_norm(out_planes, 1e-5, vb.pp("bn1"))?;

    // 构建第二个卷积层，其输入通道数和输出通道数与第一个卷积层相同
    let conv2 = conv2d(out_planes, out_planes, 3, 1, 1, vb.pp("conv2"))?;

    // 构建第二个批量归一化层，用于对第二个卷积层的输出进行归一化处理
    let bn2 = batch_norm(out_planes, 1e-5, vb.pp("bn2"))?;

    // 构建下采样层，用于调整输入的维度以匹配输出，以便进行残差相加操作
    let downsample = downsample(in_planes, out_planes, stride, vb.pp("downsample"))?;

    // 返回构建好的ResidualBlock实例
    Ok(Self { conv1, bn1, conv2, bn2, downsample })
}

没有什么可说的，直接对着模型结构写参数即可。

每个特征我们都必须实现 Module trait, 在这里，我们需要参考 pytorch 里的原始实现：

def forward(self, x: Tensor) -> Tensor:
        identity = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)

        if self.downsample is not None:
            identity = self.downsample(x)

        out += identity
        out = self.relu(out)

        return out

转换成candle 版本：

impl Module for BasicBlock {
	// 定义前向传播函数
	// 参数 xs: 输入张量
	// 返回值: 经过网络层处理后的张量
	fn forward(&self, xs: &candle_core::Tensor) -> Result<candle_core::Tensor> {
	    // 对输入张量依次进行卷积、批归一化、ReLU激活函数、卷积和批归一化操作
	    let ys = xs
	        .apply(&self.conv1)?
	        .apply_t(&self.bn1,false)?
	        .relu()?
	        .apply(&self.conv2)?
	        .apply_t(&self.bn2,false)?;
	
	    // 根据downsample的存在与否，决定是否对输入张量进行下采样，然后与之前的结果相加并应用ReLU激活函数
	    // 这里解释了为什么有一个条件分支：为了处理输入张量维度变化的情况
	    if let Some(downsample) = &self.downsample {
	        (xs.apply(downsample) + ys)?.relu()
	    } else {
	        (ys + xs)?.relu()
	    }
	}
}

AdaptiveAvgPool2d 自适应平均池化

(avgpool): AdaptiveAvgPool2d(output_size=(1, 1))

最麻烦的是这个，这是在candle 里没有对应的api（至少在笔者实现的时候是candle 0.3 版本，那时候是没有的，现在笔者写这篇文章的时候，candle 版本已经来到了 0.8）。

所以我们需要知道， AdaptiveAvgPool2d 在pytorch 里是什么样的一个操作。自适应平均池化的核心思想是：无论输入特征图的尺寸如何，都可以通过自适应地调整池化窗口的大小和步长，确保输出特征图的尺寸符合预设的目标尺寸。具体来说，AdaptiveAvgPool2d 会根据输入特征图的实际尺寸和用户指定的输出尺寸，动态计算每个池化窗口的大小和步长，然后对每个窗口内的元素求平均值，作为该窗口对应的输出值。

本质通过池化窗口的大小和步长求平均。因为刚才池化窗口大小为1，那我们需要在candle里找求平均的相关接口，比如 mean ：

/// 计算输入张量中所有元素的均值。均值计算覆盖所有输入维度，
/// 并根据 `mean_keepdim` 参数决定这些维度是被压缩还是保留。
pub fn mean<D: Dims>(&self, mean_dims: D) -> Result<Self> {
    // 将 mean_dims 转换为索引形式，并检查其有效性
    let mean_dims = mean_dims.to_indexes(self.shape(), "mean")?;
    
    // 计算需要减少的维度的大小
    let reduced_dim: usize = mean_dims.iter().map(|i| self.dims()[*i]).product();
    
    // 计算缩放因子
    let scale = 1f64 / (reduced_dim as f64);
    
    // 计算总和并应用缩放因子以获得均值
    self.sum_impl(mean_dims, false)? * scale
}

我们再看看python 的向前推导，其中有以下调用：

x = self.avgpool(x)
x = torch.flatten(x, 1)

先对张量进行平均池化，再对张量进行展平，刚好等效操作如下：

let xs = xs.mean(D::Minus1)?;
let xs = xs.mean(D::Minus1)?;

Linear 全连接层

全连接层主要是对特征进行整合，并最终输出一个固定大小的向量。这个向量通常用于分类或回归任务的最终决策

(fc): Linear(in_features=512, out_features=1000, bias=True)

这个在candle中有对应的 API。

nn::linear(in_dim, out_dim, vb.pp("fc"));

layer1-layer4

(layer1): Sequential(
    ...
)
 (layer2): Sequential(
	 ...
 )
 (layer3): Sequential(
	 ...
 )
 (layer4): Sequential(
	 ...
 )

这几个层基本是对 BasicBlock 的重复，我们直接封装一个函数

/// 创建基本层，由多个基本块组成
///
/// # 参数
///
/// - `vb`: 用于创建变量的变量构建器
/// - `in_planes`: 输入通道数
/// - `out_planes`: 输出通道数
/// - `stride`: 步幅大小，用于第一个块
/// - `cnt`: 层数，即该层中包含的基本块数量
///
/// # 返回
///
/// 返回一个包含多个基本块的序列
fn basic_layer(
    vb: VarBuilder,
    in_planes: usize,
    out_planes: usize,
    stride: usize,
    cnt: usize
) -> Result<Sequential<BasicBlock>> {
    // 初始化一个序列，用于存储多个基本块
    let mut layers = seq(cnt);
    // 遍历每个基本块的索引，创建基本块并添加到序列中
    for block_index in 0..cnt {
        // 确定输入通道数：第一个块使用输入参数in_planes，其余块使用out_planes
        let l_in = if block_index == 0 { in_planes } else { out_planes };
        // 确定步幅大小：第一个块使用输入参数stride，其余块使用1
        let stride = if block_index == 0 { stride } else { 1 };
        // 创建基本块，这里使用了变量构建器的pp方法来为每个块创建唯一的变量名前缀
        // 使用BasicBlock的new方法创建基本块，并处理可能的错误
        let layer = BasicBlock::new(vb.pp(block_index.to_string()), l_in, out_planes, stride)?;
        // 将创建的基本块添加到序列中
        layers.push(layer);
    }
    // 返回包含多个基本块的序列，表示成功创建了基本层
    Ok(layers)
}

整体结构

resnet18/resnet34 都通用的 Resnet 结构：

/// 定义ResNet网络结构
///
/// # 字段
///
/// - `conv1`: 第一个卷积层，用于初始特征提取
/// - `bn1`: 第一个批量归一化层，用于归一化卷积层的输出
/// - `layer1`: 第一个基本层，由多个基本块组成
/// - `layer2`: 第二个基本层，由多个基本块组成
/// - `layer3`: 第三个基本层，由多个基本块组成
/// - `layer4`: 第四个基本层，由多个基本块组成
/// - `linear`: 可选的全连接层，用于最终分类
#[derive(Debug, Clone)]
pub struct ResNet {
    conv1: Conv2d,
    bn1: nn::BatchNorm,
    layer1: Sequential<BasicBlock>,
    layer2: Sequential<BasicBlock>,
    layer3: Sequential<BasicBlock>,
    layer4: Sequential<BasicBlock>,
    linear: Option<Linear>,
}

创建实例

定义c1-c4 主要是因为不同的resnet 版本的深度不同，resnet 18 和resnet34 差异主要也在这。

/// 创建一个新的ResNet模型实例。
///
/// # 参数
/// - `vb`: 用于构建变量的变量构建器。
/// - `nclasses`: 可选的类别数量，用于配置模型的输出层。
/// - `c1`, `c2`, `c3`, `c4`: 每个残差块中的卷积层数量。
///
/// # 返回
/// 返回一个构建好的ResNet模型实例。
pub fn new(
    vb: VarBuilder,
    nclasses: Option<usize>,
    c1: usize,
    c2: usize,
    c3: usize,
    c4: usize
) -> Result<Self> {
    // 第一个卷积层，用于将输入图像从3通道转换为64通道，使用7x7的卷积核，步长为2。
    let conv1 = conv2d(3, 64, 7, 3, 2, vb.pp("conv1"))?;
    // 第一个批量归一化层，对卷积层的输出进行归一化处理。
    let bn1 = batch_norm(64, 1e-5, vb.pp("bn1"))?;
    // 第一个残差块，包含`c1`个卷积层，输入和输出通道数均为64，步长为1。
    let layer1 = basic_layer(vb.pp("layer1"), 64, 64, 1, c1)?;
    // 第二个残差块，包含`c2`个卷积层，输入通道数为64，输出通道数为128，步长为2。
    let layer2 = basic_layer(vb.pp("layer2"), 64, 128, 2, c2)?;
    // 第三个残差块，包含`c3`个卷积层，输入通道数为128，输出通道数为256，步长为2。
    let layer3 = basic_layer(vb.pp("layer3"), 128, 256, 2, c3)?;
    // 第四个残差块，包含`c4`个卷积层，输入通道数为256，输出通道数为512，步长为2。
    let layer4 = basic_layer(vb.pp("layer4"), 256, 512, 2, c4)?;

    // 根据类别数量决定是否添加全连接层。
    let linear = if let Some(n) = nclasses {
        // 如果提供了类别数量，则添加一个从512通道到类别数量的全连接层。
        Some(nn::linear(512, n, vb.pp("fc"))?)
    } else {
        // 如果未提供类别数量，则不添加全连接层。
        None
    };

    // 构建并返回ResNet模型实例。
    Ok(Self {
        conv1,
        bn1,
        layer1,
        layer2,
        layer3,
        layer4,
        linear,
    })
}

向前传播

python 中，resnet的向前传播如下：

    def _forward_impl(self, x: Tensor) -> Tensor:
        # See note [TorchScript super()]
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)

        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)

        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.fc(x)

        return x

转换成 candle 向前传播：

// 定义网络的前向传播函数
// 参数:
// * `xs`: 输入张量的引用
// 返回值:
// * 返回一个Result，包含前向传播后的输出张量
fn forward(&self, xs: &candle_core::Tensor) -> Result<candle_core::Tensor> {
    // 应用第一个卷积层和第一个批归一化层，随后进行ReLU激活
    let xs = xs.apply(&self.conv1)?;
    let xs = xs.apply_t(&self.bn1, false)?;
    let xs = xs.relu()?;

    // 对输入张量在两个维度上进行填充，然后应用最大池化
    let xs = xs.pad_with_same(D::Minus1, 1, 1)?;
    let xs = xs.pad_with_same(D::Minus2, 1, 1)?;
    let xs = xs.max_pool2d_with_stride(3, 2)?;

    // 应用四个网络层，这些层可能包括卷积层、批归一化层等
    let xs = xs.apply(&self.layer1)?;
    let xs = xs.apply(&self.layer2)?;
    let xs = xs.apply(&self.layer3)?;
    let xs = xs.apply(&self.layer4)?;

    // 执行全局平均池化，相当于将张量的空间维度缩减为1x1
    let xs = xs.mean(D::Minus1)?;
    let xs = xs.mean(D::Minus1)?;

    // 根据最后一个全连接层是否存在，选择应用全连接层或直接返回
    match &self.linear {
        Some(fc) => xs.apply(fc),
        None => Ok(xs),
    }
}

判断全连接层是否存在并非必要操作。这样设计的原因是在实际业务中，有时我们只需要模型的特征提取层，而不需要全连接层，因此添加了一个参数来选择性地移除全连接层。

到这里 resnet18/34 的模型结构实现完了。我们可以再提供简便的函数：

fn resnet(
    vb: VarBuilder,
    nclasses: Option<usize>,
    c1: usize,
    c2: usize,
    c3: usize,
    c4: usize
) -> Result<ResNet> {
    ResNet::new(vb, nclasses, c1, c2, c3, c4)
}

pub fn resnet18(vb: VarBuilder, num_classes: usize) -> Result<ResNet> {
    resnet(vb, Some(num_classes), 2, 2, 2, 2)
}
///去除全连接层
pub fn resnet18_no_final_layer(vb: VarBuilder) -> Result<ResNet> {
    resnet(vb, None, 2, 2, 2, 2)
}

pub fn resnet34(vb: VarBuilder, num_classes: usize) -> Result<ResNet> {
    resnet(vb, Some(num_classes), 3, 4, 6, 3)
}
///去除全连接层
pub fn resnet34_no_final_layer(vb: VarBuilder) -> Result<ResNet> {
    resnet(vb, None, 3, 4, 6, 3)
}

ResNet 50/101/152

这三个版本的残差块设计和前面的 ResNet18/34 是不一样的，并不是 BasicBlock 而是，BottleneckBlock 。

这是 BasicBlock
BasicBlock(
      (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (relu): ReLU(inplace=True)
      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)

这是 BottleneckBlock
Bottleneck(
      (conv1): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)
      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
      (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (relu): ReLU(inplace=True)
      (downsample): Sequential(
        (0): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
        (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      )
 )

因此我们同样需要重新定义 BottleneckBlock 结构：

// 适用于ResNet 50, 101, 和 152的残差块版本。
#[derive(Debug, Clone)]
pub struct BottleneckBlock {
    conv1: Conv2d,       // 第一个卷积层
    bn1: nn::BatchNorm,  // 第一个批归一化层
    conv2: Conv2d,       // 第二个卷积层
    bn2: nn::BatchNorm,  // 第二个批归一化层
    conv3: Conv2d,       // 第三个卷积层
    bn3: nn::BatchNorm,  // 第三个批归一化层
    downsample: Option<Downsample>, // 可选的下采样层
}

其他的基本逻辑基本一致，不再展开。

测试

测试环节至关重要，它能帮助我们发现代码转换过程中的遗漏，也能验证不同操作之间的等效性。

遗憾的是，candle 中没有便捷的方法来打印上述代码实现的模型结构。我们只能依赖 println!("{xs}") 来查看。

要验证模型结构的完整性，我们需要逐行对比代码。

至于验证模型向前推导逻辑的正确性，我采用以下方法：使用相同的模型权重，在可能出现问题的位置（通常是在 forward 函数中）分别打印 xs 变量，然后比较两者的输出是否匹配。

使用”匹配”这个词并不准确，因为不同的模型推理框架和底层实现语言可能会导致浮点数精度的差异。因此，严格来说，Python 版本和 Candle 版本在打印 xs 时会存在一些精度上的差异。

最后的验证

我们以这张图为例子：

写一个测试用例:

// 测试ResNet-18模型的函数
// 该函数加载预训练的ResNet-18模型，并使用它对一张图片进行分类
fn test_resnet18() -> candle_core::Result<()> {
    // 定义模型文件路径
    let model_file = "./testdata/resnet18.safetensors";
    // 定义使用CPU作为计算设备
    let device = candle_core::Device::Cpu;
    // 从模型文件中加载权重参数
    let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[model_file], DType::F32, &device)? };

    // 加载并预处理输入图片
    let image = load_image224("./testdata/mouse.png")?;

    // 构建ResNet-18模型
    let model = resnet18(vb, 1000)?;
    // 在图片数据上增加一个维度，以匹配模型输入的要求
    let image = image.unsqueeze(0)?;

    // 使用模型对图片进行前向传播，得到未经过softmax的预测值（logits）
    let logits = model.forward(&image)?;
    // 对预测值进行softmax操作，转换为概率分布
    let prs = candle_nn::ops::softmax(&logits, D::Minus1)?.i(0)?.to_vec1::<f32>()?;

    // 将概率分布和其对应的类别索引打包，并按概率降序排序
    let mut prs = prs.iter().enumerate().collect::<Vec<_>>();
    prs.sort_by(|(_, p1), (_, p2)| p2.total_cmp(p1));

    // 打印概率最大的前五个类别及其概率
    for &(category_idx, pr) in prs.iter().take(5) {
        println!("{:24}: {:.2}%", CLASSES[category_idx], 100.0 * pr);
    }

    // 函数执行成功，返回Ok
    Ok(())
}

运行测试用例，得到的结果如下，代码完美识别出它的分类是mouse 说明我们的代码已经成功：

running 1 test
test test_resnet18 ... ok

successes:

---- test_resnet18 stdout ----
mouse, computer mouse   : 90.03%
punching bag, punch bag, punching ball, punchball: 4.49%
joystick                : 1.80%
radio, wireless         : 0.44%
vacuum, vacuum cleaner  : 0.19%

总结

将 PyTorch 模型代码转换为 Candle 代码时，最大的挑战在于找不到完全对应的 API。这种情况下，我们需要深入研究 PyTorch 和 Candle 的源代码，寻找等效的实现方式。

在测试阶段也存在困难，因为缺乏完善的测试工具，只能依靠实践经验来验证。这对初学者来说确实不太友好。