Rust 能取代 Python,更好的實現神經網絡?

Rust 也能實現神經網絡?

Rust 能取代 Python,更好的实现神经网络?

作者 | Nathan J. Goldbaum

出品 | CSDN(ID:CSDNnews)

我在前一篇帖子(http://neuralnetworksanddeeplearning.com/chap1.html)中介紹了MNIST數據集(http://yann.lecun.com/exdb/mnist/)以及分辨手寫數字的問題。在這篇文章中,我將利用前一篇帖子中的代碼,通過Rust實現一個簡單的神經網絡。我的目標是探索用Rust實現數據科學工作流程的性能以及人工效率。

Rust 能取代 Python,更好的实现神经网络?

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時我在此花費了大量的時間,在此篇幅有限,我無法深入講解,如果你想了解具體的詳情,請參照這本書(http://neuralnetworksanddeeplearning.com/chap2.html)。

Rust 能取代 Python,更好的实现神经网络?

Rust的實現

首先,我們需要弄清楚如何加載數據。這個過程非常繁瑣,所以我另寫了一篇文章專門討論(https://ngoldbaum.github.io/posts/loading-mnist-data-in-rust/)。在這之後,下一步我們必須弄清楚如何用Rust表示Python代碼中的Network類。最終我決定使用struct:

use ndarray::Array2;

#[derive(Debug)]
struct Network {
num_layers: usize,
sizes: Vec<usize>,
biases: Vec<array2>>,
weights: Vec<array2>>,
}/<array2>/<array2>/<usize>

該結構的初始化與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>> = Vec::new;
let mut weights: Vec<array2>> = 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,
}
}
}/<array2>/<array2>

有一點區別在於,在Python中我們使用numpy.random.randn初始化偏差和權重,而在Rust中我們使用ndarray::Array::random函數,並以rand::distribution::Distribution為參數,允許選擇任意的分佈。在上述代碼中,我們使用了rand::distributions::StandardNormal分佈。注意,我們使用了三個不同的包中定義的接口,其中兩個ndarray本身和ndarray-rand由ndarray作者維護,另一個rand則由其他開發人員維護。

Rust 能取代 Python,更好的实现神经网络?

整體式包的優點

原則上,最好不要將隨機數生成器放到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 能取代 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>> = zero_vec_like(&self.biases);
let mut nabla_w: Vec<array2>> = 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);
}
}
}/<array2>/<array2>

該函數使用了我定義的兩個輔助函數,因此更為簡潔:

fn to_tuple(inp: &[usize]) -> (usize, usize) {
match inp {
[a, b] => (*a, *b),
_ => panic!,
}
}

fn zero_vec_like(inp: &[Array2]) -> Vec<array2>> {
inp.iter

.map(|x| Array2::zeros(to_tuple(x.shape)))
.collect
}/<array2>

與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 能取代 Python,更好的实现神经网络?

用Rust重新編寫,一切都會好起來

到此為止,我用Rust編寫的代碼運行速度超過了我最初編寫的未經優化的Python代碼。然而,從Python這樣的動態解釋語言過渡到Rust這樣的性能優先的編譯語言,應該能達到10倍或更高性能,然而我只觀察到大約2倍的提升。我該如何測量Rust代碼的性能?幸運的是,有一個非常優秀的項目flamegraph(https://github.com/ferrous-systems/flamegraph)可以很容易地為Rust項目生成火焰圖。這個工具為cargo添加了一個flamegraph子命令,因此你只需運行cargo flamegraph,就可以運行代碼,然後寫一個flamegraph的svg文件,就可以通過Web瀏覽器觀測。

Rust 能取代 Python,更好的实现神经网络?

可能你以前從未見過火焰圖,因此在此簡單地說明一下,例程中程序的運行時間比例與該例程的條形寬度成正比。主函數位於圖形的底部,主函數調用的函數堆疊在上面。你可以通過這個圖形簡單地瞭解哪些函數在程序中佔用的時間最多——圖中非常“寬”的函數都在運行中佔用了大量時間,而非常高且寬的函數棧都代表其包含非常深入的棧調用,其代碼的運行佔用了大量時間。通過以上火焰圖,我們可以看到我的程序大約一半的時間都花在了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 能取代 Python,更好的实现神经网络?

Rust是否適合數據科學工作流程?

目前,我不得不說答案是“尚未”。如果我需要編寫能夠將依賴性降到最低的、經過優化的低級代碼,那麼我肯定會使用Rust。然而,要想利用Rust完全取代Python或C++,那麼我們尚需要等待更穩定和更完善的包生態系統。

原文:https://ngoldbaum.github.io/posts/python-vs-rust-nn/

本文為 CSDN 翻譯,轉載請註明來源出處。

【End】


分享到:


相關文章: