在本篇文章中，作者將在 Rust 上移植一個簡單的神經網絡實現(xiàn)。他的目標是探索 Rust 中的數(shù)據(jù)科學工作流在性能和工程學上的表現(xiàn)。
Python 實現(xiàn)
第一章描述了一個非常簡單的單層神經網絡。這個神經網絡可以使用基于隨機梯度下降的機器學習算法，對來自于 MNIST 數(shù)據(jù)集的手寫數(shù)字進行分類。這聽起來挺復雜，這些東西也確實在上世紀 80 年代中期是最先進的，但是實際上，這全部是由一段 150 行的 Python 代碼做出來的，而且這些代碼有很多人評論過。
如果你已經知道了這一節(jié)的內容(神經網絡基礎知識)，那么我建議你可以跳過去，當然如果想再復習一下神經網絡的基礎知識也是可以看這一節(jié)的。或者不要只關注代碼，沒有必要特別細致的理解代碼為什么以現(xiàn)有的方式運行，而應該關注 Python 方法和 Rust 方法的不同。
代碼中的基礎數(shù)據(jù)容器是一個 Network 類，它表示一個對層數(shù)和每層網絡數(shù)量可控制的神經網絡。在 Network 類的內部，用 2D NumPy 數(shù)組組成的列表表示類的數(shù)據(jù)。網絡各層用一個二維的權重數(shù)組和一個一維的偏差數(shù)組來表示，這些數(shù)組被包含在叫做 biases 和 weights 的 Network 類的屬性中。這些都是二維數(shù)組列表。biases 屬性是列向量，但是會利用虛擬維度被存儲成二維的數(shù)組。Network 的構造函數(shù)如下：
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:])]
在這個簡單的實現(xiàn)中，屬性 biases 和 weights 通過描述標準正態(tài)分布來初始化。正態(tài)分布的均值為 0，標準差為 1。我們還可以清楚地看到，biases 是如何被初始化為列向量的。
Network 類公開了兩個可以被用戶直接調用的方法。第一個方法是 evaluate 方法，這個方法可以通過網絡嘗試識別一系列測試圖片中的數(shù)字，然后基于先驗已知的正確結果，對識別的結果進行打分。第二個方法是 SGD 方法，這個方法可以通過遍歷一組圖片來執(zhí)行隨機梯度下降的過程。這個過程包括：將整組圖像分解成小的類別，基于各個小類別圖像來更新網絡狀態(tài)，更新用戶指定的學習率，eta，以及在用戶隨機指定數(shù)量的一系列小類別圖像上重新運行以上的訓練步驟。每組小分類圖像和網絡的更新的核心算法如下代碼所示：
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 函數(shù)中實現(xiàn))積累成本函數(shù)的梯度估計。當程序跑完了小分類圖像集合，會根據(jù)估計梯度調整權重和偏差。因為我們希望得到小分類集合中所有估計的平均梯度，所以要更新的數(shù)據(jù)包括分母中的 len(mini_batch) 值。我們還可以通過調整學習速率和 eta 值來控制權重和偏差的更新速度，從而可以全局調整每個小分類集合更新的大小。backprop 函數(shù)開始于給定輸入圖像的網絡預期，然后通過網絡向后面運行，以在網絡中通過層來傳播這些錯誤，最后計算出神經網絡的成本梯度。這需要大量的數(shù)據(jù)調整，這也是我在移植到 Rust 時，花費了大量時間的地方。不過我認為我深入這塊花費了過于長的時間。如果你想更深入的了解細節(jié)，請看這本書的第 2 章。
Rust 實現(xiàn)
第一步要搞清楚如何加載數(shù)據(jù)。這一步太繁瑣了，我寫了一篇專門的文章來介紹。按照順序，我必須弄清楚如何用 Rust 實現(xiàn) Python 的 Network 類。最終我決定使用結構體(struct)：
use ndarray::Array2;
#[derive(Debug)]
struct Network {
num_layers: usize,
sizes: Vec
和 Python 實現(xiàn)大致一樣，依據(jù)每層神經網絡的數(shù)量初始化結構體。
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
有一個和 Python 實現(xiàn)的不同點。我們在 Python 中，使用 numpy.random.randn 來初始化權重和偏差，而在 Rust 中，我們使用 ndarray::Array::random 函數(shù)接受一個 rand::distribution::Distributionas 類型的參數(shù)和一個其他參數(shù)，并允許選擇任意分布，來完成初始化。在這種情況下，我們使用了 rand::distributions::StandardNormal 做分布。值得注意的是，這使用了在三個不同包中定義的接口：其中兩個接口，ndarray 自身和 ndarray-rand 由 ndarray 的作者維護，剩下一個由其它的開發(fā)者維護。
這些統(tǒng)一類包的優(yōu)勢
原則上，一個好處是，隨機數(shù)生成不會在 ndarray 代碼庫中被單獨隔離，如果新的隨機數(shù)分布或者功能被加到了 rand 中，ndarray 和在 Rust 系統(tǒng)中的類都可以同等的使用。另一方面，需要為了各種裝箱操作，在不同文件之間進行引用，而不是可以在一個集中的地方進行查看，這樣增加了一些學習成本。我也有個特殊情況，也算是運氣不好，rand 發(fā)布了改變公共 api 的新版本時，我正在開發(fā)這個工程。這導致了，依賴于 0.6 版本的 rand 的 ndarray-rand 和依賴于 0.7 版本的 rand 的我的工程，之間的不一致。
我了解到 cargo 和 rust 的構建系統(tǒng)可以很好地處理這類問題，但至少在這種情況下，我遇到了一個令人困惑的錯誤信息。這則錯誤信息是關于我的隨機分布如何不滿足 Distribution(分布)特征的。然而，我的分布式正確的，它的隨機分布特征滿足了 0.7 版本的 rand，而不是 ndarray-rand 依賴的 rand0.6 版本。但是因為裝箱的版本信息不出現(xiàn)在錯誤信息當中，所以會讓人非常的困惑。我最后提交了一個 issue 。我發(fā)現(xiàn)對于 Rust 語言來說，來自包裝箱的各種不一致接口的讓人困惑的錯誤信息，是一個長期存在的問題。希望在未來，Rust 可以產生更多有用的錯誤信息。
最后，作為一個新用戶，這種關注點的分離給我(的理解上)帶來了許多阻力。在 Python 中，我可以簡單的做 import numpy 操作便可以完成導入。我認為 NumPy 在完全單片化的方向上走的太遠了。它最初被編寫的時候，用 C 語言擴展的 Python 代碼，在打包和發(fā)布上，要比現(xiàn)在困難的多。我認為在一個極端方向上走的太遠，會讓一個語言或者工具的生態(tài)系統(tǒng)變得很難學習。
類型和所有權
下一步我將詳細介紹 update_mini_batch 的 Rust 版本：
impl Network {
fn update_mini_batch(
&mut self,
training_data: &[MnistImage],
mini_batch_indices: &[usize],
eta: f64,
) {
let mut nabla_b: Vec
該函數(shù)使用了我定義的兩個簡短的輔助函數(shù)，使得代碼更簡潔了一些：
fn to_tuple(inp: &[usize]) -> (usize, usize) { match inp { [a, b] => (*a, *b), _ => panic!(), } }
fn zero_vec_like(inp: &[Array2
和 Python 的實現(xiàn)版本相比較，調用 update_mini_batch 的方式有些不同。沒有直接傳遞對象列表，反而傳遞了對完整集合中全套訓練數(shù)據(jù)和一份索引的引用。這樣，更容易理解沒有觸發(fā)器的借用檢查器。
創(chuàng)建 nabla_b 和 nabla_win zero_vec_like 和我們在 Python 中使用的列表解析非常類似。有一個挫折讓我有些沮喪，因為如果我嘗試使用 Array2::zeros 創(chuàng)建一個用 0 填充的數(shù)組，并將它傳遞給一個特定形狀的 slice 或 Vec，我就會得到一個 ArrayD 的實例。為了獲得 Array2 對象(顯然這是一個二維數(shù)組，而不是一個通用的 D 緯數(shù)組)，我需要向 Array::zeros 傳遞一個元素。然而，由于 ndarray::shape 返回一個切片(slice)，我需要使用 to_tuple 函數(shù)，將這個切片轉換為一個元組。這些事情在 Python 中可以被隱藏，但是在 Rust 中，切片(slice)和元組(tuple)之間的不同(造成的影響)變得非常大，就和在這個 API 中的情況一樣。
通過反向傳播對估計的權重和偏差進行更新的代碼具有和 Python 實現(xiàn)的版本非常相似的結構。我們在小分類中訓練每個示例圖像，并獲得二次成本梯度的估計值作為偏差和權重相關的一個函數(shù)：
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; }
等到我們完成了小分類的處理，我們就會根據(jù)學習率更新權重和偏差。
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); }
這個例子說明了，Rust 對數(shù)組數(shù)據(jù)在工程上的處理和 Python 相比有區(qū)別的。首先，我們不用浮點數(shù) eta/nbatch 乘以數(shù)組，而是使用 Array::mapv，并定義一個內連閉包，以便在整個數(shù)組上以矢量化的方式進行映射。在 Python 中，由于方法調用比較慢，所以這些事情不會處理的太快。而在 Rust 中，則不會出現(xiàn)這種情況。當我們減去時，還需要借用帶 & 符號 mapv 的返回值，以免我們在迭代它時消耗數(shù)組數(shù)據(jù)。在 Rust 中，需要仔細考慮函數(shù)是否會消耗數(shù)據(jù)或者引入引用，這導致在概念上，用 Rust 編寫這種代碼，比在 Python 中要求更多。另一方面，我對我的代碼的正確性并且能夠編譯通過，有了更高的信心。我不確定的是，我寫這段代碼很費力的原因，是因為 Rust 真的更難寫，還是因為我在 Python 和 Rust 上經驗的不同。
用 Rust 重寫這些代碼，然后一切都會好起來
在這里，我留下了一些東西，比我開始使用的未經優(yōu)化的 Python 版本代碼更快。然而，相比 10 倍或是更快的速度，人們可能更期望從像 Python 這樣的動態(tài)解釋性語言轉變?yōu)橄?Rust 這樣的編譯性能導向語言，并且我也只觀察到了 2 倍的加速。可以去理解一下我為什么要測量 Rust 語言的性能表現(xiàn)。幸運的是，這里有一個非常方便的項目，可以為 Rust 工程生成火焰圖： flamegraph 。這里添加了一個 flamegraph 的子命令 cargo，因此只需要在包中執(zhí)行 cargo flamegraph 即可運行代碼，就會編寫出一個可以在瀏覽器中執(zhí)行的火焰圖 svg 文件(原圖為可交互的 svg 腳本，如果希望嘗試，可以查看原網頁)。
?
?
如果你之前還沒有看過一個火焰圖，(我解釋下)，在例程中發(fā)生的程序運行時間與與該例程的條形寬度成正比。主函數(shù)位于圖形的底部，主函數(shù)調用的函數(shù)在圖形的頂部。這樣你就可以簡單查看哪些函數(shù)占用了程序中最多的時間。圖中非常寬的東西代表了花費最多時間的地方。在調用棧中非常高和寬的函數(shù)，在代碼上花費了大量的時間。看一下上面的火焰圖，我們可以發(fā)現(xiàn)一般的時間都花費在了像名字叫 dgemm_kernel_HASWELL 的這類函數(shù)身上，這類函數(shù)是 OpenBLAS 的線性代數(shù)類庫，剩下的時間，花費在 update_mini_batch 的數(shù)組和分配數(shù)組之間的添加上。我程序的其它所有部分，對運行時間的貢獻可以忽略不計。
如果我們?yōu)?Python 代碼做一個類似的火焰圖，我們會發(fā)現(xiàn)一個相似的情況：大部分時間花費在了做線性代數(shù)函數(shù)上去(在 np.dot 反向傳播例程中調用的地方)。因此，由于不管是 Rust 還是 Python 花費的時間大部分都在數(shù)值性的線性代數(shù)庫上，我們就不能夠希望得到一個 10 倍加速的結果。
實際情況比這更糟糕。這本書中的一個練習是重寫了使用向量化矩陣乘法的 Python 代碼。在這個方法中，每個小分類中的所有樣例的反向傳播發(fā)生在單組矢量化矩陣乘法運算中。這需要能夠在 3 維和 2 維數(shù)組之間進行矩陣乘法。由于每個矩陣乘法運算使用的數(shù)據(jù)量大于非向量化的情況，OpenBLAS 可以更有效地使用 CPU 緩存和寄存器，基本上可以更好地利用我筆記本電腦上的可用 CPU 資源。重寫的 Python 版本要比 Rust 版本更快，又快了大約兩倍左右。
理論上，可以對 Rust 代碼進行相同的優(yōu)化。但是對于高于 2 維(的矩陣)的情況，ndarraycrate還不支持矩陣乘法。也可以使用像 rayon 這樣的庫在小批量更新上使用線程并行化。我在我的筆記本上嘗試這個(并行化)沒有看到任何的加速，但是可能在具有更多 CPU 線程的更強大的機器上會有作用。我也可以嘗試使用一個不同的線性代數(shù)函數(shù)實現(xiàn)，例如，有 TensorFlow 和 Torch 的 Rust 構建，但是在這種情況下，我覺得我也可以使用那些庫的 Python 構建。
Rust 是否適合數(shù)據(jù)科學的工作流?
現(xiàn)在我不得不說，答案是”未知“。在未來，當我需要編寫具有小依賴性的低級別優(yōu)化代碼時，我肯定會使用 Rust。但是，如果把 Rust 作為 Python 和 C++ 的完全替代品，還需要一個更穩(wěn)定和完善的類庫生態(tài)系統(tǒng)。
