Rust机器学习之tch-rs

tch-rs是PyTorch接口的Rust绑定，可以认为tch-rs是Rust版的PyTorch。本文将带领大家学习如何用tch-rs搭建深度神经网络识别MNIST数据集中的手写数字。

本文是“Rust替代Python进行机器学习”系列文章的第五篇，其他教程请参考下面表格目录：

数据和算法工程师偏爱Jupyter，为了跟Python保持一致的工作环境，文章中的示例都运行在Jupyter上。因此需要各位搭建Rust交互式编程环境（让Rust作为Jupyter的内核运行在Jupyter上），相关教程请参考 Rust交互式编程环境搭建

初识tch-rs

PyTorch vs. TensorFlow

在深度学习领域，最受欢迎的开源框架非TensorFlow和PyTorch莫属。这两个框架都为构建和训练深度学习模型提供了广泛的功能，并已被研发社区广泛采用。目前二者无论从功能还是性能都非常接近，但PyTorch的接口设计更加“pythonic”且支持面向对象，相比之下，虽然TensorFlow提供更多选择给开发者，但接口和设计模式稍显混乱。因此，尽管TensorFlow诞生较早，但近年来PyTorch越来越受欢迎，已经超过TensorFlow。下图是谷歌趋势绘制的二者近5年的搜索趋势：

图1. TensorFlow vs. PyTorch
PyTorch已经超过TensorFlow成为最受欢迎的开源深度学习框架

tch-rs是PyTorch接口的Rust绑定，可以认为tch-rs是Rust版的PyTorch。tch-rs由Laurent Mazare 开发，是目前最Rustacean的PyTorch绑定，它对C++实现的libtorch进行了很薄的一层封装，这样做的最大优势是封装库与原始库严格相似，从而极大地降低了学习成本。如果你对PyTorch非常熟悉，几乎可以毫不费力得迁移到tch-rs上。

安装tch-rs

安装使用tch-rs非常简单，只需要在Cargo .toml加入

[dependencies]
tch = "0.8.0"

在机器学习中，我们更喜欢使用Jupyter。如果你已经搭建好Rust交互式编程环境（可以参考 《Rust交互式编程环境搭建》），可以直接通过下面代码引入tch-rs :

:dep tch = {
   version="0.8.0"}

初次编译tch-rs时间会有点长。但好在jupyter中cell之间是共享环境的，第一次编译加载完后，后面调用都很快。

用tch-rs搭建简单神经网络

环境准备

我们首先在MNIST数据集上训练一个简单得神经网络，为此我们需要mnist包来下载MNIST数据集（MNIST数据集的版权归Yann LeCun 和 Corinna Cortes所有，我们可以在 Creative Commons Attribution-Share Alike 3.0证书下获取使用），同时还需要引入ndarray包来对图片向量数据进行一些转换操作，并最终将其转换成tch::Tensor类型。（关于ndarray的使用请参考《Rust机器学习之ndarray》）。

:dep mnist = {
   version = "0.5.0", features = ["download"]}
:dep ndarray = {
   version = "0.15.6"}

use mnist::*;
use ndarray::prelude::*;

实现思路

要完成这个神经网络的搭建，我们需要分三步：

1. 下载并解压MNIST数据集，并将数据集中的图片转换为向量，共训练、验证和测试使用；
2. 将向量转换为Tensor类型，因为tch-rs的输入数据类型为Tensor类型；
3. 实现一系列迭代，每次迭代我们将输入数据和神经网络权重矩阵相乘，然后执行反向传播算法更新权重值。

我们下面一步一步来实现。

准备数据

mnist包中的MnistBuilder结构封装了下载、解压、加载、拆分等一系列数据准备工作，我们可以通过下面代码完成数据准备工作：

const TRAIN_SIZE: usize = 50000;
const VAL_SIZE: usize = 10000;
const TEST_SIZE: usize =10000;

let Mnist {
   
    trn_img,
    trn_lbl,
    val_img, 
    val_lbl,
    tst_img,
    tst_lbl,
} = MnistBuilder::new()
    .download_and_extract()
    .label_format_digit()
    .training_set_length(TRAIN_SIZE as u32)
    .validation_set_length(VAL_SIZE as u32)
    .test_set_length(TEST_SIZE as u32)
    .finalize();

• download_and_extract()：下载并解压MNIST数据集，该方法需要启用download特性
• label_format_digit()：将标签格式设为标量数字
• training_set_length(TRAIN_SIZE as u32)：拆分训练集
• validation_set_length(VAL_SIZE as u32)：拆分验证集
• test_set_length(TEST_SIZE as u32)：拆分测试集
• finalize()：根据上面的配置获取数据（Mnist结构类型）

返回值Mnist结构包含多个数据子集，在机器学习任务中,通常包含如下3类数据：

• 训练集 - 用于训练模型
• 验证集 - 用于训练过程中验证模型效果(MNIST默认数据分割中不包含验证集)
• 测试集 - 用于训练后评估模型表现

每个子集包含2个向量，一个向量保存图片数据，另一个向量保存标签。向量中的数据都是”平展“的，假如有$60,000$张图片,那么向量中将包含$60,000 \times 28 \times 28 = 47,040,000$个元素，其中$28$是图片行列的像素数。

MNIST数据集包含70,000张手写数字图片和其对应标签。每张照片$28 \times 28$像素，灰度值0到255。标签是图片对应的数字0到9。默认情况下60,000张划为训练集，10,000张划为测试集。

转成Tensor

use tch::{
   kind, no_grad, Kind, Tensor};

pub fn image_to_tensor(data:Vec<u8>, dim1:usize, dim2:usize, dim3:usize)-> Tensor{
   
    // 将Vec转换为三维数组并将颜色值进行归一化处理 
    let inp_data: Array3<f32> = Array3::from_shape_vec((dim1, dim2, dim3), data)
        .expect("Error converting data to 3D array")
        .map(|x| *x as f32/256.0);
    // 转成Tensor
    let inp_tensor = Tensor::of_slice(inp_data.as_slice().unwrap());
    // 将Tensor转换成 [dim1, dim2*dim3] 结构的张量
    let ax1 = dim1 as i64; 
    let ax2 = (dim2 as i64)*(dim3 as i64);
    let shape: Vec<i64>  = vec![ ax1, ax2 ];
    let output_data = inp_tensor.reshape(&shape);
    println!("Output image tensor size {:?}", shape);

    output_data
}

上面的代码利用from_shape_vec将输入的Vec<u8>类型数据转换成Array3，.map(|x| *x as f32/256.0)对数值进行了归一化，并转换成f32类型。tch-rs提供了Tensor::of_slice方法，可以方便地将数组转换为torch Tensor类型。输出张量的大小为$dim1 \times (dim2 \times dim3)$，分别对应我们的训练数据集TRAIN_SIZE = 50000，HEIGHT = 28，WIDTH = 28，因此输出张量的大小为$50000 \times (28 \times 28) = 50000 \times 784$。

同理，我们需要将标记数据也转成Tensor，它的大小为dim1——因此，对应训练集标记数据我们需要一个大小为50000的张量。代码如下：

pub fn labels_to_tensor(data:Vec<u8>, dim1:usize, dim2:usize)-> Tensor{
   
    let inp_data: Array2<i64> = Array2::from_shape_vec((dim1, dim2), data)
        .expect("Error converting data to 2D array")
        .map(|x| *x as i64);

    let output_data = Tensor::of_slice(inp_data.as_slice().unwrap());
    println!("Output label tensor size {:?}", output_data.size());

    output_data
}

构建模型

现在，我们可以开始着手构建我们的线性神经网络模型了。

首先我们将权重矩阵和误差矩阵设为0：

let mut ws = Tensor::zeros(&[(HEIGHT*WIDTH) as i64, LABELS], kind::FLOAT_CPU).set_requires_grad(true);
    let mut bs = Tensor::zeros(&[LABELS], kind::FLOAT_CPU).set_requires_grad(true);

然后循环迭代训练线性神经网络

const LABELS: i64 = 10; // 标签类别数量
const HEIGHT: usize = 28; 
const WIDTH: usize = 28;
const N_EPOCHS: i64 = 200; // 迭代次数
const THRES: f64 = 0.001; // 阈值

let mut loss_diff;
let mut curr_loss = 0.0;

// 开始训练
'train: for epoch in 1..N_EPOCHS{
   
    // neural network multiplication
    let logits = train_data.matmul(&ws) + &bs; 
    // 用log softmax计算loss
    let loss = logits.log_softmax(-1, Kind::Float).nll_loss(&train_lbl);
    // 处理梯度
    ws.zero_grad();
    bs.zero_grad();
    loss.backward();
    // 反向传播
    no_grad(|| {
   
        ws += ws.grad()*(-1);
        bs += bs.grad()*(-1);
    });
    // 验证
    let val_logits = val_data.matmul(&ws) + &bs;
    let val_accuracy = val_logits
            .argmax(Some(-1), false)
            .eq_tensor(&val_lbl)
            .to_kind(Kind::Float)
            .mean(Kind::Float)
            .double_value(&[]);

    println!(
            "epoch: {:4} train loss: {:8.5} val acc: {:5.2}%",
            epoch,
            loss.double_value(&[]),
            100. * val_accuracy
    );
    // 判断是否达到精度要求 
    if epoch == 1{
   
        curr_loss = loss.double_value(&[]);
    } else {
   
        loss_diff = (loss.double_value(&[]) - curr_loss).abs(); 
        curr_loss = loss.double_value(&[]); 
        // 如果loss小于阈值则停止循环
        if loss_diff < THRES {
   
            println!("Target accuracy reached, early stopping");
            break 'train;
        }
    }
} 

// 在测试集上测试模型效果
let test_logits = test_data.matmul(&ws) + &bs; 
let test_accuracy = test_logits
        .argmax(Some(-1), false)
        .eq_tensor(&test_lbl)
        .to_kind(Kind::Float)
        .mean(Kind::Float)
        .double_value(&[]);
println!("Final test accuracy {:5.2}%", 100.*test_accuracy);

上面代码主体逻辑是一个循环，我们将其命名为'train。循环中我们监控每次迭代的loss，如果连续两次循环的loss差小于给定阈值THRES则结束循环（这里的处理不一定合理，但是为了演示简单起见，我们暂且这样处理）。整体逻辑非常简单，就是最最简单的神经网络，相信大家都能理解其逻辑，我这里不做过多的赘述。

我们执行上面代码即可训练模型，由于模型简单，在我的笔记本上大约十几秒即可训练完成，最终准确率90.45%。

用tch-rs搭建序贯神经网络

我们再来看一下序贯神经网络的实现。

首先，我们需要引入tch::nn::Module，然后实现fn net(vs: &nn::Path) -> impl Module函数。该函数接收nn::Path输入参数，表示运行神经网络的硬件信息（例如CPU还是GPU），返回一个Module实现。

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

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

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

接着我们通过如下代码创建神经网络:

// 创建变量保存CUDA是否可用
let vs = nn::VarStore::new(Device::cuda_if_available());
// 创建序贯网络
let net = net(&vs.root());
// 创建优化器
let mut opt = nn::Adam::default().build(&vs, 1e-4)?;

这里我们使用Adam优化器。然后，我们可以简单地按照PyTorch的步骤进行操作，我们需要多轮迭代，并使用优化器的backward_step方法执行反向传播，代码如下：

for epoch in 1..N_EPOCHS {
   
        let loss = net.forward(&train_data).cross_entropy_for_logits(&train_lbl);
        // 反向传播 
        opt.backward_step(&loss);
        // 计算测试集上的精度
        let val_accuracy = net.forward(&val_data).accuracy_for_logits(&val_lbl);
        println!(
            "epoch: {:4} train loss: {:8.5} val acc: {:5.2}%",
            epoch,
            f64::from(&loss),
            100. * f64::from(&val_accuracy),
        );
    }

经过大约1分钟的训练，最终模型准确率85.50%

用tch-rs搭建卷积神经网络

我们日常用的最多的神经网络当属卷积神经网络，文章最后我们看一下如何用tch-rs实现卷积神经网络。

首先我们需要先引入nn::ModuleT，该模块特性是一个附加的训练参数，通常用于区分训练和评估之间的网络行为。然后，我们定义结构体Net，它由两个conv2d层和两个线性层组成。

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

#[derive(Debug)]
struct Net {
   
    conv1: nn::Conv2D,
    conv2: nn::Conv2D,
    fc1: nn::Linear,
    fc2: nn::Linear,
}

Net结构的实现定义了网络如何构成。两个卷积层的步长（Stride）分别为1和32，填充（Padding）分别为32和64，扩张（Dilation ）分别为5和5。线性层接收1024个输入，最终层返回10个元素的输出。

impl Net {
   
    fn new(vs: &nn::Path) -> Net {
   
        let conv1 = nn::conv2d(vs, 1, 32, 5, Default::default());
        let conv2 = nn::conv2d(vs, 32, 64, 5, Default::default());
        let fc1 = nn::linear(vs, 1024, 1024, Default::default());
        let fc2 = nn::linear(vs, 1024, 10, Default::default());
        Net {
    conv1, conv2, fc1, fc2 }
    }
}

最后，我们要实现Net的ModuleT模块特性。这里前向步骤forward_t接收一个额外的布尔参数train表示是否为训练集，返回一个Tensor张量。前向步骤会用到卷积层以及max_pool_2d和dropout。dropout仅用于训练目的，因此要传入布尔变量train。

impl nn::ModuleT for Net {
   
    fn forward_t(&self, xs: &Tensor, train: bool) -> Tensor {
   
        xs.view([-1, 1, 28, 28])
            .apply(&self.conv1)
            .max_pool2d_default(2)
            .apply(&self.conv2)
            .max_pool2d_default(2)
            .view([-1, 1024])
            .apply(&self.fc1)
            .relu()
            .dropout(0.5, train)
            .apply(&self.fc2)
    }
}

为了提高训练性能，我们使用张量批处理来训练卷积层。为此，我们额外实现一个函数generate_random_index，将输入张量随机拆分为指定大小的批次：

const BATCH_SIZE: i64 = 256;

pub fn generate_random_index(ArraySize: i64, BatchSize: i64)-> Tensor{
   
    let random_idxs = Tensor::randint(ArraySize, &[BatchSize], kind::INT64_CPU);
    random_idxs
}

训练过程依然是一个循环迭代。输入数据被拆分为n_it个批次，对每一批数据我们通过网络计算loss并用backward_step反向传播误差。代码如下：

let n_it = (TRAIN_SIZE as i64) / BATCH_SIZE;

for epoch in 1..N_EPOCHS {
   
        // generate random idxs for batch size 
        // run all the images divided in batches  -> for loop
        for i in 1..n_it {
   
            let batch_idxs = generate_random_index(TRAIN_SIZE as i64, BATCH_SIZE); 
            let batch_images = train_data.index_select(0, &batch_idxs).to_device(vs.device()).to_kind(Kind::Float); 
            let batch_lbls = train_lbl.index_select(0, &batch_idxs).to_device(vs.device()).to_kind(Kind::Int64);
            // compute the loss 
            let loss = net.forward_t(&batch_images, true).cross_entropy_for_logits(&batch_lbls);
            opt.backward_step(&loss);
        }
        // compute accuracy 
        let val_accuracy =
            net.batch_accuracy_for_logits(&val_data, &val_lbl, vs.device(), 1024);
        println!("epoch: {:4} test acc: {:5.2}%", epoch, 100. * val_accuracy,);
    }

在我的笔记本电脑上运行卷积网络需要几分钟，验证准确率达到97.40%。

总结

整体上tch-rs的使用思路和PyTorch是一致的，因为本身tch-rs就是PyTorch的C++库libtorch的绑定。如果你熟练使用PyTorch，那么用tch-rs上手会非常快。关键是用tch-rs能够带给你更快的速度，这在大规模项目中是一个巨大的优势。