tch-rs

Rust 对 PyTorch C++ api 的绑定。tch crate 的目标是提供对 C++ PyTorch api（即 libtorch）的一些薄封装。它旨在尽可能接近原始的 C++ api。然后可以在此之上开发更多符合 Rust 习惯的绑定。文档可以在 docs.rs 上找到。

变更日志

在 libtorch 之上的 C api 代码生成部分来自 ocaml-torch.

入门指南

此 crate 需要 C++ PyTorch 库（libtorch）版本 v2.4.0 在您的系统上可用。您可以选择以下方式：

• 使用系统范围内的 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 must contain `include` directory.
export LIBTORCH_INCLUDE=/path/to/libtorch/
# LIBTORCH_LIB must contain `lib` directory.
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的调试和发布版本不兼容ABI。如果使用了错误的libtorch版本，这可能导致某些段错误。

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

静态链接

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

git clone -b v2.4.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
# export LIBTORCH to point at the build directory in pytorch-static.

示例

基本张量运算

此crate提供了一个封装PyTorch张量的张量类型。以下是一个执行一些张量运算的最小示例。

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 {
        // Dummy mini-batches made of zeros.
        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

这将打印出图像的前 5 个 imagenet 类别。该示例的代码相当简单。

    // First the image is loaded and resized to 224x224.
    let image = imagenet::load_image_and_resize(image_file)?;

    // A variable store is created to hold the model parameters.
    let vs = tch::nn::VarStore::new(tch::Device::Cpu);

    // Then the model is built on this variable store, and the weights are loaded.
    let resnet18 = tch::vision::resnet::resnet18(vs.root(), imagenet::CLASS_COUNT);
    vs.load(weight_file)?;

    // Apply the forward pass of the model to get the logits and convert them
    // to probabilities via a softmax.
    let output = resnet18
        .forward_t(&image.unsqueeze(0), /*train=*/ false)
        .softmax(-1);

    // Finally print the top 5 categories and their associated probabilities.
    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<()> {
	// Create the model and load the pre-trained weights
	let mut vs = VarStore::new(Device::cuda_if_available());
	let model = resnet18(&vs.root(), 1000);
	vs.load("resnet18.safetensors")?;
	
	// Load the image file and resize it to the usual imagenet dimension of 224x224.
	let image = imagenet::load_image_and_resize224("dog.jpg")?
		.to_device(vs.device());

	// Apply the forward pass of the model to get the logits
	let output = image
		.unsqueeze(0)
		.apply_t(&model, false)
		.softmax(-1, Kind::Float);
	
	// Print the top 5 categories for this image.
    for (probability, class) in imagenet::top(&output, 5).iter() {
        println!("{:50} {:5.2}%", class, 100.0 * probability)
    }
    
    Ok(())
}

其他示例包括

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

外部材料

• 一个教程展示了如何使用Torch计算期权价格和greeks。
• tchrs-opencv-webcam-inference使用tch-rs和opencv在基于mobilenet v3的Python训练模型上进行Webcam流的推理。

常见问题解答

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

请参阅此线程中的详细内容。

如何在M1/M2 Mac上实现此功能？

查看此问题。

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

查看此问题，这可能是由于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安装的路径。

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