Rust 也能实现神经网络？

作者 | Nathan J. Goldbaum

译者 | 弯月，责编 | 屠敏

出品 | CSDN（ID：CSDNnews）

以下为译文：

我在前一篇帖子（）中介绍了MNIST数据集（）以及分辨手写数字的问题。在这篇文章中，我将利用前一篇帖子中的代码，通过Rust实现一个简单的神经网络。我的目标是探索用Rust实现数据科学工作流程的性能以及人工效率。

Python的实现

我在前一篇帖子中描述了一个非常简单的单层神经网络，其可以利用基于随机梯度下降的学习算法对MNIST数据集中的手写数字进行分类。听起来有点复杂，但实际上只有150行Python代码，以及大量注释。

如果你想深入了解神经网络的基础知识，请仔细阅读我的前一篇帖子。而且请不要只关注代码，理解代码工作原理的细节并不是非常重要，你需要了解Python和Rust的实现差异。

在前一篇帖子中，Python代码的基本数据容器是一个Network类，它表示一个神经网络，其层数和每层神经元数可以自由控制。在内部，Network类由NumPy二维数组的列表表示。该网络的每一层都由一个表示权重的二维数组和一个表示偏差的一维数组组成，分别包含在Network类的属性weights和biases中。两者都是二维数组的列表。偏差是列向量，但仍然添加了一个无用的维度，以二维数组的形式存储。Network类的初始化程序如下所示：

class Network(object):
def __init__(self, sizes):
"""The list ``sizes`` contains the number of neurons in the
respective layers of the network. For example, if the list
was [2, 3, 1] then it would be a three-layer network, with the
first layer containing 2 neurons, the second layer 3 neurons,
and the third layer 1 neuron. The biases and weights for the
network are initialized randomly, using a Gaussian
distribution with mean 0, and variance 1. Note that the first
layer is assumed to be an input layer, and by convention we
won't set any biases for those neurons, since biases are only
ever used in computing the outputs from later layers."""
self.num_layers = len(sizes)
self.sizes = sizes
self.biases = [np.random.randn(y, 1) for y in sizes[1:]]
self.weights = [np.random.randn(y, x)
for x, y in zip(sizes[:-1], sizes[1:])]

在这个简单的实现中，权重和偏差的初始化呈标准正态分布——即均值为零，标准差为1的正态分布。我们可以看到，偏差明确地初始化为列向量。

这个Network类公开了两个用户可以直接调用的方法。第一个是evaluate方法，它要求网络尝试识别一组测试图像中的数字，然后根据已知的正确答案对结果进行评分。第二个是SGD方法，它通过迭代一组图像来运行随机梯度下降的学习过程，将整组图像分解成小批次，然后根据每一小批次的图像以及用户指定的学习速率eta更新该网络的状态；最后再根据用户指定的迭代次数，随机选择一组小批次图像，重新运行这个训练过程。该算法的核心（每一小批次图像处理以及神经网络的状态更新）代码如下所示：

def update_mini_batch(self, mini_batch, eta):
"""Update the network's weights and biases by applying
gradient descent using backpropagation to a single mini batch.
The ``mini_batch`` is a list of tuples ``(x, y)``, and ``eta``
is the learning rate."""
nabla_b = [np.zeros(b.shape) for b in self.biases]
nabla_w = [np.zeros(w.shape) for w in self.weights]
for x, y in mini_batch:
delta_nabla_b, delta_nabla_w = self.backprop(x, y)
nabla_b = [nb+dnb for nb, dnb in zip(nabla_b, delta_nabla_b)]
nabla_w = [nw+dnw for nw, dnw in zip(nabla_w, delta_nabla_w)]
self.weights = [w-(eta/len(mini_batch))*nw
for w, nw in zip(self.weights, nabla_w)]
self.biases = [b-(eta/len(mini_batch))*nb
for b, nb in zip(self.biases, nabla_b)]

我们可以针对小批次中的每个训练图像，通过反向传播（在backprop函数中实现）求出代价函数的梯度的估计值的总和。在处理完所有的小批次后，我们可以根据估计的梯度调整权重和偏差。更新时在分母中加入了len(mini_batch)，因为我们想要小批次中所有估计的平均梯度。我们还可以通过调整学习速率eta来控制权重和偏差的更新速度，eta可以在全局范围内调整每个小批次更新的大小。

backprop函数在计算该神经网络的代价函数的梯度时，首先从输入图像的正确输出开始，然后将错误反向传播至网络的各层。这需要大量的数据调整，在将代码移植到Rust时我在此花费了大量的时间，在此篇幅有限，我无法深入讲解，如果你想了解具体的详情，请参照这本书（）。

Rust的实现

首先，我们需要弄清楚如何加载数据。这个过程非常繁琐，所以我另写了一篇文章专门讨论（）。在这之后，下一步我们必须弄清楚如何用Rust表示Python代码中的Network类。最终我决定使用struct：

use ndarray::Array2;
#[derive(Debug)]
struct Network {
num_layers: usize,
sizes: Vec<usize>,
biases: Vec<Array2<f64>>,
weights: Vec<Array2<f64>>,
}

该结构的初始化与Python的实现大致相同：根据每层中的神经元数量进行初始化。

use rand::distributions::StandardNormal;
use ndarray::{Array, Array2};
use ndarray_rand::RandomExt;
impl Network {
fn new(sizes: &[usize]) -> Network {
let num_layers = sizes.len;
let mut biases: Vec<Array2<f64>> = Vec::new;
let mut weights: Vec<Array2<f64>> = Vec::new;
for i in 1..num_layers {
biases.push(Array::random((sizes[i], 1), StandardNormal));
weights.push(Array::random((sizes[i], sizes[i - 1]), StandardNormal));
}
Network {
num_layers: num_layers,
sizes: sizes.to_owned,
biases: biases,
weights: weights,
}
}
}

有一点区别在于，在Python中我们使用numpy.random.randn初始化偏差和权重，而在Rust中我们使用ndarray::Array::random函数，并以rand::distribution::Distribution为参数，允许选择任意的分布。在上述代码中，我们使用了rand::distributions::StandardNormal分布。注意，我们使用了三个不同的包中定义的接口，其中两个ndarray本身和ndarray-rand由ndarray作者维护，另一个rand则由其他开发人员维护。

整体式包的优点

原则上，最好不要将随机数生成器放到ndarray代码库中，这样当rand函数支持新的随机分布时，ndarray以及Rust生态系统中所有需要随机数的包都会受益。另一方面，这确实会增加一些认知开销，因为没有集中的位置，查阅文档时需要参考多个包的文档。我的情况有点特殊，我没想到做这个项目的时候，恰逢rand发布改变了其公共API的版本。导致ndarray-rand（依赖于rand版本0.6）和我的项目所依赖的版本0.7之间产生了不兼容性。

我听说cargo和Rust的构建系统可以很好地处理这类问题，但至少我遇到了一个非常令人困惑的错误信息：我传入的随机数分布不能满足Distribution这个trait的要求。虽然这话不假——它符合0.7版本的rand，但不符合ndarray-rand要求的0.6版本的rand，但这依然非常令人费解，因为错误信息中没有给出各种包的版本号。最后我报告了这个问题。我发现这些有关API版本不兼容的错误消息是Rust语言长期存在的一个问题。希望将来Rust可以显示更多有用的错误信息。

最后，这种关注点的分离给我这个新用户带来了很大困难。在Python中，我可以简单通过import numpy完成。我确实认为NumPy在整体式上走得太远了（当时打包和分发带有C扩展的Python代码与现在相比太难了），但我也认为在另一个极端上渐行渐远，会导致语言或生态系统的学习难度增大。

类型和所有权

下面我将详细介绍一下Rust版本的update_mini_batch：

impl Network {
fn update_mini_batch(
&mut self,
training_data: &[MnistImage],
mini_batch_indices: &[usize],
eta: f64,
) {
let mut nabla_b: Vec<Array2<f64>> = zero_vec_like(&self.biases);
let mut nabla_w: Vec<Array2<f64>> = zero_vec_like(&self.weights);
for i in mini_batch_indices {
let (delta_nabla_b, delta_nabla_w) = self.backprop(&training_data[*i]);
for (nb, dnb) in nabla_b.iter_mut.zip(delta_nabla_b.iter) {
*nb += dnb;
}
for (nw, dnw) in nabla_w.iter_mut.zip(delta_nabla_w.iter) {
*nw += dnw;
}
}
let nbatch = mini_batch_indices.len as f64;
for (w, nw) in self.weights.iter_mut.zip(nabla_w.iter) {
*w -= &nw.mapv(|x| x * eta / nbatch);
}
for (b, nb) in self.biases.iter_mut.zip(nabla_b.iter) {
*b -= &nb.mapv(|x| x * eta / nbatch);
}
}
}

该函数使用了我定义的两个辅助函数，因此更为简洁：

fn to_tuple(inp: &[usize]) -> (usize, usize) {
match inp {
[a, b] => (*a, *b),
_ => panic!,
}
}
fn zero_vec_like(inp: &[Array2<f64>]) -> Vec<Array2<f64>> {
inp.iter
.map(|x| Array2::zeros(to_tuple(x.shape)))
.collect
}

与Python实现相比，调用update_mini_batch的接口有点不同。这里，我们没有直接传递对象列表，而是传递了整套训练数据的引用以及数据集中的索引的切片。由于这种做法不会触发借用检查，因此更容易理解。

在zero_vec_like中创建nabla_b和nabla_w与我们在Python中使用的列表非常相似。其中有一个波折让我有些沮丧，本来我想设法使用Array2::zeros创建一个初始化为零的数组，并将其传递给图像的切片或Vec，这样我就可以得到一个ArrayD实例。如果想获得一个Array2（显然这是一个二维数组，而不是一个通用的D维数组），我需要将一个元组传递给Array::zeros。然而，由于ndarray::shape会返回一个切片，我需要通过to_tuple函数手动将切片转换为元组。这种情况在Python很容易处理，但在Rust中，元组和切片之间的差异非常重要，就像在这个API中一样。

利用反向传播估计权重和偏差更新的代码与python的实现结构非常相似。我们分批训练每个示例图像，并获得二次成本梯度的估计值作为偏差和权重的函数：

let (delta_nabla_b, delta_nabla_w) = self.backprop(&training_data[*i]);

然后累加这些估计值：

for (nb, dnb) in nabla_b.iter_mut.zip(delta_nabla_b.iter) {
*nb += dnb;
}
for (nw, dnw) in nabla_w.iter_mut.zip(delta_nabla_w.iter) {
*nw += dnw;
}

在处理完小批次后，我们根据学习速率调整权重和偏差：

let nbatch = mini_batch_indices.len as f64;
for (w, nw) in self.weights.iter_mut.zip(nabla_w.iter) {
*w -= &nw.mapv(|x| x * eta / nbatch);
}
for (b, nb) in self.biases.iter_mut.zip(nabla_b.iter) {
*b -= &nb.mapv(|x| x * eta / nbatch);
}

这个例子说明与Python相比，在Rust中使用数组数据所付出的人力有非常大的区别。首先，我们没有让这个数组乘以浮点数eta / nbatch，而是使用了Array::mapv，并定义了一个闭包，以矢量化的方式映射了整个数组。这种做法在Python中会很慢，因为函数调用非常慢。然而，在Rust中没有太大的区别。在做减法时，我们还需要通过&借用mapv的返回值，以免在迭代时消耗数组数据。在编写Rust代码时需要仔细考虑函数是否消耗数据或引用，因此在编写类似于Python的代码时，Rust的要求更高。另一方面，我更加确信我的代码在编译时是正确的。我不确定这段代码是否有必要，因为Rust真的很难写，可能是因为我的Rust编程经验远不及Python。

用Rust重新编写，一切都会好起来

到此为止，我用Rust编写的代码运行速度超过了我最初编写的未经优化的Python代码。然而，从Python这样的动态解释语言过渡到Rust这样的性能优先的编译语言，应该能达到10倍或更高性能，然而我只观察到大约2倍的提升。我该如何测量Rust代码的性能？幸运的是，有一个非常优秀的项目flamegraph（）可以很容易地为Rust项目生成火焰图。这个工具为cargo添加了一个flamegraph子命令，因此你只需运行cargo flamegraph，就可以运行代码，然后写一个flamegraph的svg文件，就可以通过Web浏览器观测。

可能你以前从未见过火焰图，因此在此简单地说明一下，例程中程序的运行时间比例与该例程的条形宽度成正比。主函数位于图形的底部，主函数调用的函数堆叠在上面。你可以通过这个图形简单地了解哪些函数在程序中占用的时间最多——图中非常“宽”的函数都在运行中占用了大量时间，而非常高且宽的函数栈都代表其包含非常深入的栈调用，其代码的运行占用了大量时间。通过以上火焰图，我们可以看到我的程序大约一半的时间都花在了dgemm_kernel_HASWELL等函数上，这些是OpenBLAS线性代数库中的函数。其余的时间都花在了`update_mini_batch和分配数组中等数组操作上，而程序中其他部分的运行时间可以忽略不计。

如果我们为Python代码制作了一个类似的火焰图，则也会看到一个类似的模式——大部分时间花在线性代数上（在反向传播例程中调用np.dot）。因此，由于Rust或Python中的大部分时间都花在数值线性代数库中，所以我们永远也无法得到10倍的提速。

实际情况可能比这更糟。上述我提到的书中有一个练习是使用向量化矩阵乘法重写Python代码。在这个方法中，每个小批次中所有图像的反向传播都需要通过一组矢量化矩阵乘法运算完成。这需要在二维和三维数组间运行矩阵乘法。由于每个矩阵乘法运算使用的数据量大于非向量化的情况，因此OpenBLAS能够更有效地使用CPU缓存和寄存器，最终可以更好地利用我的笔记本电脑上的CPU资源。重写的Python版本比Rust版本更快，但也只有大约两倍左右。

原则上，我们可以用相同的方式优化Rust代码，但是ndarray包还不支持高于二维的矩阵乘法。我们也可以利用rayon等库实现小批次更新线程的并行化。我在自己的笔记本电脑上试了试，并没有看到任何提速，但可能更强大的机器有更多CPU线程。我还尝试了使用使用不同的低级线性代数实现，例如，利用Rust版的tensorflow和torch，但当时我觉得我完全可以利用Python版的这些库。

Rust是否适合数据科学工作流程？

目前，我不得不说答案是“尚未”。如果我需要编写能够将依赖性降到最低的、经过优化的低级代码，那么我肯定会使用Rust。然而，要想利用Rust完全取代Python或C++，那么我们尚需要等待更稳定和更完善的包生态系统。

原文：

本文为 CSDN 翻译，转载请注明来源出处。

【End】