Rust机器学习之tch-rs
tch-rs是PyTorch接口的Rust绑定，可以认为tch-rs是Rust版的PyTorch。本文将带领大家学习如何用tch-rs搭建深度神经网络识别MNIST数据集中的手写数字。
本文是“Rust替代Python进行机器学习”系列文章的第五篇，其他教程请参考下面表格目录：
| Python库 | Rust替代方案 | 教程 |
|---|---|---|
numpy |
ndarray |
Rust机器学习之ndarray |
pandas |
Polars |
Rust机器学习之Polars |
scikit-learn |
Linfa |
Rust机器学习之Linfa |
matplotlib |
plotters |
Rust机器学习之plotters |
pytorch |
tch-rs |
Rust机器学习之tch-rs |
networks |
petgraph |
Rust机器学习之petgraph |
数据和算法工程师偏爱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::*;
实现思路
要完成这个神经网络的搭建，我们需要分三步：
- 下载并解压MNIST数据集，并将数据集中的图片转换为向量，共训练、验证和测试使用；
- 将向量转换为
Tensor类型，因为tch-rs的输入数据类型为Tensor类型； - 实现一系列迭代，每次迭代我们将输入数据和神经网络权重矩阵相乘，然后执行反向传播算法更新权重值。
我们下面一步一步来实现。
准备数据
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能够带给你更快的速度，这在大规模项目中是一个巨大的优势。
