Rust 机器学习领域约束与实现：高效内存、GPU加速与模型部署最佳实践

domain-ml by zhanghandong/rust-skills

437 周安装量

941 GitHub Stars

GitHub

安装命令

npx skills add https://github.com/zhanghandong/rust-skills --skill domain-ml

AI/机器学习性能优化 Rust

🇨🇳中文介绍

机器学习领域

第 3 层：领域约束

领域约束 → 设计影响

领域规则	设计约束	Rust 实现含义
大数据量	高效内存使用	零拷贝、流式处理
GPU 加速	支持 CUDA/Metal	candle, tch-rs
模型可移植性	标准格式	ONNX
批处理	吞吐量优先于延迟	批量推理
数值精度	浮点数处理	ndarray, 谨慎使用 f32/f64
可复现性	确定性	设定种子的随机数、版本控制

关键约束

内存效率

RULE: Avoid copying large tensors
WHY: Memory bandwidth is bottleneck
RUST: References, views, in-place ops

广告位招租

在这里展示您的产品或服务

触达数万 AI 开发者，精准高效

联系我们

用例	推荐	原因
仅推理	tract (ONNX)	轻量级、可移植
训练 + 推理	candle, burn	纯 Rust，支持 GPU
PyTorch 模型	tch-rs	直接绑定
数据管道	polars	快速、惰性求值

用途	库
张量	ndarray
ONNX 推理	tract
机器学习框架	candle, burn
PyTorch 绑定	tch-rs
数据处理	polars
嵌入	fastembed

模式	目的	实现
模型加载	一次加载，重复使用	`OnceLock<Model>`
批处理	提高吞吐量	收集然后处理
流式处理	处理大数据	基于迭代器
GPU 异步	并行化	数据加载与计算并行

代码模式：推理服务器

use std::sync::OnceLock;
use tract_onnx::prelude::*;

static MODEL: OnceLock<SimplePlan<TypedFact, Box<dyn TypedOp>, Graph<TypedFact, Box<dyn TypedOp>>>> = OnceLock::new();

fn get_model() -> &'static SimplePlan<...> {
    MODEL.get_or_init(|| {
        tract_onnx::onnx()
            .model_for_path("model.onnx")
            .unwrap()
            .into_optimized()
            .unwrap()
            .into_runnable()
            .unwrap()
    })
}

async fn predict(input: Vec<f32>) -> anyhow::Result<Vec<f32>> {
    let model = get_model();
    let input = tract_ndarray::arr1(&input).into_shape((1, input.len()))?;
    let result = model.run(tvec!(input.into()))?;
    Ok(result[0].to_array_view::<f32>()?.iter().copied().collect())
}

代码模式：批量推理

async fn batch_predict(inputs: Vec<Vec<f32>>, batch_size: usize) -> Vec<Vec<f32>> {
    let mut results = Vec::with_capacity(inputs.len());

    for batch in inputs.chunks(batch_size) {
        // Stack inputs into batch tensor
        let batch_tensor = stack_inputs(batch);

        // Run inference on batch
        let batch_output = model.run(batch_tensor).await;

        // Unstack results
        results.extend(unstack_outputs(batch_output));
    }

    results
}

错误	违反的领域原则	修复方法
克隆张量	浪费内存	使用视图
单次推理	GPU 利用率低	批处理
每个请求都加载模型	速度慢	单例模式
同步数据加载	GPU 空闲	异步管道

约束	第 2 层模式	第 1 层实现
内存效率	零拷贝	ndarray 视图
模型单例	惰性初始化	OnceLock
批处理	分块迭代	chunks() + parallel
GPU 异步	并发加载	tokio::spawn + GPU

何时使用	参见
性能	m10-performance
惰性初始化	m12-lifecycle
异步模式	m07-concurrency
内存效率	m01-ownership

🇺🇸English

Machine Learning Domain

Layer 3: Domain Constraints

Domain Constraints → Design Implications

Domain Rule	Design Constraint	Rust Implication
Large data	Efficient memory	Zero-copy, streaming
GPU acceleration	CUDA/Metal support	candle, tch-rs
Model portability	Standard formats	ONNX
Batch processing	Throughput over latency	Batched inference
Numerical precision	Float handling	ndarray, careful f32/f64
Reproducibility	Deterministic	Seeded random, versioning

Critical Constraints

Memory Efficiency

RULE: Avoid copying large tensors
WHY: Memory bandwidth is bottleneck
RUST: References, views, in-place ops

GPU Utilization

RULE: Batch operations for GPU efficiency
WHY: GPU overhead per kernel launch
RUST: Batch sizes, async data loading

Model Portability

RULE: Use standard model formats
WHY: Train in Python, deploy in Rust
RUST: ONNX via tract or candle

Trace Down ↓

From constraints to design (Layer 2):

"Need efficient data pipelines"
    ↓ m10-performance: Streaming, batching
    ↓ polars: Lazy evaluation

"Need GPU inference"
    ↓ m07-concurrency: Async data loading
    ↓ candle/tch-rs: CUDA backend

"Need model loading"
    ↓ m12-lifecycle: Lazy init, caching
    ↓ tract: ONNX runtime

Use Case → Framework

Use Case	Recommended	Why
Inference only	tract (ONNX)	Lightweight, portable
Training + inference	candle, burn	Pure Rust, GPU
PyTorch models	tch-rs	Direct bindings
Data pipelines	polars	Fast, lazy eval

Key Crates

Purpose	Crate
Tensors	ndarray
ONNX inference	tract
ML framework	candle, burn
PyTorch bindings	tch-rs
Data processing	polars
Embeddings	fastembed

Design Patterns

Pattern	Purpose	Implementation
Model loading	Once, reuse	`OnceLock<Model>`
Batching	Throughput	Collect then process
Streaming	Large data	Iterator-based
GPU async	Parallelism	Data loading parallel to compute

Code Pattern: Inference Server

use std::sync::OnceLock;
use tract_onnx::prelude::*;

static MODEL: OnceLock<SimplePlan<TypedFact, Box<dyn TypedOp>, Graph<TypedFact, Box<dyn TypedOp>>>> = OnceLock::new();

fn get_model() -> &'static SimplePlan<...> {
    MODEL.get_or_init(|| {
        tract_onnx::onnx()
            .model_for_path("model.onnx")
            .unwrap()
            .into_optimized()
            .unwrap()
            .into_runnable()
            .unwrap()
    })
}

async fn predict(input: Vec<f32>) -> anyhow::Result<Vec<f32>> {
    let model = get_model();
    let input = tract_ndarray::arr1(&input).into_shape((1, input.len()))?;
    let result = model.run(tvec!(input.into()))?;
    Ok(result[0].to_array_view::<f32>()?.iter().copied().collect())
}

Code Pattern: Batched Inference

async fn batch_predict(inputs: Vec<Vec<f32>>, batch_size: usize) -> Vec<Vec<f32>> {
    let mut results = Vec::with_capacity(inputs.len());

    for batch in inputs.chunks(batch_size) {
        // Stack inputs into batch tensor
        let batch_tensor = stack_inputs(batch);

        // Run inference on batch
        let batch_output = model.run(batch_tensor).await;

        // Unstack results
        results.extend(unstack_outputs(batch_output));
    }

    results
}

Common Mistakes

Mistake	Domain Violation	Fix
Clone tensors	Memory waste	Use views
Single inference	GPU underutilized	Batch processing
Load model per request	Slow	Singleton pattern
Sync data loading	GPU idle	Async pipeline

Trace to Layer 1

Constraint	Layer 2 Pattern	Layer 1 Implementation
Memory efficiency	Zero-copy	ndarray views
Model singleton	Lazy init	OnceLock
Batch processing	Chunked iteration	chunks() + parallel
GPU async	Concurrent loading	tokio::spawn + GPU

Related Skills

When	See
Performance	m10-performance
Lazy initialization	m12-lifecycle
Async patterns	m07-concurrency
Memory efficiency	m01-ownership

Weekly Installs

433

Repository

zhanghandong/rust-skills

GitHub Stars

920

First Seen

Jan 20, 2026

Security Audits

Gen Agent Trust HubPass SocketPass SnykPass

Installed on

opencode395

codex379

gemini-cli373

github-copilot363

amp323

kimi-cli319

超能力技能使用指南：AI助手技能调用优先级与工作流程详解

41,800 周安装