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替代方案教程
numpyndarrayRust机器学习之ndarray
pandasPolars Rust机器学习之Polars
scikit-learnLinfaRust机器学习之Linfa
matplotlibplottersRust机器学习之plotters
pytorchtch-rsRust机器学习之tch-rs
networkspetgraphRust机器学习之petgraph

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

文章目录

初识tch-rs

PyTorch vs. TensorFlow

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

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

tch-rs是PyTorch接口的Rust绑定,可以认为tch-rs是Rust版的PyTorch。tch-rsLaurent 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 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 = 50000HEIGHT = 28WIDTH = 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-近邻算法实践学习

深度学习之概述(Overview)

Python大数据与机器学习之NumPy初体验

机器学习之学习路线

深度学习之梯度

机器学习之激活函数