尧图网页设计 
Burn 文本分类实战在 AG News 与 DbPedia 上训练 Transformer 模型并用 LoRA 微调【免费下载链接】burnBurn is a next generation tensor library and Deep Learning Framework that doesnt compromise on flexibility, efficiency and portability.项目地址: https://gitcode.com/GitHub_Trending/bu/burn本文以 burn 仓库中的 text-classification 示例为主体讲解如何在 Rust 深度学习框架 Burn 中完成一个完整的文本分类闭环使用 HuggingFace 数据集加载 AG News4 类与 DbPedia14 类语料、用 Transformer Encoder 构建分类模型、在多种计算后端Torch CPU/GPU、WGPU、CUDA、Metal、Flex上训练与推理以及加载model.bpk预训练权重后用 LoRA 做跨类别数迁移的微调。读完后你可以直接复制文中的命令在本仓库中跑通训练、推理和微调三条流水线。一、任务与数据集该示例支持两个经典文本分类基准AG News来自 2000 多个新闻源的新闻标题语料分为 4 个类别World、Sports、Business、Technology。DbPedia从 Wikipedia 抽取的多类别文本分类数据集包含 14 个类别如 Company、EducationalInstitution、Artist 等。从源码看两个数据集均通过 dataset.rs 中的HuggingfaceDatasetLoader拉取并落为 SQLite 存储AG News 对应fancyzhx/ag_newsL66-L72DbPedia 对应fancyzhx/dbpedia_14L141-L147。因此示例依赖 HuggingFacedatasets工具链运行前需确保本机已安装 Python。每个数据集都实现了统一的TextClassificationDatasettraitdataset.rs L20-L23提供两个静态方法num_classes()返回类别总数AG News 为 4、DbPedia 为 14class_name(label)把标签映射回类别名。这个 trait 是后文训练、推理和微调代码泛化的关键——同一条流水线只需替换数据集类型即可。数据加载链路为DatasetSQLite 行级读取→TokenizerBERT 分词→Batcher组批 padding mask→DataLoader下面三节分别展开。二、模型结构Transformer Encoder 首 token 池化模型定义在 model.rs。整体结构是一个双嵌入 Transformer 编码器 线性分类头的浅层分类器#[derive(Module, Debug)] pub struct TextClassificationModel { transformer: TransformerEncoder, embedding_token: Embedding, embedding_pos: Embedding, output: Linear, n_classes: usize, }embedding_token词表大小 ×d_model的词嵌入embedding_pos最大序列长度 ×d_model的位置嵌入。二者相加后除以 2再送入 Transformermodel.rs L81-L86transformerTransformerEncoder通过TransformerEncoderInput::new(embedding).mask_pad(mask_pad)传入 padding mask 以屏蔽填充位L89-L91;outputLinear(d_model → n_classes)分类头。分类策略上模型对编码器输出取序列第一个位置的向量output.slice([0..batch_size, 0..1])model.rs L94-L96作为整句表示这类似于 BERT 的[CLS]池化思路。由于采用固定索引的位置嵌入策略配置要求必须给定最大序列长度SeqLengthOption::NoMax会在初始化时直接 panicmodel.rs L46-L51。训练前向由TrainSteptrait 驱动step中执行forward后调用item.loss.backward()得到梯度再包装为TrainOutput返回model.rs L151-L162。损失函数是CrossEntropyLossConfig验证侧则实现InferenceStep只跑前向。推理专用方法infer与forward流程一致只是最后对分类 logits 做softmax返回概率分布model.rs L111-L139。此外模型还有一个专为微调准备的方法reset_headmodel.rs L141-L147重建一个指向新类别数的线性头并把其weight与bias显式置为require_grad(true)让新头从随机权重开始学习——这一点是 LoRA 微调跨类别数迁移的前提第四节会用到。三、训练流水线配置、数据加载与指标训练入口是 training.rs 的泛型函数trainD: TextClassificationDataset核心由实验配置结构体控制#[derive(Config, Debug)] pub struct ExperimentConfig { pub transformer: TransformerEncoderConfig, pub optimizer: AdamConfig, #[config(default SeqLengthOption::Fixed(256))] pub seq_length: SeqLengthOption, #[config(default 32)] pub batch_size: usize, #[config(default 5)] pub num_epochs: usize, }默认超参training.rs L27-L37为序列长度 256、batch size 32、训练 5 个 epoch。各示例中实际使用的 Transformer 与优化器配置不同AG News 训练ag-news-train.rs L60-L74TransformerEncoderConfig::new(128, 512, 4, 4)较小模型d_model128、ffn512、4 层 4 头并启用with_norm_first(true)、with_quiet_softmax(true)优化器为AdamConfig::new().with_weight_decay(Some(WeightDecayConfig::new(5e-5)))。DbPedia 训练db-pedia-train.rs L16-L29TransformerEncoderConfig::new(256, 1024, 8, 4)d_model256、ffn1024、8 头 4 层同样配置 Adam weight decay 5e-5。数据侧train函数内部完成这些组装training.rs L44-L77创建BertCasedTokenizer默认分词器下文第四节用TextClassificationBatcher::new(tokenizer, seq_length)构造组批器在strategy.main_device().autodiff()设备上下文上初始化模型用DataLoaderBuilder分别构建训练/测试数据加载器batch_size(config.batch_size)、num_workers(1)并分别用SamplerDataset::new(dataset, 25_000)训练与SamplerDataset::new(dataset, 2500)测试控制每轮采样的最大样本数避免整轮跑完超大数据集。优化与学习率调度training.rs L79-L87let optim config.optimizer.init(); let lr_scheduler NoamLrSchedulerConfig::new(1e-2) .with_warmup_steps(1000) .with_model_size(config.transformer.d_model) .init() .unwrap();即经典的 Noam 调度器峰值学习率 1e-21000 步 warmup学习率曲线随d_model缩放。训练循环由SupervisedTraining构建挂载了丰富的指标training.rs L89-L102CudaMetric训练/验证、IterationSpeedMetric训练、LossMetric训练/验证数值、AccuracyMetric训练/验证数值、LearningRateMetric训练并启用with_default_checkpointers()做检查点。最终通过training.launch(Learner::new(model, optim, lr_scheduler))启动。训练结束后产物统一写入 artifact 目录AG News 为/tmp/text-classification-ag-newsDbPedia 为/tmp/text-classification-db-pediaconfig.save(format!({artifact_dir}/config.json)).unwrap(); result.model.into_record().save(format!({artifact_dir}/model)).unwrap();即同时保存实验配置config.json与模型权重model.bpkburn 的Record::save会在路径后补.bpk扩展名。这两个文件正是推理与微调阶段重新加载模型的输入。分词与组批细节分词器tokenizer.rsBertCasedTokenizer基于 Rust 的tokenizers库加载bert-base-cased预训练分词模型tokenizer.rs L37-L43。Tokenizertrait 约定了encode/decode/vocab_size/pad_token四个接口且要求Send Sync以支持跨线程使用。词表大小在模型初始化时通过tokenizer.vocab_size()注入。组批器batcher.rsTextClassificationBatcher分别实现了两个Batcher——训练批TextClassificationTrainingBatch { tokens, labels, mask_pad }与推理批TextClassificationInferenceBatch { tokens, mask_pad }。核心是把变长 token 序列通过generate_padding_mask(pad_token, tokens_list, seq_length, device)统一 padding 到seq_length并生成布尔掩码batcher.rs L65-L78该掩码随后注入 Transformer 的注意力计算。四、跨后端运行训练与推理示例通过 Cargo feature 切换计算后端Cargo.toml 中定义了tch-cpu、tch-gpu、wgpu、vulkan、flex、remote、cuda、rocm、metal等 feature分别映射到burn/tch、burn/wgpu、burn/flex等依赖另有f16半精度、flex32与ddp多卡分布式标记位。以 ag-news-train.rs 为例main函数按启用的 feature 分发到对应设备模块tch-gpu走Device::libtorch_cudamacOS 上自动落到libtorch_mps、tch-cpu走Device::libtorch()、wgpu/metal/vulkan走Device::wgpu(DefaultDevice)、cuda默认枚举所有 CUDA 卡并以MultiDeviceOptimSharded多设备训练启用ddpfeature 时切换为ExecutionStrategy::ddpag-news-train.rs L24-L50、flex走Device::flex()。数值类型也由 feature 决定默认f32开启f16时为f16flex32时为flex32L16-L22。各后端的完整运行命令如下均在仓库根目录执行务必加--release加速训练Torch GPU 后端export TORCH_CUDA_VERSIONcu128 # 设置 CUDA 版本CUDA 用户 # AG News cargo run --example ag-news-train --release --features tch-gpu # 训练 cargo run --example ag-news-infer --release --features tch-gpu # 推理 # DbPedia cargo run --example db-pedia-train --release --features tch-gpu # 训练 cargo run --example db-pedia-infer --release --features tch-gpu # 推理若 CUDA 设备支持 FP16可追加f16feature 以半精度运行并非所有设备都表现良好。Torch CPU 后端cargo run --example ag-news-train --release --features tch-cpu cargo run --example ag-news-infer --release --features tch-cpu cargo run --example db-pedia-train --release --features tch-cpu cargo run --example db-pedia-infer --release --features tch-cpuFlex 后端cargo run --example ag-news-train --release --features flex cargo run --example ag-news-infer --release --features flex cargo run --example db-pedia-train --release --features flex cargo run --example db-pedia-infer --release --features flexWGPU 后端cargo run --example ag-news-train --release --features wgpu cargo run --example ag-news-infer --release --features wgpu cargo run --example db-pedia-train --release --features wgpu cargo run --example db-pedia-infer --release --features wgpuCUDA 后端# 可追加 f16 feature 以半精度运行 cargo run --example ag-news-train --release --features cuda cargo run --example ag-news-infer --release --features cudaMetal 后端# 可追加 f16 feature 以半精度运行 cargo run --example ag-news-train --release --features metal cargo run --example ag-news-infer --release --features metal五、推理加载权重并对样本文本分类推理入口是 inference.rs 的泛型函数inferD: TextClassificationDataset流程inference.rs L16-L71用ExperimentConfig::load({artifact_dir}/config.json)还原训练时的实验配置重建BertCasedTokenizer与TextClassificationBatcher复用训练时的seq_length保证 token 长度一致通过ModuleRecord::load({artifact_dir}/model)加载权重再用TextClassificationModelConfig::new(config.transformer, n_classes, tokenizer.vocab_size(), config.seq_length).init(device).load_record(record)实例化模型并灌入权重将样本文本经 batcher 组批后调用model.infer逐条取出 softmax 概率、用argmax得到类别索引并通过D::class_name打印类别名。以 ag-news-infer.rs 为例它内置了三条样本文本分别对应体育、国际、科技新闻并从/tmp/text-classification-ag-news加载产物ag-news-infer.rs L17-L36text_classification::inference::infer::AgNewsDataset( device, /tmp/text-classification-ag-news, vec![ Jays power up to take finale ....to_string(), Yemen Sentences 15 Militants on Terror Charges ....to_string(), IBM puts grids to work at U.S. Open ....to_string(), ], );因此推理必须先完成对应数据集的训练使config.json与model.bpk存在于同一 artifact 目录。六、LoRA 微调用 DbPedia 预训练权重迁移到 AG NewsREADME 中的 Finetuning Using LoRA 一节对应的实现是 finetune.rs 的lora_finetuning函数把在 DbPedia14 类上预训练的模型加载进来用 LoRA低秩适配Low-Rank Adaptation只训练低秩增量矩阵并把分类头换成 AG News 的 4 类头实现在参数开销远低于全量微调的前提下完成跨数据集迁移。关键源码步骤finetune.rs L20-L52// 1. 按“预训练时的类别数”初始化模型结构 let model TextClassificationModelConfig::new( config.transformer.clone(), num_class_before, // 传入 14DbPedia 的类别数 tokenizer.vocab_size(), config.seq_length, ) .init(strategy.main_device().clone().autodiff()); // 2. 从磁盘加载预训练权重 let model model.load_file(model.bpk); // 3. 对注意力模块的 query/value/output 与 feed-forward 权重施加 LoRA let r 8.0; let mut model model.apply_lora(Lora::new(r as usize, 2.0 * r)); // rank8, alpha16 // 4. 用当前数据集的类别数重置分类头随机初始化、可训练 model.reset_head(D::num_classes());要点说明权重文件要求model.load_file(model.bpk)是相对路径运行前需要在当前工作目录放置一个model.bpk即先跑完db-pedia-train之一并把产物/tmp/text-classification-db-pedia/model.bpk复制过来——这正是 README 中Such a file can be obtained by running thedb-pedia-trainexample的由来结构匹配因为权重来自 DbPedia 训练ag-news-finetune.rs里初始化的 Transformer 必须与预训练时一致TransformerEncoderConfig::new(256, 1024, 8, 4)与 db-pedia-train.rs L18 相同并把num_class_before设为 14ag-news-finetune.rs L60-L73LoRA 参数Lora::new(8, 16)表示秩 r8、缩放系数 alpha2r16apply_lora是 burnModuletrait 提供的方法从源码注释看它作用于注意力 Q/V/O 与 feed-forward 权重这些低秩分支成为主要可训练参数主干权重保持冻结加载优化配置差异微调不复用 Noam 调度器而是直接用标量学习率1e-3构造Learner::new(model, optim, lr_scheduler)finetune.rs L67-L68其余数据加载训练采样 25_000、测试采样 2500、指标挂载与产物保存逻辑与第三节的全量训练保持一致。运行方式与训练相同任选上面列出的后端 feature例如用 WGPUcargo run --example ag-news-finetune --release --features wgpu微调产物同样输出到/tmp/text-classification-ag-newsconfig.json model.bpk随后可用ag-news-infer推理验证。七、性能相关配置burn.toml示例目录下的 burn.toml 提供了运行期性能调优配置使用 CubeCl 系后端WGPU/CUDA/Metal 等时生效[fusion] logger { file /tmp/fusion.log, level full } [cubecl.autotune] level balanced cache target [cubecl.memory] persistent_memory enabled [cubecl.streaming] max_streams 8其中fusion打开算子融合日志输出到/tmp/fusion.logcubecl.autotune的level balanced并缓存到 target 目录persistent_memory enabled复用持久内存池max_streams 8限制并发流数量。若需要观察融合行为或排查性能问题可直接查看对应日志文件。八、小结一条泛型流水线覆盖三件事train/infer/lora_finetuning都以D: TextClassificationDataset泛型化AG News 与 DbPedia 只是 trait 的两个实现切换数据集无需改动训练逻辑产物闭环训练输出config.jsonmodel.bpk→ 推理按配置重建模型并加载 → 微调再加载model.bpk并施加 LoRA 与reset_head多后端等价同一份 Rust 代码通过 Cargo feature 编译进不同后端--release是训练速度的前提f16提供半精度选项可深入阅读的源码入口model.rs模型结构与前向、training.rs实验配置与训练装配、finetune.rsLoRA 微调、batcher.rs 与 tokenizer.rs数据管线。【免费下载链接】burnBurn is a next generation tensor library and Deep Learning Framework that doesnt compromise on flexibility, efficiency and portability.项目地址: https://gitcode.com/GitHub_Trending/bu/burn创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 Burn 文本分类实战：在 AG News 与 DbPedia 上训练 Transformer 模型并用 LoRA 微调
Burn 文本分类实战在 AG News 与 DbPedia 上训练 Transformer 模型并用 LoRA 微调【免费下载链接】burnBurn is a next generation tensor library and Deep Learning Framework that doesnt compromise on flexibility, efficiency and portability.项目地址: https://gitcode.com/GitHub_Trending/bu/burn本文以 burn 仓库中的 text-classification 示例为主体讲解如何在 Rust 深度学习框架 Burn 中完成一个完整的文本分类闭环使用 HuggingFace 数据集加载 AG News4 类与 DbPedia14 类语料、用 Transformer Encoder 构建分类模型、在多种计算后端Torch CPU/GPU、WGPU、CUDA、Metal、Flex上训练与推理以及加载model.bpk预训练权重后用 LoRA 做跨类别数迁移的微调。读完后你可以直接复制文中的命令在本仓库中跑通训练、推理和微调三条流水线。一、任务与数据集该示例支持两个经典文本分类基准AG News来自 2000 多个新闻源的新闻标题语料分为 4 个类别World、Sports、Business、Technology。DbPedia从 Wikipedia 抽取的多类别文本分类数据集包含 14 个类别如 Company、EducationalInstitution、Artist 等。从源码看两个数据集均通过 dataset.rs 中的HuggingfaceDatasetLoader拉取并落为 SQLite 存储AG News 对应fancyzhx/ag_newsL66-L72DbPedia 对应fancyzhx/dbpedia_14L141-L147。因此示例依赖 HuggingFacedatasets工具链运行前需确保本机已安装 Python。每个数据集都实现了统一的TextClassificationDatasettraitdataset.rs L20-L23提供两个静态方法num_classes()返回类别总数AG News 为 4、DbPedia 为 14class_name(label)把标签映射回类别名。这个 trait 是后文训练、推理和微调代码泛化的关键——同一条流水线只需替换数据集类型即可。数据加载链路为DatasetSQLite 行级读取→TokenizerBERT 分词→Batcher组批 padding mask→DataLoader下面三节分别展开。二、模型结构Transformer Encoder 首 token 池化模型定义在 model.rs。整体结构是一个双嵌入 Transformer 编码器 线性分类头的浅层分类器#[derive(Module, Debug)] pub struct TextClassificationModel { transformer: TransformerEncoder, embedding_token: Embedding, embedding_pos: Embedding, output: Linear, n_classes: usize, }embedding_token词表大小 ×d_model的词嵌入embedding_pos最大序列长度 ×d_model的位置嵌入。二者相加后除以 2再送入 Transformermodel.rs L81-L86transformerTransformerEncoder通过TransformerEncoderInput::new(embedding).mask_pad(mask_pad)传入 padding mask 以屏蔽填充位L89-L91;outputLinear(d_model → n_classes)分类头。分类策略上模型对编码器输出取序列第一个位置的向量output.slice([0..batch_size, 0..1])model.rs L94-L96作为整句表示这类似于 BERT 的[CLS]池化思路。由于采用固定索引的位置嵌入策略配置要求必须给定最大序列长度SeqLengthOption::NoMax会在初始化时直接 panicmodel.rs L46-L51。训练前向由TrainSteptrait 驱动step中执行forward后调用item.loss.backward()得到梯度再包装为TrainOutput返回model.rs L151-L162。损失函数是CrossEntropyLossConfig验证侧则实现InferenceStep只跑前向。推理专用方法infer与forward流程一致只是最后对分类 logits 做softmax返回概率分布model.rs L111-L139。此外模型还有一个专为微调准备的方法reset_headmodel.rs L141-L147重建一个指向新类别数的线性头并把其weight与bias显式置为require_grad(true)让新头从随机权重开始学习——这一点是 LoRA 微调跨类别数迁移的前提第四节会用到。三、训练流水线配置、数据加载与指标训练入口是 training.rs 的泛型函数trainD: TextClassificationDataset核心由实验配置结构体控制#[derive(Config, Debug)] pub struct ExperimentConfig { pub transformer: TransformerEncoderConfig, pub optimizer: AdamConfig, #[config(default SeqLengthOption::Fixed(256))] pub seq_length: SeqLengthOption, #[config(default 32)] pub batch_size: usize, #[config(default 5)] pub num_epochs: usize, }默认超参training.rs L27-L37为序列长度 256、batch size 32、训练 5 个 epoch。各示例中实际使用的 Transformer 与优化器配置不同AG News 训练ag-news-train.rs L60-L74TransformerEncoderConfig::new(128, 512, 4, 4)较小模型d_model128、ffn512、4 层 4 头并启用with_norm_first(true)、with_quiet_softmax(true)优化器为AdamConfig::new().with_weight_decay(Some(WeightDecayConfig::new(5e-5)))。DbPedia 训练db-pedia-train.rs L16-L29TransformerEncoderConfig::new(256, 1024, 8, 4)d_model256、ffn1024、8 头 4 层同样配置 Adam weight decay 5e-5。数据侧train函数内部完成这些组装training.rs L44-L77创建BertCasedTokenizer默认分词器下文第四节用TextClassificationBatcher::new(tokenizer, seq_length)构造组批器在strategy.main_device().autodiff()设备上下文上初始化模型用DataLoaderBuilder分别构建训练/测试数据加载器batch_size(config.batch_size)、num_workers(1)并分别用SamplerDataset::new(dataset, 25_000)训练与SamplerDataset::new(dataset, 2500)测试控制每轮采样的最大样本数避免整轮跑完超大数据集。优化与学习率调度training.rs L79-L87let optim config.optimizer.init(); let lr_scheduler NoamLrSchedulerConfig::new(1e-2) .with_warmup_steps(1000) .with_model_size(config.transformer.d_model) .init() .unwrap();即经典的 Noam 调度器峰值学习率 1e-21000 步 warmup学习率曲线随d_model缩放。训练循环由SupervisedTraining构建挂载了丰富的指标training.rs L89-L102CudaMetric训练/验证、IterationSpeedMetric训练、LossMetric训练/验证数值、AccuracyMetric训练/验证数值、LearningRateMetric训练并启用with_default_checkpointers()做检查点。最终通过training.launch(Learner::new(model, optim, lr_scheduler))启动。训练结束后产物统一写入 artifact 目录AG News 为/tmp/text-classification-ag-newsDbPedia 为/tmp/text-classification-db-pediaconfig.save(format!({artifact_dir}/config.json)).unwrap(); result.model.into_record().save(format!({artifact_dir}/model)).unwrap();即同时保存实验配置config.json与模型权重model.bpkburn 的Record::save会在路径后补.bpk扩展名。这两个文件正是推理与微调阶段重新加载模型的输入。分词与组批细节分词器tokenizer.rsBertCasedTokenizer基于 Rust 的tokenizers库加载bert-base-cased预训练分词模型tokenizer.rs L37-L43。Tokenizertrait 约定了encode/decode/vocab_size/pad_token四个接口且要求Send Sync以支持跨线程使用。词表大小在模型初始化时通过tokenizer.vocab_size()注入。组批器batcher.rsTextClassificationBatcher分别实现了两个Batcher——训练批TextClassificationTrainingBatch { tokens, labels, mask_pad }与推理批TextClassificationInferenceBatch { tokens, mask_pad }。核心是把变长 token 序列通过generate_padding_mask(pad_token, tokens_list, seq_length, device)统一 padding 到seq_length并生成布尔掩码batcher.rs L65-L78该掩码随后注入 Transformer 的注意力计算。四、跨后端运行训练与推理示例通过 Cargo feature 切换计算后端Cargo.toml 中定义了tch-cpu、tch-gpu、wgpu、vulkan、flex、remote、cuda、rocm、metal等 feature分别映射到burn/tch、burn/wgpu、burn/flex等依赖另有f16半精度、flex32与ddp多卡分布式标记位。以 ag-news-train.rs 为例main函数按启用的 feature 分发到对应设备模块tch-gpu走Device::libtorch_cudamacOS 上自动落到libtorch_mps、tch-cpu走Device::libtorch()、wgpu/metal/vulkan走Device::wgpu(DefaultDevice)、cuda默认枚举所有 CUDA 卡并以MultiDeviceOptimSharded多设备训练启用ddpfeature 时切换为ExecutionStrategy::ddpag-news-train.rs L24-L50、flex走Device::flex()。数值类型也由 feature 决定默认f32开启f16时为f16flex32时为flex32L16-L22。各后端的完整运行命令如下均在仓库根目录执行务必加--release加速训练Torch GPU 后端export TORCH_CUDA_VERSIONcu128 # 设置 CUDA 版本CUDA 用户 # AG News cargo run --example ag-news-train --release --features tch-gpu # 训练 cargo run --example ag-news-infer --release --features tch-gpu # 推理 # DbPedia cargo run --example db-pedia-train --release --features tch-gpu # 训练 cargo run --example db-pedia-infer --release --features tch-gpu # 推理若 CUDA 设备支持 FP16可追加f16feature 以半精度运行并非所有设备都表现良好。Torch CPU 后端cargo run --example ag-news-train --release --features tch-cpu cargo run --example ag-news-infer --release --features tch-cpu cargo run --example db-pedia-train --release --features tch-cpu cargo run --example db-pedia-infer --release --features tch-cpuFlex 后端cargo run --example ag-news-train --release --features flex cargo run --example ag-news-infer --release --features flex cargo run --example db-pedia-train --release --features flex cargo run --example db-pedia-infer --release --features flexWGPU 后端cargo run --example ag-news-train --release --features wgpu cargo run --example ag-news-infer --release --features wgpu cargo run --example db-pedia-train --release --features wgpu cargo run --example db-pedia-infer --release --features wgpuCUDA 后端# 可追加 f16 feature 以半精度运行 cargo run --example ag-news-train --release --features cuda cargo run --example ag-news-infer --release --features cudaMetal 后端# 可追加 f16 feature 以半精度运行 cargo run --example ag-news-train --release --features metal cargo run --example ag-news-infer --release --features metal五、推理加载权重并对样本文本分类推理入口是 inference.rs 的泛型函数inferD: TextClassificationDataset流程inference.rs L16-L71用ExperimentConfig::load({artifact_dir}/config.json)还原训练时的实验配置重建BertCasedTokenizer与TextClassificationBatcher复用训练时的seq_length保证 token 长度一致通过ModuleRecord::load({artifact_dir}/model)加载权重再用TextClassificationModelConfig::new(config.transformer, n_classes, tokenizer.vocab_size(), config.seq_length).init(device).load_record(record)实例化模型并灌入权重将样本文本经 batcher 组批后调用model.infer逐条取出 softmax 概率、用argmax得到类别索引并通过D::class_name打印类别名。以 ag-news-infer.rs 为例它内置了三条样本文本分别对应体育、国际、科技新闻并从/tmp/text-classification-ag-news加载产物ag-news-infer.rs L17-L36text_classification::inference::infer::AgNewsDataset( device, /tmp/text-classification-ag-news, vec![ Jays power up to take finale ....to_string(), Yemen Sentences 15 Militants on Terror Charges ....to_string(), IBM puts grids to work at U.S. Open ....to_string(), ], );因此推理必须先完成对应数据集的训练使config.json与model.bpk存在于同一 artifact 目录。六、LoRA 微调用 DbPedia 预训练权重迁移到 AG NewsREADME 中的 Finetuning Using LoRA 一节对应的实现是 finetune.rs 的lora_finetuning函数把在 DbPedia14 类上预训练的模型加载进来用 LoRA低秩适配Low-Rank Adaptation只训练低秩增量矩阵并把分类头换成 AG News 的 4 类头实现在参数开销远低于全量微调的前提下完成跨数据集迁移。关键源码步骤finetune.rs L20-L52// 1. 按“预训练时的类别数”初始化模型结构 let model TextClassificationModelConfig::new( config.transformer.clone(), num_class_before, // 传入 14DbPedia 的类别数 tokenizer.vocab_size(), config.seq_length, ) .init(strategy.main_device().clone().autodiff()); // 2. 从磁盘加载预训练权重 let model model.load_file(model.bpk); // 3. 对注意力模块的 query/value/output 与 feed-forward 权重施加 LoRA let r 8.0; let mut model model.apply_lora(Lora::new(r as usize, 2.0 * r)); // rank8, alpha16 // 4. 用当前数据集的类别数重置分类头随机初始化、可训练 model.reset_head(D::num_classes());要点说明权重文件要求model.load_file(model.bpk)是相对路径运行前需要在当前工作目录放置一个model.bpk即先跑完db-pedia-train之一并把产物/tmp/text-classification-db-pedia/model.bpk复制过来——这正是 README 中Such a file can be obtained by running thedb-pedia-trainexample的由来结构匹配因为权重来自 DbPedia 训练ag-news-finetune.rs里初始化的 Transformer 必须与预训练时一致TransformerEncoderConfig::new(256, 1024, 8, 4)与 db-pedia-train.rs L18 相同并把num_class_before设为 14ag-news-finetune.rs L60-L73LoRA 参数Lora::new(8, 16)表示秩 r8、缩放系数 alpha2r16apply_lora是 burnModuletrait 提供的方法从源码注释看它作用于注意力 Q/V/O 与 feed-forward 权重这些低秩分支成为主要可训练参数主干权重保持冻结加载优化配置差异微调不复用 Noam 调度器而是直接用标量学习率1e-3构造Learner::new(model, optim, lr_scheduler)finetune.rs L67-L68其余数据加载训练采样 25_000、测试采样 2500、指标挂载与产物保存逻辑与第三节的全量训练保持一致。运行方式与训练相同任选上面列出的后端 feature例如用 WGPUcargo run --example ag-news-finetune --release --features wgpu微调产物同样输出到/tmp/text-classification-ag-newsconfig.json model.bpk随后可用ag-news-infer推理验证。七、性能相关配置burn.toml示例目录下的 burn.toml 提供了运行期性能调优配置使用 CubeCl 系后端WGPU/CUDA/Metal 等时生效[fusion] logger { file /tmp/fusion.log, level full } [cubecl.autotune] level balanced cache target [cubecl.memory] persistent_memory enabled [cubecl.streaming] max_streams 8其中fusion打开算子融合日志输出到/tmp/fusion.logcubecl.autotune的level balanced并缓存到 target 目录persistent_memory enabled复用持久内存池max_streams 8限制并发流数量。若需要观察融合行为或排查性能问题可直接查看对应日志文件。八、小结一条泛型流水线覆盖三件事train/infer/lora_finetuning都以D: TextClassificationDataset泛型化AG News 与 DbPedia 只是 trait 的两个实现切换数据集无需改动训练逻辑产物闭环训练输出config.jsonmodel.bpk→ 推理按配置重建模型并加载 → 微调再加载model.bpk并施加 LoRA 与reset_head多后端等价同一份 Rust 代码通过 Cargo feature 编译进不同后端--release是训练速度的前提f16提供半精度选项可深入阅读的源码入口model.rs模型结构与前向、training.rs实验配置与训练装配、finetune.rsLoRA 微调、batcher.rs 与 tokenizer.rs数据管线。【免费下载链接】burnBurn is a next generation tensor library and Deep Learning Framework that doesnt compromise on flexibility, efficiency and portability.项目地址: https://gitcode.com/GitHub_Trending/bu/burn创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考