tch-rs: Rust绑定PyTorch C++ API的强大工具

tch-rs:将PyTorch的强大功能带入Rust世界

在机器学习和深度学习领域,PyTorch无疑是最受欢迎和广泛使用的框架之一。而随着Rust语言的日益流行,越来越多的开发者希望能够在Rust中使用PyTorch的强大功能。tch-rs应运而生,它为PyTorch的C++ API提供了Rust绑定,让Rust开发者能够方便地使用PyTorch构建和训练神经网络模型。

tch-rs的设计理念

tch-rs的核心目标是为PyTorch的C++ API提供一层薄薄的包装。它努力保持与原始C++ API的紧密联系,这意味着熟悉PyTorch C++的开发者可以轻松过渡到使用tch-rs。同时,tch-rs也为Rust开发者提供了一个基础,他们可以在此基础上开发更加符合Rust习惯的高级API。

tch-rs的主要特点包括:

与PyTorch C++ API的高度兼容性
提供Rust风格的API封装
支持自动微分和梯度下降优化
包含多种预训练模型的使用示例
丰富的文档和教程资源

安装和配置

要使用tch-rs,首先需要在系统上安装PyTorch的C++库(libtorch)。tch-rs支持多种安装和配置方式:

使用系统全局安装的libtorch(默认方式)
手动安装libtorch并通过环境变量指定路径
使用Python版PyTorch安装
通过特性标志自动下载预编译的libtorch二进制文件

对于Linux和macOS用户,可以通过设置环境变量来指定libtorch的路径:

export LIBTORCH=/path/to/libtorch
export LD_LIBRARY_PATH=${LIBTORCH}/lib:$LD_LIBRARY_PATH

Windows用户则需要在系统环境变量中设置相应的路径。

基本张量操作

tch-rs提供了一个包装PyTorch张量的Tensor类型。以下是一个简单的示例,展示了如何执行基本的张量操作:

use tch::Tensor;

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

这段代码创建了一个一维张量,将其所有元素乘以2,然后打印结果。

使用梯度下降训练模型

PyTorch的一个强大特性是自动微分,这使得使用梯度下降算法训练模型变得非常简单。tch-rs完全支持这一功能。以下是一个使用梯度下降优化简单模型的示例:

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);
    }
}

这个例子定义了一个简单的模块,其中包含两个可训练参数。然后,我们创建了一个随机梯度下降(SGD)优化器,并在一个循环中执行前向传播、计算损失和反向传播步骤。

构建神经网络

tch-rs的nn模块提供了构建复杂神经网络架构所需的工具。以下是一个在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(())
}

这个例子展示了如何定义一个简单的前馈神经网络,并使用Adam优化器在MNIST数据集上训练它。

MNIST Digits

使用预训练模型

tch-rs还支持使用预训练的计算机视觉模型。以下是一个使用预训练ResNet-18模型进行图像分类的例子:

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)?;
let output = resnet18
    .forward_t(&image.unsqueeze(0), /*train=*/ false)
    .softmax(-1);
for (probability, class) in imagenet::top(&output, 5).iter() {
    println!("{:50} {:5.2}%", class, 100.0 * probability)
}