tch-rs
为 PyTorch 的 C++ API 提供的 Rust 绑定。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 必须包含 `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 的调试和发布版本构建不兼容 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
# 导出 LIBTORCH 以指向 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
使用了两个初始值为 0 的变量 x1
和 x2
。应用于张量 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(())
}
有关训练循环的更多详细信息,可以在 详细教程 中找到。
使用一些预训练模型
预训练模型示例 演示了如何在图像上使用一些预训练的计算机视觉模型。 可以在这里下载权重(从 PyTorch 实现中提取):resnet18.ot 和 resnet34.ot。
然后可以通过以下命令运行示例:
cargo run --example pretrained-models -- resnet18.ot tiger.jpg
这应该会打印图像的前 5 个 imagenet 类别。该示例的代码非常简单。
// 首先加载图像并将其调整为 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)?;
<SOURCE_TEXT>
// 应用模型的前向传递以获取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模型。
- 一个min-GPT简化版类似于minGPT。
- 一个稳定扩散 实现,遵循huggingface's diffusers库的路线。
外部材料:
- 一个教程展示了如何使用Torch计算期权价格和greeks。
- tchrs-opencv-webcam-inference使用
tch-rs
和opencv
在网络摄像头上进行推理 一些基于mobilenet v3的Python训练模型。
常见问题
从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安装路径。
# 对于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。 </SOURCE_TEXT>