Rust机器学习之tch-rs
Posted JarodYv
tags:
篇首语:本文由小常识网(cha138.com)小编为大家整理,主要介绍了Rust机器学习之tch-rs相关的知识,希望对你有一定的参考价值。
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年的搜索趋势:
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 60,000张图片,那么向量中将包含 60 , 000 × 28 × 28 = 47 , 040 , 000 60,000 \\times 28 \\times 28 = 47,040,000 60,000×28×28=47,040,000个元素,其中 28 28 28是图片行列的像素数。
MNIST数据集包含70,000张手写数字图片和其对应标签。每张照片 28 × 28 28 \\times 28 28×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类型。输出张量的大小为
d
i
m
1
×
(
d
i
m
2
×
d
i
m
3
)
dim1 \\times (dim2 \\times dim3)
dim1×(dim2×dim3),分别对应我们的训练数据集TRAIN_SIZE = 50000
,HEIGHT = 28
,WIDTH = 28
,因此输出张量的大小为
50000
×
(
28
×
28
)
=
50000
×
784
50000 \\times (28 \\times 28) = 50000 \\times 784
50000×(28×28)=50000×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机器学习之k-近邻算法实践学习