移植 yolov10 到 candle 不见得有多复杂，主要难度来源于 ultralytics 这个项目本身的代码可读性太差。ultralytics 为了适配 yolov5-v12，还有 yoloe、yolo-world 这些不同的变种，ultralytics代码库中充斥着大量的 if 判断，再加上训练/推理模式混合在一起，使得代码看起来非常混乱。
我个人不是太认可这种代码组织方式，实在是太乱了。牺牲一些抽象换来更好的可读性不好吗？毕竟，模型不同版本之间本身没有什么联系。
不过，在经过一番努力之后，我成功地将 yolov10 移植到了 rust ai 框架 candle 中。
分析
ultralytics 通过配置文件定义模型结构，再通过yaml_model_load(cfg) 加载配置文件，解析模型。这里以 yolov10s 为例，因为 yolov10s 尺寸较小，准确率也能接受。
yolov10s 尺寸的结构定义如下：
backbone:
# [from, repeats, module, args]
- [-1, 1, Conv, [64, 3, 2]] # 0-P1/2
- [-1, 1, Conv, [128, 3, 2]] # 1-P2/4
- [-1, 3, C2f, [128, True]]
- [-1, 1, Conv, [256, 3, 2]] # 3-P3/8
- [-1, 6, C2f, [256, True]]
- [-1, 1, SCDown, [512, 3, 2]] # 5-P4/16
- [-1, 6, C2f, [512, True]]
- [-1, 1, SCDown, [1024, 3, 2]] # 7-P5/32
- [-1, 3, C2fCIB, [1024, True, True]]
- [-1, 1, SPPF, [1024, 5]] # 9
- [-1, 1, PSA, [1024]] # 10
# YOLOv10.0n head
head:
- [-1, 1, nn.Upsample, [None, 2, "nearest"]]
- [[-1, 6], 1, Concat, [1]] # cat backbone P4
- [-1, 3, C2f, [512]] # 13
- [-1, 1, nn.Upsample, [None, 2, "nearest"]]
- [[-1, 4], 1, Concat, [1]] # cat backbone P3
- [-1, 3, C2f, [256]] # 16 (P3/8-small)
- [-1, 1, Conv, [256, 3, 2]]
- [[-1, 13], 1, Concat, [1]] # cat head P4
- [-1, 3, C2f, [512]] # 19 (P4/16-medium)
- [-1, 1, SCDown, [512, 3, 2]]
- [[-1, 10], 1, Concat, [1]] # cat head P5
- [-1, 3, C2fCIB, [1024, True, True]] # 22 (P5/32-large)
- [[16, 19, 22], 1, v10Detect, [nc]] # Detect(P3, P4, P5)这是一个 23 层的模型结构定义。
yolov10 是一个NMS-Free 模型，具体来说，就是训练期间结合了一对多(one2many)和一对一(one2one)的策略，对于训练，走的一对多(one2many)，对于推理，走的一对一(one2one)。因此我们不需要理会(one2many) 的内容。
实现思路
首先， yolov8 有纯 rust 语言实现，代码在这里https://github.com/huggingface/candle。比较遗憾的是，candle 官方项目里的 yolov8 并没有基于ultralytics 这个项目实现，原因大概是许可证问题。因此我们会发现，candle 里的 yolov8 的模型节点名和ultralytics完全对不上，我们也无法直接导出ultralytics的模型再使用 candle 进行推理 (当然，修改节点就能满足要求了)。
但是一些模型我们可以直接复用。因此我们移植工作大概内容如下：
- 实现从
ultralytics导出 safetensors 权重 - 实现 yolov8 中没有的模块,如
SCDown、C2fCIB、PSA、v10Detect等。 - 适配
ultralytics的权重节点名称。 - 实现 v10 版本的后处理。
导出权重
candle 使用的权重格式是 safetensors，因此我们首先需要将ultralytics导出的权重转换成 safetensors。
核心代码如下：
model = YOLO(model_path)
print(model.model)
tensors = model.model.state_dict() # type: ignore
for k, v in tensors.items():
print(str(k), v.shape)
# 保存为safetensors格式
save_model(model.model, output_path) # type: ignore我们需要保存 2 个 print 的输出，一个是权重的形状，一个是权重的 key。主要用于后面适配权重的节点名称。这个很重要。
实现模块
SCDown、C2fCIB 是比较正常的卷积模块，不展开。
PSA 全称Position-Sensitive Attention, 是一个实现位置敏感注意力机制的神经网络模块，用于增强特征提取和处理能力。没深入之前，我以为 yolov10 是没有注意力模块的。移植的时候才发现，我的理解有错误。不过yolov10 也只有这么一个模块，不像 yolov12，注意力机制模块贯穿整个模型。
PSA 最核心的是 Attention 模块。 Attention 多头注意力机制增强表达能力，Rust 实现如下：
#[derive(Clone, Debug)]
pub struct Attention{
qkv: ConvBlock,
proj: ConvBlock,
pe: ConvBlock,
num_heads: usize,
key_dim: usize,
scale: f64,
head_dim: usize,
}
impl Attention {
/// num_heads=8, attn_ratio=0.5
pub fn load(vb:VarBuilder,dim:usize,num_heads:usize,attn_ratio:f64)->Result<Self> {
let head_dim = dim / num_heads;
let key_dim = (head_dim as f64 * attn_ratio) as usize;
let scale = (key_dim as f64).powf(-0.5);
let nh_kd = key_dim * num_heads;
let h = dim + nh_kd * 2;
let qkv = ConvBlock::load(vb.pp("qkv"), dim, h, 1, 1, None, None, false)?;
let proj = ConvBlock::load(vb.pp("proj"), dim, dim, 1, 1, None, None, false)?;
let pe = ConvBlock::load(vb.pp("pe"), dim, dim, 3, 1, None, Some(dim), false)?;
Ok(
Self {
qkv,
proj,
pe,
num_heads,
key_dim,
scale,
head_dim,
}
)
}
}其他模块没有特别的地方。
碰到的实现问题
candle 和 pytorch 终究是不一样的，很多操作 API 存在不一样的地方，之前我自己有记录过一些不一样的地方。这次移植的时候再次碰到了一些问题。
nn.ModuleList
这个模块属于容器，在 candle 中是不存在的。一般我喜欢使用 Vec<Box<dyn Module>> 来代替。
比如 python 版本：
self.m = nn.ModuleList(CIB(self.c, self.c, shortcut, e=1.0, lk=lk) for _ in range(n))Rust 语言版本：
let mut cib = Vec::with_capacity(n);
for idx in 0..n {
// CIB(self.c, self.c, shortcut, e=1.0, lk=lk)
let b = CIB::load(vb.pp(format!("m.{idx}")), c, c, shortcut, 1f64, lk)?;
cib.push(b)
}torch.view
candle 不存在 view 函数，因此需要使用 reshape 函数代替。
split
同样地，candle 也不存在 split 函数。yolov10 中使用得最多的方式是指定尺寸列表分割，比如：
q,k,v = x.split([self.key_dim, self.key_dim, self.head_dim], dim=2)在 rust 中，有类似分割张量的函数我能想到的就是 narrow：
let q = rs.narrow(2, 0, self.key_dim)?; // 从dim=2的第0个位置开始，取key_dim个元素
let k = rs.narrow(2, self.key_dim, self.key_dim)?; // 从dim=2的第key_dim个位置开始，取key_dim个元素
let v = rs.narrow(2, self.key_dim * 2, self.head_dim)?; // 从dim=2的第key_dim*2个位置开始，取head_dim个元素矩阵乘法操作
python 可以实现类似的k @ v 这种矩阵乘法，在 rust 中，需要使用 matmul 函数。
let kv = k.matmul(v)?;有 padding 的 nn.MaxPool2d
有 padding 的 nn.MaxPool2d，比如nn.MaxPool2d(kernel_size=k, stride=1, padding=k // 2),在 rust 上可以这么实现：
x
.pad_with_zeros(2, self.k / 2, self.k / 2)?
.pad_with_zeros(3, self.k / 2, self.k / 2)?
.max_pool2d_with_stride(self.k, 1)?;v10postprocess
yolov10 移植难度最大的其实是 v10postprocess 模块。
首先，candle 没有内容 topk。不过好在 deepseek 的 candle 实现里有一个 topk 函数,比较接近pytorch 的 topk 函数，可以满足需求。
pub struct TopKOutput {
pub values: Tensor,
pub indices: Tensor,
}
pub trait TopKLastDimOp {
fn topk(&self, topk: usize) -> Result<TopKOutput>;
}
impl TopKLastDimOp for Tensor {
fn topk(&self, topk: usize) -> Result<TopKOutput> {
// Sorted descending
let sorted_indices = self.arg_sort_last_dim(false)?;
// 获取最后一维的大小
let last_dim_size = sorted_indices.dim(D::Minus1)?;
// 确保不超过最后一维的实际大小，符合PyTorch的torch.topk行为
let actual_topk = topk.min(last_dim_size);
let topk_indices = sorted_indices
.narrow(D::Minus1, 0, actual_topk)?
.contiguous()?;
Ok(TopKOutput {
values: self.gather(&topk_indices, D::Minus1)?,
indices: topk_indices,
})
}
}其次，candle 中没找到取模运算（求余数）。就是下面这个操作:
result = index % nc # 计算 index ÷ nc 的余数并且这还是一个支持广播机制的取模运算。我的实现如下：
pub trait TensorRemOps {
fn broadcast_rem(&self, other: &Tensor) -> Result<Tensor>;
}
impl TensorRemOps for Tensor {
fn broadcast_rem(&self, other: &Tensor) -> Result<Tensor> {
// 获取广播后的形状
let broadcast_shape = broadcast_shape(self.shape(), other.shape())?;
// 对两个张量进行广播扩展
let self_expanded = self.expand(&broadcast_shape)?;
let other_expanded = other.expand(&broadcast_shape)?;
// 转换为二维数组进行逐元素取模运算
let self_data = self_expanded.to_vec2::<u32>()?;
let other_data = other_expanded.to_vec2::<u32>()?;
// 执行逐元素取模运算
let result: Vec<Vec<u32>> = self_data
.into_iter()
.zip(other_data.into_iter())
.map(|(row1, row2)| {
row1.into_iter().zip(row2.into_iter())
.map(|(a, b)| a % b)
.collect()
})
.collect();
// 将结果展平并重塑为原来的形状
let flat_result: Vec<u32> = result.into_iter().flatten().collect();
Tensor::from_vec(flat_result, &broadcast_shape, self.device())
}
}到这里，结构上的障碍就基本解决完了。
下面是把模块按原版的要求像砌墙一样垒起来。
通过 print(model.model) 观察 yolov10 模型结构，我们会发现，它是一个非常扁平的结构。从 0 到 23层，如果按 rust 的代码组织方式，不太友好。因此，我把 yolov10 的大模块分成了 backbone、neck、head 三大部分。
backbone：0 - 10 层neck： 11 - 22 层head： 23 层
pub struct YoloV10 {
backbone: Backbone,
neck: YoloNeck,
head: V10DetectionHead,
}可以不按原版本的结构来组织代码的原因是，v.pp()支持点语法取值，比如 model.model.0.0.0.weight。
这种组织方式还有利于与 pytorch 版本做权重值对齐。
对齐节点
对齐节点主要靠三个内容：
yolov10s.yaml文件，用于查看整体结构print(model.model)用于对齐节点层级for k, v in tensors.items(): \print(str(k), v.shape)用于对齐节点细节。
写一个调用用例进行测试：
let vb = unsafe {
VarBuilder::from_mmaped_safetensors(
&["yolov10s.safetensors"],
DType::F32,
&device,
)
}?;
let xs = vec![1f32; 640 * 640 * 3];
let image_t = candle_core::Tensor::from_vec(xs, (1, 3, 640, 640), &device)?;
let output = yolo.forward(&image_t)?;如果我们运行上面的代码不再出错，那么恭喜！模型的节点对齐了！
检测头
前面讲过，yolov10 有一对多(one2many)和一对一(one2one) 两个策略。在训练阶段，返回是这样的：
if self.training: # Training path
return {"one2many": x, "one2one": one2one}但是我们只是做推理，推理的时候，处理的其实是 one2one 分支。所以，其实 one2many 的内容是不需要的。
后处理
原来我们使用 onnx 实现 yolov10 的时候，后处理非常简单，那是因为在导出 onnx 的时候，v10postprocess 模块也进入了 onnx 的计算图中。但其实如果算上 v10postprocess 模块的话，个人认为 yolov10 的后处理要比 yolov8 复杂。在上面的实现细节中，就直接碰到了 2 个在 candle 中没有的操作。
在 candle 中，我们当然不得不自己实现 v10postprocess 模块。
fn v10postprocess(preds: &Tensor, max_det: usize, nc: usize) -> Result<Tensor> {
let preds_shape = preds.dims();
assert!(4 + nc == preds_shape[preds_shape.len() - 1]);
// Split boxes and scores
let boxes = preds.i((.., .., ..4))?;
let scores = preds.i((.., .., 4..))?;
let amax_scores = scores.max(D::Minus1)?;
// max_scores, index = torch.topk(max_scores, max_det, dim=-1)
let TopKOutput {
values: _max_scores,
indices: topk_indices,
} = amax_scores.topk(max_det)?;
let index = topk_indices.unsqueeze(D::Minus1)?; // Equivalent to index.unsqueeze(-1)
let boxes = boxes.contiguous()?.gather(&index.repeat((1, 1, 4))?, 1)?;
let scores = scores.contiguous()?.gather(&index.repeat((1, 1, nc))?, 1)?;
// scores, index = torch.topk(scores.flatten(1), max_det, dim=-1)
let scores_flat = scores.flatten(1, 2)?;
let TopKOutput {
values: scores,
indices: index,
} = scores_flat.topk(max_det)?;
let nc_tensor = Tensor::from_slice(&[nc as u32], 1, scores.device())?;
let index_div = index.broadcast_div(&nc_tensor)?;
// 使用 gather 代替高级索引 boxes[i, index // nc]
let boxes_indices = index_div.unsqueeze(2)?.repeat((1, 1, 4))?; // [batch_size, max_det, 4]
let boxes_gathered = boxes.gather(&boxes_indices, 1)?; // [batch_size, max_det, 4]
// scores[..., None] - 添加新轴
let scores_expanded = scores.unsqueeze(2)?; // [batch_size, max_det, 1]
// (index % nc)[..., None].float() - 取模并添加新轴
let index_mod = index.broadcast_rem(&nc_tensor)?; // index % nc
let index_mod_expanded = index_mod.unsqueeze(2)?.to_dtype(DType::F32)?; // [batch_size, max_det, 1]
// 在最后一个维度上连接所有张量
// python: torch.cat([boxes[i, index // nc], scores[..., None], (index % nc)[..., None].float()], dim=-1)
// 最后一维应该是 4 + 1 + 1 = 6
let result = Tensor::cat(&[&boxes_gathered, &scores_expanded, &index_mod_expanded], 2)?;
Ok(result)
}需要注意的是，candle 的 gather 要求输入的必须是连续存储的张量。因此，在 v10postprocess 中，我们需要使用 contiguous() 方法来确保张量在内存中以连续的方式存储。
对齐结果
由于ultralytics 代码比较复杂，使用同一张图片来进行对比结果进行验证其实不太可行。原因是我们很难做到和ultralytics一样的预处理。为了减少预处理带来的误差，我们最好保证输入模型的张量值是一致的。
幸好，ultralytics支持纯张量输入,这样就不会进入预处理步骤。
one = torch.ones(1, 3, 640, 640)
results = model(one)同样地， candle 中可以这样实现：
let xs = vec![1f32; 640 * 640 * 3];
let image_t = candle_core::Tensor::from_vec(xs, (1, 3, 640, 640), &device)?;如何对比两者输出？
这里主要是 ultralytics 的问题，两个办法，第一个办法是使用钩子函数：
# 定义钩子函数
def hook(module, input, output):
print(f"Layer: {module.__class__.__name__}")
print(f"Input shape: {input[0].shape}")
print(f"Output shape: {output.shape}")
print(f"Output sample: {output}")
return output
layer = model.model.model[0] # type: ignore # model.model是实际的网络结构
handle = layer.register_forward_hook(hook)其中model.model.model[0] 中的 0 为模型的第一层。yolo10s 有 24 层[0-23]。
第二个办法是找到ultralytics 执行层的主函数。这个在ultralytics/nn/tasks.py 中可以找到_predict_once。
# 第 180 行
x = m(x) # run
if m.i == 0:
print("x::",x)m.i 也是模型的层数[0-23]。
误差
对齐结果的时候不需要要求值完全一致。事实上，这是完全不可能的事。
误差主要是硬件差异(CPU、GPU)、浮点数精度差异、框架差异、算子实现差异等影响。甚至随着模型的层数加深，误差会越来越大。
比如我目前的这个 rust 版本实现，在经过后处理之后，误差已经来到了 0.1 这个数量级。
下面是 candle 实现的效果,基本和ultralytics一致。
如果你对这个实现感兴趣，可以到这里看源代码：https://gitcode.com/tunzei/yolov10。
当然，我的实现是有一些问题的，比如最后的累计误差有一些大了，主要是有部分算子是按自己的理解实现的，并未真正参考 pytorch的实现。
又比如 ultralytics 有部分模块做了算子融合，这无疑可以加速推理。虽然我没有做速度上的对比，但是应该推理速度上会比 ultralytics 要慢一些。