在本文中，分享尝试使用Rust从头创建一个小型机器学习(ML)框架的经验。

对于这个实验，有以下目标：

1.调查将Python + PyTorch转换为Rust + LibTorch（PyTorch的C++后端库）是否能带来明显的速度提升，特别是在模型训练过程中。我们知道，ML模型正在变得越来越大，因此需要越来越多的计算能力来训练(有时对普通人来说是不可行的)。减轻硬件需求不断增加的一种方式是找到一种使算法更具计算效率的方法。知道在PyTorch中，Python只是作为LibTorch之上的一层，最大的问题是用Rust替换Python顶层是否值得。计划是使用Tch-rs Rust crate只是接触到LibTorch DLL的张量和Autograd功能，因此作为我们的“梯度计算器”，然后从头开始在Rust中开发其余部分。

2.希望保持代码足够简单，以便能够清楚地理解所执行的所有线性代数，并允许我在需要时轻松扩展它。

3.框架必须尽可能地允许我按照标准Python/PyTorch的类似结构定义ML模型。

这篇文章的目的并不是教授Rust本身，而是展示如何将Rust应用于机器学习以及其中的好处。

直接跳到最终结果，我的小框架如下所示的神经网络模型:

清单1 -定义神经网络模型

struct MyModel {

    l1: Linear,

    l2: Linear,

}



impl MyModel {

    fn new (mem: &mut Memory) -> MyModel {

        let l1 = Linear::new(mem, 784, 128);

        let l2 = Linear::new(mem, 128, 10);

        Self {

            l1: l1,

            l2: l2,

        }

    }

}



impl Compute for MyModel {

    fn forward (&self, mem: &Memory, input: &Tensor) -> Tensor {

        let mut o = self.l1.forward(mem, input);

        o = o.relu();

        o = self.l2.forward(mem, &o);

        o

    }

}

然后像这样实例化和训练模型。

清单2 -实例化和训练神经网络模型

fn main() {

   let (x, y) = load_mnist();



   let mut m = Memory::new();

   let mymodel = MyModel::new(&mut m);

   train(&mut m, &x, &y, &mymodel, 100, 128, cross_entropy, 0.3);

   let out = mymodel.forward(&m, &x);

   println!("Training Accuracy: {}", accuracy(&y, &out));

}

对于PyTorch用户来说，上面的内容与在Python中定义和训练神经网络的方式非常直观地相似。上面的示例显示了一个神经网络模型，然后将其用于分类(该模型应用于Mnist数据集，我将使用该数据集作为基准数据集来比较Rust - Python模型版本)。

在代码的第一个块中，创建了一个MyModel结构体，它包含两层线性类型。

第二个块是MyModel结构实现，它定义了一个关联函数new。这个函数初始化这两层并返回一个结构的新实例。

最后，第三个代码块为MyModel实现了Compute trait，其中定义了forward方法。在main函数中，加载了Mnist数据集，初始化内存，实例化了MyModel，然后使用100个Epochs、批量大小为128、交叉熵损失和学习率为0.3进行训练。

很直观吧?这就是使用我的小框架在Rust中创建和训练新模型所需要的。然而，我们现在开始在引擎盖下看看是什么使上述成为可能。

查看上面的代码，如果你习惯于在PyTorch中构建ML模型，一个明显的问题可能会弹出-内存引用在做什么?我在下面解释。

前传

从机器学习的文献中，我们知道神经网络的训练机制是通过迭代地进行两个步骤来进行的，这些步骤在一定数量的epochs（通常也是在一定数量的批次中）中进行，即前向传播和反向传播（两个步骤）。

在前向传递中，我们将输入和随后的计算推入网络中的所有层，其中对于每一层，我们都有：

公式1

其中w提供线性函数的权重，b提供偏差，然后通过激活函数传递，比如提供非线性的Sigmoid。

有了这些信息，我们现在可以创建我们的Linear层(下面的清单3)。你可以注意到，定义层的结构遵循定义模型的结构(上面的清单1)，并实现相同的功能和特征。

在线性层的情况下，该结构包含一个名为params的字段。params字段是一个HashMap类型的集合，其中key的类型是String，它存储参数名称，value的类型是usize，它保存特定参数(它是PyTorch张量)在Memory中的位置，反过来充当我们所有参数的存储。

清单3 -定义神经网络层，在本例中是线性层

trait Compute {

    fn forward (&self, mem: &Memory, input: &Tensor) -> Tensor;

}



struct Linear {

    params: HashMap,

}



impl Linear {

    fn new (mem: &mut Memory, ninputs: i64, noutputs: i64) -> Self {

       let mut p = HashMap::new();

       p.insert("W".to_string(), mem.new_push(&[ninputs,noutputs], true));

       p.insert("b".to_string(), mem.new_push(&[1, noutputs], true));



       Self {

           params: p,

       }

    }

}



impl Compute for Linear {

    fn forward (&self, mem: &Memory, input: &Tensor) -> Tensor {

        let w = mem.get(self.params.get(&"W".to_string()).unwrap());

        let b = mem.get(self.params.get(&"b".to_string()).unwrap());

        input.matmul(w) + b

    }

}

根据公式1，在我们的关联函数new中，我们在 HashMap 上插入线性层所需的两个参数“W”和“ b” 。

稍后介绍的memm .new_push()方法，以所需的大小创建各自的张量，将它们推入内存存储，并返回它们的位置。insert方法中的布尔参数定义了我们需要计算这些参数的梯度。通过这种方式，每一层将包含参数名称和它们在内存结构中各自的张量存储位置。

与MyModel定义类似，我们随后为我们的线性层实现Compute特性。这需要定义forward函数，该函数在训练过程的前向传播过程中被调用。

在这个函数中，我们首先使用get方法从我们的张量存储中获得对两个张量参数的引用，然后计算我们的线性函数(公式1)。与PyTorch类似，从我们的神经网络中，我们输出非规范化预测(logits)并在稍后的误差计算期间执行规范化(在本例中为Softmax)。

有人可能会问，为什么采用这种方法来表示线性层，而不是直接在一两行代码中硬编码公式1 ?

采用这种方法的目的是，如果需要定义额外的神经网络层类型，例如CNN或LSTM层，那么只需精确地复制上述线性层结构并在相关函数中注入额外的参数和计算，就可以立即将其包含在模型中(如清单1所示)。

此外，这种将所有张量推入中心存储的方法在反向传播步骤中将变得很方便，我将在下面进行讨论。

误差计算

在向前传播结束时，我们需要计算我们的预测和目标之间的误差。

下面是均方误差计算的代码，通常用于回归，而交叉熵损失通常用于分类。

清单4 -均方误差和交叉熵损失函数

fn mse(target: &Tensor, pred: &Tensor) -> Tensor {

   (target - pred).square().mean(Kind::Float)

}



fn cross_entropy (target: &Tensor, pred: &Tensor) -> Tensor {

   let loss = pred.log_softmax(-1, Kind::Float).nll_loss(target);

   loss

}

这样向前传播就完成了，我们现在开始向后传播。

反向传播

在反向传播中，我们需要使用梯度来更新模型的参数，其中每个梯度都是损失函数对于每个相应参数的导数。第一步，我们得到梯度:

公式2

其中m '表示小批量的大小。对于每个小批量，参数更新如下:

公式3

其中epsilon是学习率。

这就是我使用来自LibTorch的Autograd功能来获取梯度的地方。在PyTorch中，我们通常在损失上应用backward方法来计算导数，然后调用优化器的step函数将梯度应用到模型参数上。在这里，同样的过程发生，唯一的区别是我们不能直接应用step函数来应用梯度，因为我们没有从nn.Module类继承并使用PyTorch优化器，这是我们在Python中通常所做的。因此，我们需要自己处理step部分。

在下面的代码片段中（清单5），我们展示了我们的张量内存实现，它也适用于梯度步骤功能。张量存储被实现为一个结构体，有两个字段：size保存当前存储的张量数量，而values是一个张量的向量。在实现块中，new方法处理存储的初始化，push、new_push和get方法处理张量的传递（后两个在前面的线性层中使用过）。

清单5 -张量存储-内存

struct Memory {

    size: usize,

    values: Vec,

}



impl Memory {



    fn new() -> Self {

        let v = Vec::new();

        Self {size: 0,

            values: v}

}



    fn push (&mut self, value: Tensor) -> usize {

        self.values.push(value);

        self.size += 1;

        self.size-1

}



    fn new_push (&mut self, size: &[i64], requires_grad: bool) -> usize {

        let t = Tensor::randn(size, (Kind::Float, Device::Cpu)).requires_grad_(requires_grad);

        self.push(t)

}



    fn get (&self, addr: &usize) -> &Tensor {

        &self.values[*addr]

}



    fn apply_grads_sgd(&mut self, learning_rate: f32) {

        let mut g = Tensor::new();

        self.values

        .iter_mut()

        .for_each(|t| {

            if t.requires_grad() {

                g = t.grad();

                t.set_data(&(t.data() - learning_rate*&g));

                t.zero_grad();

            }

        });

    }



    fn apply_grads_sgd_momentum(&mut self, learning_rate: f32) {

        let mut g: Tensor = Tensor::new();

        let mut velocity: Vec= Tensor::zeros(&[self.size as i64], (Kind::Float, Device::Cpu)).split(1, 0);

        let mut vcounter = 0;

        const BETA:f32 = 0.9;



        self.values

        .iter_mut()

        .for_each(|t| {

            if t.requires_grad() {

                g = t.grad();

                velocity[vcounter] = BETA * &velocity[vcounter] + (1.0 - BETA) * &g;

                t.set_data(&(t.data() - learning_rate * &velocity[vcounter]));

                t.zero_grad();

            }

            vcounter += 1;

        });

    }

}

上述代码中的后两种方法分别实现了基本梯度下降和动量的梯度下降。这些方法假设已经调用了生成梯度的方向步骤，所以我们在处理的是PyTorch中的步骤函数调用。

该过程包括循环遍历存储中的每个张量，使用grad方法获得计算的梯度，然后通过调用set_data方法应用参数更新规则。可以很容易地引入其他方法并实现其他算法，如Rmsprop和Adam。

训练循环

在训练循环中，我们将前面讨论的所有内容整合到我们的学习过程中。像往常一样，我们对每个epoch应用一个循环，然后对每个minibatch进行循环，对每个minibatch进行前向传播，计算误差，在错误上调用反向方法来生成梯度，然后应用梯度(清单6)。

清单6 -训练循环。

fn train(mem: &mut Memory, x: &Tensor, y: &Tensor, model: &dyn Compute, epochs: i64, batch_size: i64, errfunc: F, learning_rate: f32)

    where F: Fn(&Tensor, &Tensor)-> Tensor

        {

        let mut error = Tensor::from(0.0);

        let mut batch_error = Tensor::from(0.0);

        let mut pred = Tensor::from(0.0);

        for epoch in 0..epochs {

            batch_error = Tensor::from(0.0);

            for (batchx, batchy) in get_batches(&x, &y, batch_size, true) {

                pred = model.forward(mem, &batchx);

                error = errfunc(&batchy, &pred);

                batch_error += error.detach();

                error.backward();

                mem.apply_grads_sgd_momentum(learning_rate);

            }

            println!("Epoch: {:?} Error: {:?}", epoch, batch_error/batch_size);

        }

}

在PyTorch中，我们有Dataset和Dataloader类来处理数据的小批处理机制。

下面的Rust函数(清单7)接受对完整数据集的引用，然后返回一个迭代器，该迭代器允许训练函数(清单6)遍历小批量。

清单7 -小批处理。

fn get_batches(x: &Tensor, y: &Tensor, batch_size: i64, shuffle: bool) -> impl Iterator {

   let num_rows = x.size()[0];

   let num_batches = (num_rows + batch_size - 1) / batch_size;



   let indices = if shuffle {

      Tensor::randperm(num_rows as i64, (Kind::Int64, Device::Cpu))

   } else

    {

       let rng = (0..num_rows).collect::>();

       Tensor::from_slice(&rng)

   };

   let x = x.index_select(0, &indices);

   let y = y.index_select(0, &indices);



   (0..num_batches).map(move |i| {

      let start = i * batch_size;

      let end = (start + batch_size).min(num_rows);

      let batchx: Tensor = x.narrow(0, start, end - start);

      let batchy: Tensor = y.narrow(0, start, end - start);

      (batchx, batchy)

   })

}

最终辅助函数

运行完整代码所需的最后两个函数是两个辅助函数(清单8)。

第一个函数从我命名为data的目录加载数据集(你必须首先下载Mnist数据集)。

第二个函数计算模型的准确度，接受对目标和预测的引用作为参数。

清单8 -最后两个helper函数。

fn load_mnist() -> (Tensor, Tensor) {

   let m = vision::mnist::load_dir("data").unwrap();

   let x = m.train_images;

   let y = m.train_labels;

   (x, y)

}



fn accuracy(target: &Tensor, pred: &Tensor) -> f64 {

   let yhat = pred.argmax(1,true).squeeze();

   let eq = target.eq_tensor(&yhat);

   let accuracy: f64 = (eq.sum(Kind::Int64) / target.size()[0]).double_value(&[]).into();

   accuracy

}

你唯一需要的导入是:

清单9 -所需的导入

use std::{collections::HashMap};

use tch::{Tensor, Kind, Device, vision, Scalar};

在运行代码之前，你还需要从PyTorch网站下载LibTorch c++库。

总结

用Rust开发上面的代码一开始肯定要花更多的时间，特别是对于Rust新手，但是，一旦构建了所有的库组件和管道代码，只需要测试/创建新模型(如清单1所示)，那么在我看来，它就会变得和用Python一样简单。

并且，训练速度的提高是不容忽视的——它确实可以节省长时间的训练，尤其是随着 ML 模型的复杂性、更大的数据集或巨大的迭代学习过程的增加，比如强化学习。

 

来源：https://medium.com/better-programming/boosting-machine-learning-performance-with-rust-aab1f3ae1424