在上一篇文章(第1部分)中，做了用Rust从头开发机器学习框架的实验。实验的主要目的是衡量通过将Rust与PyTorch结合使用而不是Python等同物来获得的模型训练速度改进。在本文中，我将继续在此基础上进行构建，主要目标是能够定义和训练卷积神经网络(CNN)。与上一篇文章一样，继续使用Tch-rs Rust crate作为PyTorch c++库LibTorch的包装器，主要是访问张量、线性代数和自grad函数，其余部分从头开发。

本文的最终结果是允许我们在Rust中定义卷积神经网络(CNN)，如下所示:

清单1 -定义CNN模型。

struct MyModel {

    l1: Conv2d,

    l2: Conv2d,

    l3: Linear,

    l4: Linear,

}



impl MyModel {

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

       let l1 = Conv2d::new(mem, 5, 1, 10, 1);

       let l2 = Conv2d::new(mem, 5, 10, 20, 1);

       let l3 = Linear::new(mem, 320, 64);

       let l4 = Linear::new(mem, 64, 10);

       Self {

           l1: l1,

           l2: l2,

           l3: l3,

           l4: l4,

       }

    }

}



impl Compute for MyModel {

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

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

       o = o.max_pool2d_default(2);

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

       o = o.max_pool2d_default(2);

       o = o.flat_view();

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

       o = o.relu();

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

       o

    }

}

然后按如下方式实例化和训练：

清单2 -训练CNN模型。

fn main() {

   let (mut x, y) = load_mnist();

   x = x / 250.0;

   x = x.view([-1, 1, 28, 28]);



   let mut m = Memory::new();

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

   train(&mut m, &x, &y, &mymodel, 20, 512, cross_entropy, 0.0001);

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

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

}

为了尽量与Python等效代码保持相似，上面的清单1对于Python-pytorch用户应该非常直观。在MyModel结构中，我们现在能够添加Conv2D图层，然后在相关的new函数中初始化它们。在Compute trait实现中，forward函数被定义并通过所有层获取输入，包括中间的MaxPooling函数。在main函数(清单2)中，与上一篇文章类似，我们正在训练模型并将其应用于Mnist数据集。

接下来，我将描述能够以这种方式定义和训练CNN的原理。

试验卷积核

卷积网络的独特特征是，在某些层(至少一个层)中，我们应用卷积而不是一般的矩阵乘法。卷积操作的目的是利用核从输入图像中提取感兴趣的某些特征。核是一个矩阵，它在图像(输入)的子部分之间滑动并相乘，使得输出是以某种理想的方式对输入进行变换(见下图)。



在二维情况下，我们使用二维图像I作为输入，我们通常也使用二维核K，导致以下卷积计算:

公式1

从公式1可以推断，从计算的角度来看，应用卷积的算法非常昂贵，因为有大量的循环和矩阵乘法。更糟糕的是，对于网络中的每个卷积层和每个训练示例/批次，这必须重复多次计算。因此，在扩展第1部分的库来处理CNN之前，第一步是研究一种计算卷积的有效方法。

寻找计算卷积的有效方法是一个研究得很好的问题。在研究了不同的选项后，其中包括一些纯Rust版本，但需要从PyTorch张量进行中间数据转换，所以选择使用LibTorch c++卷积函数。为了实验这个函数，我想创建一个小的玩具程序，它获取彩色图像，将其转换为灰度，然后应用一些已知的内核来执行边缘检测。

首先要求Microsoft Bing chat 为我生成图像。一旦我对图像感到满意，我想首先使用高斯核应用卷积函数，然后是拉普拉斯核。

高斯核

拉普拉斯核

内核使用LibTorch c++方法conv2d进行应用，该方法通过Tch-rs公开为：

清单3 - 公开Tch-rs的LibTorch conv2d方法

pub fn conv2d>(

   &self,

   weight: &Tensor,

   bias: Option,

   stride: impl IntList,

   padding: impl IntList,

   dilation: impl IntList,

   groups: i64

 ) -> Tensor

我最后的玩具程序如下所示：

清单4 -获取图像并应用卷积操作进行边缘检测

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



fn rgb_to_grayscale(tensor: &Tensor) -> Tensor {

   let red_channel = tensor.get(0);

   let green_channel = tensor.get(1);

   let blue_channel = tensor.get(2);



   // 使用练度公式计算灰度张量

   let grayscale = (red_channel * 0.2989) + (green_channel * 0.5870) + (blue_channel * 0.1140);

   grayscale.unsqueeze(0)

}



fn main() {

   let mut img = image::load("mypic.jpg").expect("Failed to open image");

   img = rgb_to_grayscale(&img).reshape(&[1,1,1024,1024]);

   let bias: Tensor = Tensor::full(&[1], 0.0, (Kind::Float, Device::Cpu));



   // 定义并应用高斯核

   let mut k1 = [-1.0, 0.0, 1.0, -2.0, 0.0, 2.0, -1.0, 0.0, 1.0];

   for element in k1.iter_mut() {

       *element /= 16.0;

   }

   let kernel1 = Tensor::from_slice(&k1)

                       .reshape(&[1,1,3,3])

                       .to_kind(Kind::Float);

   img = img.conv2d(&kernel1, Some(&bias), &[1], &[0], &[1], 1);

 

   // 定义并应用拉普拉斯核

   let k2 = [0.0, 1.0, 0.0, 1.0, -4.0, 1.0, 0.0, 1.0, 0.0];

   let kernel2 = Tensor::from_slice(&k2)

                       .reshape(&[1,1,3,3])

                       .to_kind(Kind::Float);

   img = img.conv2d(&kernel2, Some(&bias), &[1], &[0], &[1], 1);

   image::save(&img, "filtered.jpg");

}

运行结果如下:



在这个玩具程序中，我们应用我们选择的内核进行卷积，将原始图像转换为边缘(我们想要的特征)。下面，将描述如何将这个想法融入到CNN中，主要区别在于核矩阵的值是由网络在训练期间选择的——即网络自己通过调整核来决定从图像中选择哪些特征。

卷积层

CNN的典型结构包括许多卷积层，每个层后面都有一个子采样层(池化)，然后通常进入到完全连接的层。池化对参数减少有很大的贡献，因为它对输入数据进行了下采样。下图描绘了最早的CNN之一，名为LeNet-5。



在第1部分中，已经定义了一个包含全连接层的简单框架。同样的,我们现在需要做的是在我们的框架中注入一个卷积的定义层,以便在定义新的网络架构时（如清单1中）可用。另一件需要记住的事情是，在我们的玩具程序（清单4）中，我们将卷积核矩阵和偏置设置为固定值，但现在我们需要将它们定义为由训练算法训练的网络参数，因此我们需要跟踪它们的梯度并相应地进行更新。

新的卷积层Conv2d的定义如下：

清单5 -新的Conv2d层

pub struct Conv2d {

   params: HashMap,

}



impl Conv2d {

    pub fn new (mem: &mut Memory, kernel_size: i64, in_channel: i64, out_channel: i64, stride: i64) -> Self {

       let mut p = HashMap::new();

       p.insert("kernel".to_string(), mem.new_push(&[out_channel, in_channel, kernel_size, kernel_size], true));

       p.insert("bias".to_string(), mem.push(Tensor::full(&[out_channel], 0.0, (Kind::Float, Device::Cpu)).requires_grad_(true)));

       p.insert("stride".to_string(), mem.push(Tensor::from(stride as i64)));

       Self {

             params: p,

       }

    }

}



impl Compute for Conv2d {

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

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

        let stride: i64 = mem.get(self.params.get(&"stride".to_string()).unwrap()).int64_value(&[]);

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

        input.conv2d(&kernel, Some(bias), &[stride], 0, &[1], 1)

    }

}

在第1部分中的方法，该结构包含一个名为params的字段。params字段是HashMap类型的集合，其中Key的类型是String，它存储参数名称，value的类型是usize，它保存特定参数(它是PyTorch张量)在Memory中的位置，它反过来充当我们所有模型参数的存储。在卷积层的情况下，在我们的关联函数new中，我们在HashMap中插入两个参数“Kernel”和“bias”，它们设置为required_gradient标志True。我还注入了一个参数“Stride”，但这不是一个可训练的参数。

然后我们为卷积层实现Compute特性。这需要定义forward函数，该函数在训练过程的前向传递过程中被调用。在这个函数中，我们首先使用get方法从我们的张量存储中获得对核、偏差和跨幅张量的引用，然后调用我们的Conv2d函数(就像我们在玩具程序中所做的那样，但是在这种情况下，网络会告诉我们使用哪个核)。填充被硬编码为零值，但是如果你愿意，也可以将其作为类似Stride的参数添加。

这是第1部分中的小框架中需要添加的内容，以便能够定义和训练清单1 - 2所示的CNN。

亚当优化

Adam算法于2015年首次发表，基本结合了Momentum和RMSprop训练算法的思想。原论文中的算法如下:



在第1部分中，我们实现了张量内存，它也满足PyTorch等效梯度阶进函数(方法apply_grads_sgd和apply_grads_sgd_momentum)。因此，在Memory 结构实现中添加了一个新方法，使用Adam执行梯度更新:

清单6 -我们的Adam实现。

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

        let mut g = Tensor::new();

        const BETA:f32 = 0.9;



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

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

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

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

        let mut counter = 0;



        self.values

        .iter_mut()

        .for_each(|t| {

            if t.requires_grad() {

                g = t.grad();

                mom[counter] = BETA * &mom[counter] + (1.0 - BETA) * &g;

                velocity[counter] = BETA * &velocity[counter] + (1.0 - BETA) * (&g.pow(&Tensor::from(2)));

                mom_corr[counter] = &mom[counter] / (Tensor::from(1.0 - BETA).pow(&Tensor::from(2)));

                vel_corr[counter] = &velocity[counter] / (Tensor::from(1.0 - BETA).pow(&Tensor::from(2)));

                

                t.set_data(&(t.data() - learning_rate *(&mom_corr[counter]/(&velocity[counter].sqrt() + 0.0000001))));

                t.zero_grad();

            }

            counter += 1;

        });

}

结论

在多次运行测试之后，Rust训练的平均速度比Python提高了60%。虽然这是一个显着的改进，但这比FFN的情况下实现的要少(我在第1部分中的发现)。我将这种较低的改进归因于CNN中最昂贵的计算是卷积，并且如上所述，我选择使用LibTorch c++ conv2d函数，它与Python等效函数调用的函数相同。然而，话虽如此，将模型训练的时间减少一半以上仍然是不容忽视的——这通常意味着节省几个小时，甚至几天!

参考文献

1.Ian Goodfellow、Yoshua Bengio 和 Aaron Courville，《深度学习》，麻省理工学院出版社，2016 年。

2.Pavel Karas 和 David Svoboda，卷积高效计算算法，数字信号处理设计和架构，美国纽约州纽约市：IntechOpen，2013 年 1 月。

3.LeCun, Y.、Bottou, L.、Bengio, Y. 和 Haffner, P.，Gradient-based learning applied to document recognition， IEEE 86 会刊，2278–2324，1998。

4.Diederik P Kingma 和 Jimmy Ba，Adam：一种随机优化方法，载于第三届国际学习表示会议 (ICLR) 论文集，2015 年。

 

来源：https://medium.com/@vincevella/boosting-machine-learning-performance-with-rust-part-2-5e4c2d0adf19