Rust的Candle中ModernBertSequentialClassification张量适配问题

核心错误分析与修复 #

1. 索引越界错误（index out of bounds: the len is 1 but the index is 1） #

原因：标签张量形状不匹配。你的代码中单个标签是标量（[]形状），而模型计算损失时期望标签为一维张量（[batch_size]），后续批量处理时维度混乱引发索引错误。

修复：创建标签张量时将单个值包装为一维数组，确保形状为[1]（单样本），后续批量堆叠后会自动转为[batch_size]：

// 原代码
// let tensor_label = Tensor::new(label_value, device)?;
// 修改后
let tensor_label = Tensor::new(&[label_value], device)?;

2. 矩阵乘法形状不匹配（shape mismatch in matmul, lhs: [1, 20], rhs: [756, 756]） #

核心原因：手动设置的hidden_size=256与ModernBERT-base预训练模型的原生参数不兼容。该模型的预训练权重中，词嵌入层及Transformer层的隐藏维度均为768，强行修改维度会导致参数加载后形状完全不匹配，引发矩阵乘法错误。

配套修复步骤：

（1）统一序列长度并配置tokenizer

必须将所有样本的输入截断/填充到你设置的seq_len=20，否则样本长度不一致无法堆叠为批量张量：

use tokenizers::{PaddingParams, TruncationParams, TruncationStrategy};

// 初始化tokenizer时配置自动填充与截断
let mut tokenizer = Tokenizer::from_pretrained("answerdotai/ModernBERT-base", None)?;
tokenizer.with_padding(Some(PaddingParams {
    max_length: Some(20),
    pad_id: *tokenizer.get_vocab().get("[PAD]").unwrap() as u32,
    pad_token: "[PAD]".to_string(),
    ..Default::default()
}))?;
tokenizer.with_truncation(Some(TruncationParams {
    max_length: 20,
    strategy: TruncationStrategy::LongestFirst,
    ..Default::default()
}))?;

（2）生成符合要求的输入张量

确保input_ids和attention_mask的形状为[batch_size, seq_len]（单样本时为[1, 20]）：

// 原代码
// let tensor_input_ids = Tensor::new(input_ids, device)?;
// 修改后
let tensor_input_ids = Tensor::new(input_ids, device)?.unsqueeze(0)?; // 新增批量维度

// attention_mask做同样处理
let tensor_mask = Tensor::new(mask, device)?.unsqueeze(0)?;

（3）使用预训练模型的原生配置

禁止手动修改hidden_size，直接加载模型的默认配置：

use candle_transformers::models::bert::{BertConfig, BertWeights};

// 加载预训练模型的原生配置
let config = BertConfig::from_pretrained("answerdotai/ModernBERT-base")?;
// 初始化分类模型（num_labels为你的任务类别数）
let mut model = ModernBertSequentialClassification::new(&config, num_labels, device)?;
// 加载预训练权重
let weights = BertWeights::from_pretrained("answerdotai/ModernBERT-base", device)?;
model.load_weights(&weights)?;

修正后的get_train_data函数示例 #

fn get_train_data(
    tokenizer: &Tokenizer,
    device: &Device,
) -> Result<Vec<(Tensor, Tensor, Tensor)>, Error> {
    let sentences: Vec<&str> = vec![
        "The new smartphone features a foldable display and 5G support.",
        "The government announced new economic policies today.",
        "Regular exercise and a balanced diet are key to staying healthy.",
        "The latest action movie broke box office records this weekend.",
    ];

    let labels: Vec<u32> = vec![1, 2, 3, 4];

    let mut features: Vec<(Tensor, Tensor, Tensor)> = Vec::with_capacity(sentences.len());

    for (idx, text) in sentences.iter().enumerate() {
        let encoding = tokenizer.encode(*text, true)?;

        // 生成[1, 20]形状的input_ids张量
        let input_ids = encoding.get_ids();
        let tensor_input_ids = Tensor::new(input_ids, device)?.unsqueeze(0)?;

        // 生成[1, 20]形状的attention_mask张量
        let mask = encoding.get_attention_mask();
        let tensor_mask = Tensor::new(mask, device)?.unsqueeze(0)?;

        // 生成[1]形状的标签张量
        let label_value = labels[idx];
        let tensor_label = Tensor::new(&[label_value], device)?;

        features.push((tensor_input_ids, tensor_mask, tensor_label));
    }

    Ok(features)
}

关键规则总结 #

• 预训练模型参数不可随意修改：hidden_size、num_attention_heads等核心参数必须与预训练权重完全一致，否则必然引发形状不匹配。
• 张量形状严格对齐：
  • input_ids/attention_mask：[batch_size, seq_len]
  • 标签张量：[batch_size]
  • 模型输出logits：[batch_size, num_labels]
• 批量处理用堆叠而非拼接：合并多个单样本张量时，使用Tensor::stack(&tensors, 0)?生成批量张量，避免维度混乱。

内容的提问来源于stack exchange，提问作者whitebox3