tch-rs

PyTorch C++ API 的 Rust 绑定。tch crate 的目标是为 PyTorch 的 C++ API（即 libtorch）提供一些轻量级的封装。它力求尽可能贴近原始的 C++ API，以便在此基础上进一步开发更符合 Rust 习惯的绑定。相关文档可在 docs.rs 上找到。

变更日志

基于 libtorch 的 C API 代码生成部分源自 ocaml-torch。

快速入门

本 crate 需要系统中已安装 v2.11.0 版本的 PyTorch C++ 库（libtorch）。您可以选择以下方式之一：

• 使用系统全局的 libtorch 安装（默认）。
• 手动安装 libtorch，并通过 LIBTORCH 环境变量告知构建脚本其位置。
• 使用 Python 安装的 PyTorch，为此设置 LIBTORCH_USE_PYTORCH=1。
• 如果无法找到系统全局的 libtorch 且未设置 LIBTORCH，构建脚本可以通过启用 download-libtorch 功能下载预编译的 libtorch 二进制文件。默认使用 CPU 版本；可通过设置 TORCH_CUDA_VERSION 环境变量为 cu117 来获取使用 CUDA 11.7 的预编译二进制文件。

系统全局的 Libtorch

在 Linux 平台上，构建脚本会在 /usr/lib/libtorch.so 中查找系统全局的 libtorch 库。

Python PyTorch 安装

如果设置了 LIBTORCH_USE_PYTORCH 环境变量，则会调用当前激活的 Python 解释器来获取 torch Python 包的相关信息，并链接到该版本。

手动安装 Libtorch

• 从 PyTorch 官网的下载页面 获取 libtorch，并解压 ZIP 文件。
• 对于 Linux 和 macOS 用户，在 .bashrc 或等效配置文件中添加以下内容，其中 /path/to/libtorch 是解压后创建的目录路径：

export LIBTORCH=/path/to/libtorch

还可以分别指定头文件和共享库的位置：

# LIBTORCH_INCLUDE 必须包含 `include` 目录。
export LIBTORCH_INCLUDE=/path/to/libtorch/
# LIBTORCH_LIB 必须包含 `lib` 目录。
export LIBTORCH_LIB=/path/to/libtorch/

• 对于 Windows 用户，假设 X:\path\to\libtorch 是解压后的 libtorch 目录：

  • 打开控制面板 -> 查看高级系统设置 -> 环境变量。
  • 创建 LIBTORCH 变量并将其设置为 X:\path\to\libtorch。
  • 将 X:\path\to\libtorch\lib 追加到 Path 变量中。

  如果您希望临时设置环境变量，可以在 PowerShell 中运行：

$Env:LIBTORCH = "X:\path\to\libtorch"
$Env:Path += ";X:\path\to\libtorch\lib"

• 此时您应该可以运行一些示例，例如 cargo run --example basics。

Windows 特别说明

根据 PyTorch 文档，Windows 下的 Debug 和 Release 构建不兼容 ABI。如果使用了错误版本的 libtorch，可能会导致段错误。

建议使用 MSVC Rust 工具链（例如通过 rustup install stable-x86_64-pc-windows-msvc 安装），而不是基于 MinGW 的工具链，因为 PyTorch 与 MinGW 存在兼容性问题。

静态链接

当设置环境变量 LIBTORCH_STATIC=1 时，将静态链接 libtorch，而非使用动态库。预编译的工件默认不包含 libtorch.a，因此需要手动编译，例如：

git clone -b v2.11.0 --recurse-submodule https://github.com/pytorch/pytorch.git pytorch-static --depth 1
cd pytorch-static
USE_CUDA=OFF BUILD_SHARED_LIBS=OFF python setup.py build
# 将 LIBTORCH 导出指向 pytorch-static 中的构建目录。

示例

基本张量操作

本 crate 提供了一个包装 PyTorch 张量的 Tensor 类型。以下是一个执行基本张量操作的最小示例。

use tch::Tensor;

fn main() {
    let t = Tensor::from_slice(&[3, 1, 4, 1, 5]);
    let t = t * 2;
    t.print();
}

使用梯度下降训练模型

PyTorch 为其支持的大多数张量运算提供了自动微分功能，这通常用于通过梯度下降训练模型。优化过程是在通过 nn::VarStore 创建的变量上进行的，这些变量需要定义其形状和初始值。

在下面的示例中，my_module 使用两个变量 x1 和 x2，它们的初始值均为 0。前向传播作用于张量 xs 后，返回 xs * x1 + exp(xs) * x2。

模型生成后，创建一个 nn::Sgd 优化器。在训练循环的每一步中：

• 对一个小批量数据应用前向传播。
• 计算损失，即模型输出与小批量真实标签之间的均方误差。
• 最后执行一次优化步骤：计算梯度，并相应地更新 VarStore 中的变量。

use tch::nn::{Module, OptimizerConfig};
use tch::{kind, nn, Device, Tensor};

fn my_module(p: nn::Path, dim: i64) -> impl nn::Module {
    let x1 = p.zeros("x1", &[dim]);
    let x2 = p.zeros("x2", &[dim]);
    nn::func(move |xs| xs * &x1 + xs.exp() * &x2)
}

fn gradient_descent() {
    let vs = nn::VarStore::new(Device::Cpu);
    let my_module = my_module(vs.root(), 7);
    let mut opt = nn::Sgd::default().build(&vs, 1e-2).unwrap();
    for _idx in 1..50 {
        // 由零组成的虚拟小批量。
        let xs = Tensor::zeros(&[7], kind::FLOAT_CPU);
        let ys = Tensor::zeros(&[7], kind::FLOAT_CPU);
        let loss = (my_module.forward(&xs) - ys).pow_tensor_scalar(2).sum(kind::Kind::Float);
        opt.backward_step(&loss);
    }
}

编写一个简单的神经网络

nn API 可以用于构建神经网络架构。例如，以下代码定义了一个包含一个隐藏层的简单模型，并使用 Adam 优化器在 MNIST 数据集上对其进行训练。

use anyhow::Result;
use tch::{nn, nn::Module, nn::OptimizerConfig, Device};

const IMAGE_DIM: i64 = 784;
const HIDDEN_NODES: i64 = 128;
const LABELS: i64 = 10;

fn net(vs: &nn::Path) -> impl Module {
    nn::seq()
        .add(nn::linear(
            vs / "layer1",
            IMAGE_DIM,
            HIDDEN_NODES,
            Default::default(),
        ))
        .add_fn(|xs| xs.relu())
        .add(nn::linear(vs, HIDDEN_NODES, LABELS, Default::default()))
}

pub fn run() -> Result<()> {
    let m = tch::vision::mnist::load_dir("data")?;
    let vs = nn::VarStore::new(Device::Cpu);
    let net = net(&vs.root());
    let mut opt = nn::Adam::default().build(&vs, 1e-3)?;
    for epoch in 1..200 {
        let loss = net
            .forward(&m.train_images)
            .cross_entropy_for_logits(&m.train_labels);
        opt.backward_step(&loss);
        let test_accuracy = net
            .forward(&m.test_images)
            .accuracy_for_logits(&m.test_labels);
        println!(
            "epoch: {:4} train loss: {:8.5} test acc: {:5.2}%",
            epoch,
            f64::from(&loss),
            100. * f64::from(&test_accuracy),
        );
    }
    Ok(())
}

有关训练循环的更多详细信息，请参阅详细教程。

使用预训练模型

pretrained-models 示例展示了如何在图像上使用一些预训练的计算机视觉模型。这些权重是从 PyTorch 实现中提取的，可以在这里下载：resnet18.ot 和 resnet34.ot。

然后可以通过以下命令运行该示例：

cargo run --example pretrained-models -- resnet18.ot tiger.jpg

这将打印出图像对应的 ImageNet 前 5 类别及其概率。该示例的代码非常简单。

    // 首先加载图像并将其调整为 224x224 大小。
    let image = imagenet::load_image_and_resize(image_file)?;

    // 创建一个变量存储来保存模型参数。
    let vs = tch::nn::VarStore::new(tch::Device::Cpu);

    // 然后在此变量存储上构建模型，并加载权重。
    let resnet18 = tch::vision::resnet::resnet18(vs.root(), imagenet::CLASS_COUNT);
    vs.load(weight_file)?;

    // 对模型执行前向传播以获得 logits，并通过 softmax 将其转换为概率。
    let output = resnet18
        .forward_t(&image.unsqueeze(0), /*train=*/ false)
        .softmax(-1);

    // 最后打印出前 5 类别及其对应的概率。
    for (probability, class) in imagenet::top(&output, 5).iter() {
        println!("{:50} {:5.2}%", class, 100.0 * probability)
    }

使用 SafeTensors 从 PyTorch 导入预训练权重

safetensors 是 HuggingFace 推出的一种新的简单张量存储格式。它不依赖于 Python 的 pickle 模块，因此张量不会绑定到保存模型时使用的特定类和确切的目录结构。此外，它是零拷贝的，这意味着读取文件所需的内存与原始文件大小相同。

有关 safetensors 的更多信息，请参阅：https://github.com/huggingface/safetensors

安装 safetensors

您可以通过 pip 包管理器安装 safetensors：

pip install safetensors

在 PyTorch 中导出权重

import torchvision
from safetensors import torch as stt

model = torchvision.models.resnet18(pretrained=True)
stt.save_file(model.state_dict(), 'resnet18.safetensors')

注意：导出文件的文件名必须以 .safetensors 为后缀，这样才能被 tch 正确解码。

在 tch 中导入权重

use anyhow::Result;
use tch::{
	Device,
	Kind,
	nn::VarStore,
	vision::{
		imagenet,
		resnet::resnet18,
	}
};

fn main() -> Result<()> {
	// 创建模型并加载预训练权重
	let mut vs = VarStore::new(Device::cuda_if_available());
	let model = resnet18(&vs.root(), 1000);
	vs.load("resnet18.safetensors")?;
	
	// 加载图像文件，并将其调整为 ImageNet 常用的 224x224 尺寸。
	let image = imagenet::load_image_and_resize224("dog.jpg")?
		.to_device(vs.device());

	// 对模型进行前向传播，得到 logits
	let output = image
		.unsqueeze(0)
		.apply_t(&model, false)
		.softmax(-1, Kind::Float);
	
	// 打印该图像的前 5 类别。
    for (probability, class) in imagenet::top(&output, 5).iter() {
        println!("{:50} {:5.2}%", class, 100.0 * probability)
    }
    
    Ok(())
}

更多示例包括：

• 一个简化的 char-rnn 示例，展示了使用循环神经网络进行字符级语言建模。
• 神经风格迁移 使用预训练的 VGG-16 模型，将一张图像转换为另一张图像的风格（预训练权重：vgg16.ot）。
• 一些在 CIFAR-10 数据集上的 ResNet 示例。
• 一个 教程，演示如何使用 TorchScript JIT 部署和运行由 Python 训练的模型。
• 一些使用 OpenAI Gym 环境的 强化学习示例，其中包括策略梯度示例以及可在 Atari 游戏上运行的 A2C 实现。
• 一个 迁移学习教程，展示了如何在一个非常小的数据集上对预训练的 ResNet 模型进行微调。
• 一个类似于 minGPT 的 简化版 GPT。
• 一个遵循 HuggingFace diffusers 库思路的 Stable Diffusion 实现。

外部资料：

• 一个 教程，展示如何使用 Torch 计算期权价格和希腊字母。
• tchrs-opencv-webcam-inference 使用 tch-rs 和 opencv 对基于 mobilenet v3 的 Python 训练模型进行网络摄像头推理。

常见问题解答

从 Python 到 Rust 模型转换的最佳实践是什么？

请参阅 此讨论 以获取详细信息。

如何在 M1/M2 Mac 上运行此代码？

请查看 此问题。

编译速度很慢，每次运行 cargo 时似乎都会重新构建 torch-sys。

请参阅 此问题，这可能是由于 rust-analyzer 不了解正确的环境变量（如 LIBTORCH 和 LD_LIBRARY_PATH）所致。

从 Python 调用 Rust/tch 代码。

可以通过 PyO3 从 Python 调用 Rust/tch 代码，tch-ext 提供了一个这样的 Python 扩展示例。

加载共享库时出错。

如果您在运行生成的二进制文件时遇到找不到某些共享库的错误（例如： error while loading shared libraries: libtorch_cpu.so: cannot open shared object file: No such file or directory）， 您可以尝试将以下内容添加到您的 .bashrc 文件中，其中 /path/to/libtorch 是您 libtorch 的安装路径。

# 对于 Linux
export LD_LIBRARY_PATH=/path/to/libtorch/lib:$LD_LIBRARY_PATH
# 对于 macOS
export DYLD_LIBRARY_PATH=/path/to/libtorch/lib:$DYLD_LIBRARY_PATH

许可证

tch-rs 根据您的选择，同时采用 MIT 许可证和 Apache 许可证（版本 2.0）进行分发。

有关详细信息，请参阅 LICENSE-APACHE 和 LICENSE-MIT。

tch-rs 快速上手指南

tch-rs 是 PyTorch C++ API (libtorch) 的 Rust 绑定库。它旨在提供轻量级的封装，尽可能贴近原始的 C++ 接口，让 Rust 开发者能够利用 PyTorch 的强大功能进行深度学习开发。

环境准备

系统要求

• Rust: 建议安装最新稳定版 (rustup)。
• 操作系统: Linux, macOS, Windows。
• 编译器:
  • Linux/macOS: GCC 或 Clang。
  • Windows: 强烈推荐使用 MSVC 工具链 (例如 stable-x86_64-pc-windows-msvc)。MinGW 可能与 PyTorch 存在兼容性问题。

核心依赖：LibTorch

本库依赖 C++ 版的 PyTorch 库 (libtorch)，版本需为 v2.11.0。你有以下几种方式获取：

1. 自动下载（推荐新手）：启用 download-libtorch 特性，构建脚本会自动下载预编译的二进制文件。
2. 使用现有 Python 环境：如果你已安装 torch Python 包，可复用其底层库。
3. 手动安装：从官网下载并配置环境变量。

国内加速提示：如果自动下载速度慢，建议先手动下载 libtorch 压缩包（见下文“手动安装”），或通过设置环境变量指向国内镜像源下载的文件。

安装步骤

方法一：自动下载（最简单）

在 Cargo.toml 中添加依赖并启用下载特性：

[dependencies]
tch = { version = "0.15", features = ["download-libtorch"] }

注：默认下载 CPU 版本。如需 CUDA 支持，请在运行 cargo 前设置环境变量：

export TORCH_CUDA_VERSION=cu117
cargo build

方法二：复用 Python PyTorch 环境

如果你已经通过 pip 安装了 PyTorch：

1. 确保 torch 已安装：pip install torch
2. 在 Cargo.toml 中仅添加依赖（无需特殊 feature）：

   [dependencies]
   tch = "0.15"
   

3. 构建前设置环境变量告诉构建脚本使用 Python 的库：

   export LIBTORCH_USE_PYTORCH=1
   cargo build
   

方法三：手动安装 LibTorch

1. 下载：访问 PyTorch 官网下载页 或国内镜像，下载 LibTorch (C++) 版本 v2.11.0 (选择对应的 CPU 或 CUDA 版本)。

2. 解压：将下载的 zip 文件解压到本地目录，例如 /path/to/libtorch。

3. 配置环境变量：

   Linux / macOS: 在 .bashrc 或 .zshrc 中添加：

   export LIBTORCH=/path/to/libtorch
   # 如果需要区分头文件和库文件路径（通常不需要，指向根目录即可）
   # export LIBTORCH_INCLUDE=/path/to/libtorch/include
   # export LIBTORCH_LIB=/path/to/libtorch/lib
   

   然后执行 source ~/.bashrc。

   Windows (PowerShell):

   $Env:LIBTORCH = "X:\path\to\libtorch"
   $Env:Path += ";X:\path\to\libtorch\lib"
   

   (永久设置可在“系统属性 -> 环境变量”中配置 LIBTORCH 变量，并将 lib 目录追加到 Path 中)

4. 构建项目：

   cargo build
   

基本使用

1. 基础张量操作

创建一个简单的张量并进行数学运算：

use tch::Tensor;

fn main() {
    // 从切片创建张量
    let t = Tensor::from_slice(&[3, 1, 4, 1, 5]);
    
    // 执行乘法操作
    let t = t * 2;
    
    // 打印结果
    t.print();
}

2. 简单的梯度下降训练

利用 tch 的自动微分功能训练一个简单的模型：

use tch::nn::{Module, OptimizerConfig};
use tch::{kind, nn, Device, Tensor};

// 定义一个简单的模块：f(x) = x * x1 + exp(x) * x2
fn my_module(p: nn::Path, dim: i64) -> impl nn::Module {
    let x1 = p.zeros("x1", &[dim]);
    let x2 = p.zeros("x2", &[dim]);
    nn::func(move |xs| xs * &x1 + xs.exp() * &x2)
}

fn gradient_descent() {
    // 创建变量存储区
    let vs = nn::VarStore::new(Device::Cpu);
    
    // 实例化模型
    let my_module = my_module(vs.root(), 7);
    
    // 创建 SGD 优化器
    let mut opt = nn::Sgd::default().build(&vs, 1e-2).unwrap();
    
    // 训练循环
    for _idx in 1..50 {
        // 构造虚拟数据 (全零)
        let xs = Tensor::zeros(&[7], kind::FLOAT_CPU);
        let ys = Tensor::zeros(&[7], kind::FLOAT_CPU);
        
        // 前向传播并计算损失 (MSE)
        let loss = (my_module.forward(&xs) - ys).pow_tensor_scalar(2).sum(kind::Kind::Float);
        
        // 反向传播并更新参数
        opt.backward_step(&loss);
    }
}

fn main() {
    gradient_descent();
}

3. 运行示例

安装完成后，你可以直接运行官方提供的示例来验证环境：

# 运行基础示例
cargo run --example basics

# 运行 MNIST 训练示例 (需准备数据)
cargo run --example mnist