程式碼最佳化與程式加速指南——針對數值最佳化和深度學習領域

BigdogManLuo發表於2023-03-06

背景


當需要處理規模較大、任務較複雜的最佳化問題或訓練神經網路時,我們經常會遇到程式執行時間長或無法完成的情況。然而,這不一定是由於問題規模大或計算機硬體能力的限制。即使嘗試使用更高效能的伺服器或計算機,也不能保證能夠有效地加速程式碼執行。因為高效能的硬體通常需要與為高效能運算而設計的程式碼相匹配。

本文旨在為程式加速提供一些程式碼方面的最佳化思路,透過最佳化程式碼結構、設計高效能運算方案,來有效加速程式執行,提高程式執行效率。需要注意的是,本文只涉及程式碼層面的加速方案,不包括演算法、硬體等方面的最佳化措施。文章的撰寫基於個人經驗,如果有不足之處,敬請指出。

本文將介紹一些常見的程式碼最佳化技巧,例如使用向量化和並行化的方式來加速計算、減少記憶體的佔用以及利用編譯器最佳化程式碼的迴圈結構等。
主要思路

程式碼最佳化


簡單來說,實現程式最佳化的方式主要有兩種思路。一種是並行化任務,編寫並行化程式碼,以利用多核CPU或GPU的平行計算能力來加速程式執行。另一種則是利用編譯器的程式碼最佳化機制,將Python、MATLAB等需要直譯器執行的程式碼部分編譯成機器程式碼,以實現更快的程式執行速度。

1 並行化

程式最佳化的另一種思路是透過並行化加速程式。並行化需要軟體與硬體配合,但前提是總任務能夠被分解為同時進行的子任務。並行化有兩種方式,一種依靠多核CPU實現多程式操作,一種依靠GPU完成。下面我們將分別介紹這兩種方式的實現方法。

1.1 CPU多程式操作

多程式操作是利用CPU多核心的特性來實現並行化計算。在多程式操作中,程式將被分解成多個子任務,每個子任務都在獨立的程式中執行。這些程式可以並行地執行不同的任務,從而加速程式的執行。多程式操作可以使用Python的multi-processing庫來實現。

1.1.1 多程式

當我們在電腦上執行程式時,實際上是啟動了一個程式。程式是指在作業系統中執行的一個程式,它佔據著系統的一些資源,如記憶體、CPU時間等。在傳統的單程式計算模型中,所有的任務都在同一個程式中執行,如果任務需要進行復雜的計算,就會耗費大量時間,且無法利用多核CPU的優勢。而多程式技術可以將任務分解為多個子任務,每個子任務在一個獨立的程式中執行,從而實現平行計算,提高了計算效率。

1.1.2 多程式的Python實現

在使用multi-processing庫進行多程式操作時,首先需要將任務分解成多個子任務,並將每個子任務交給不同的程式來執行。

這是multiprocessing庫的程式碼框架

import multiprocessing

def worker(num):
    """子程式要執行的任務"""
    print(f"Worker {num} is running")
    return

if __name__ == "__main__":
    processes = []
    num_processes = 4

    # 建立多個子程式
    for i in range(num_processes):
        p = multiprocessing.Process(target=worker, args=(i,))
        processes.append(p)
        p.start()

    # 等待所有子程式完成
    for p in processes:
        p.join()

    print("All workers are done")

在這個例子中,我們首先定義了一個worker函式,它是每個子程式要執行的任務。然後在主程式(if name=="main")中,我們建立了num_processes個子程式,並將它們新增到processes列表中。接著,我們遍歷processes列表,用Process類中的start()方法啟動每個子程式,並等待它們完成。最後輸出"All workers are done",表示所有子程式都已經執行完畢。

這個示例程式碼只是一個簡單的例子,實際應用中,我們可以根據具體情況編寫更復雜的子程式任務函式,並使用multiprocessing庫中提供的各種工具實現更復雜的多程式操作。

1.2 GPU

相比於CPU,GPU有著更多的計算核心和更高的計算能力,可以更好地支援並行化計算。因此,利用GPU進行程式碼最佳化可以大幅度提高程式的執行效率。而要在Python中使用GPU進行程式碼最佳化,則需要使用GPU程式設計框架,比較常見的有NVIDIA開發的CUDA框架以及OpenCL框架。

1.2.1 GPU平行計算的Python實現

其中,CUDA是由NVIDIA推出的GPU程式設計框架,它可以讓開發者利用GPU的平行計算能力,加速各種計算密集型任務。CUDA提供了一組API,可以方便地進行GPU的程式設計,支援C/C++、Python、Java等多種程式語言。在使用CUDA時,需要使用CUDA工具包,其中包括CUDA驅動程式、CUDA執行時庫和CUDA工具。

以下是一個簡單的利用CUDA框架實現向量加法的例子:

import numpy as np
from numba import cuda

# 定義向量加法函式
@cuda.jit
def vector_add(a, b, c):
    i = cuda.grid(1)
    if i < len(c):
        c[i] = a[i] + b[i]

# 定義主程式
if __name__ == '__main__':
    # 定義向量大小
    n = 100000

    # 在主機上生成隨機向量
    a = np.random.randn(n).astype(np.float32)
    b = np.random.randn(n).astype(np.float32)
    c = np.zeros(n, dtype=np.float32)

    # 將資料傳輸到GPU視訊記憶體中
    d_a = cuda.to_device(a)
    d_b = cuda.to_device(b)
    d_c = cuda.to_device(c)

    # 定義執行緒塊和執行緒數量
    threads_per_block = 64
    blocks_per_grid = (n + (threads_per_block - 1)) // threads_per_block

    # 執行向量加法操作
    vector_add[blocks_per_grid, threads_per_block](d_a, d_b, d_c)

    # 將結果從GPU視訊記憶體中傳輸回主機記憶體
    d_c.copy_to_host(c)

    # 列印結果
    print(c)

在上述程式碼中,我們首先定義了一個vector_add函式,用於將兩個向量相加並將結果儲存在第三個向量中。然後我們生成了兩個隨機向量,並將其傳輸到GPU視訊記憶體中。接著,我們定義了執行緒塊和執行緒數量,並在GPU部分,我們需要注意的是程式碼的向量化,也就是多利用矩陣運算編寫程式碼,而不是多重for loop的序列計算。這樣可以充分利用GPU的平行計算能力,加速神經網路的訓練過程。例如,可以使用cuBLAS庫來實現矩陣乘法,使用cuDNN庫來實現卷積操作等。執行了向量加法操作。最後,我們將結果從GPU視訊記憶體中傳輸回主機記憶體,並列印出結果。

值得注意的是,由於CUDA框架需要GPU的支援,因此我們在使用CUDA框架時需要保證系統中有可用的GPU。

1.2.2利用GPU訓練神經網路

當涉及到訓練深度神經網路時,GPU可以發揮其優勢,因為神經網路的訓練往往需要進行大量的矩陣運算,這正是GPU擅長的任務。為了使用GPU進行深度神經網路的訓練,需要利用一些特定的框架和庫,比如TensorFlow和PyTorch。
在利用CUDA框架進行神經網路訓練時,首先,需要定義一個在GPU上執行的張量(tensor)來儲存神經網路的引數。然後,將資料載入到GPU記憶體中,並將計算任務分配到GPU核心上。
下面是一個簡單的基於PyTorch和CUDA的神經網路訓練程式碼示例:

import torch
import torch.nn as nn
import torch.optim as optim

# 定義神經網路模型
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

# 定義訓練資料和標籤
inputs = torch.randn(1, 3, 32, 32)
labels = torch.randn(1, 10)

# 將神經網路模型移動到GPU上
net = Net().cuda()

# 將訓練資料和標籤移動到GPU上
inputs = inputs.cuda()
labels = labels.cuda()

# 定義損失函式和最佳化器
criterion = nn.MSELoss()
optimizer = optim.SGD(net.parameters(), lr=0.01)

# 開始訓練
for epoch in range(100):
    optimizer.zero_grad()
    outputs = net(inputs)
    loss = criterion(outputs, labels)
    loss.backward()
    optimizer.step()
    print('Epoch %d, Loss: %.4f' % (epoch+1, loss.item()))

在上面的程式碼中,我們透過xxx.cuda()的方法將模型、訓練資料和標籤都載入到GPU記憶體中,從而使得訓練迴圈中涉及的運算全部在GPU上完成。

1.2.3 注意事項

值得特別注意的是,如果要利用GPU實現加速,那麼一定要保證程式碼的向量化。簡單來說,就是儘可能用矩陣運算的方式來表示數值計算過程,而不是用多重for loop的形式。

這是因為,GPU相較於CPU只是強在並行上,但是計算能力是差於CPU的,多層for迴圈這樣的序列計算並不適合GPU。用沐神的話說,如果把CPU比作是一個大學生的話,GPU就像是一群小學生。大學生可以做微積分這樣的任務,小學生只能做加減乘除。但如果把一個微積分這樣的任務分解成多個加減乘除,那麼一群小學生這樣的GPU的優勢就體現出來了。

舉一個簡單的例子,例如我們要設計一個這樣的損失函式:

\[L_\theta=-\sum_{i=1}^{N}\sum_{j=1}^{N_i}w_{ij}log(p(x_i^j|M)) \]

如果沒有刻意注意程式碼的向量化,那麼很容易想到的思路是這樣的:

	def lossFunc(y_pred,sols,objs):
	batch size=y_pred.shape[0]
	loss=torch.tensor(0.0)
	for i in range(batch_size):
		nSols=sols[i].shape[0] #當前batch(MIP)下的可行解數目
		nVars=sols[i].shape[1]
		#目標斷數歸-化
		den=objs[i].sum()
		for l in range(objs[i].shape[0]):
			objs[i][l]=objs[i][l]/den
		den=sum(exp(-objs[i]))#計算wii係數的分母
		sum1=torch.tensor(0.0)
		for j in range(nSols):
			#計算權重wij
			w=exp(-objs[il[j])/den
			#計算可行解生成的機率
			P=torch.tensor(1.9)
			for k in range(nVars):
				if sols[il[j,k]==1:
					P=p*y_pred[i][k]
				elif sols[il[j,k]==0:
					P=p*(1-y_pred[i][k])
			#計算求和
			sum1+=w*p
		loss+=sum1
	loss=-loss
	return loss

但實際上用多重for迴圈設計的損失函式並不利於在GPU上工作,甚至比不上在CPU上執行這樣的程式碼。最後測試結果發現如果使用這樣的損失函式,一個epoch大約跑1個小時....所以也不可能把神經網路訓練出來了。

def lossFunc(y,sols,objs):
    """
    損失函式
    Parameters
    ----------
    y : 神經網路輸出 batch_size x nVars
    sols : 可行解集合 batch_size x nSols x nVars
    objs : 可行解對應的目標函式值 batch_size x nSols

    Returns
    -------
    loss : 在當前batch上的損失
    """
    objs=objs/15
    eObjs=exp(-objs)
    den=eObjs.sum(axis=1)
    den=den.unsqueeze(1)
    w=eObjs/den
    y=y.unsqueeze(1)
    p=y*sols+(1-y)*(1-sols)
    p=log(p+1e-45)
    P=p.sum(axis=2)
    loss=-(w*P).sum()
    return loss

所以正確的寫法應該是這樣的,不僅在GPU上加速效果明顯,並且看起來簡潔多了...但是缺點就是,得不斷利用矩陣運算的廣播機制,並且最好在小規模的例子上多手算幾遍,邊算邊設計。

2 編譯器加速

2.1 原理

在計算機程式設計中,編譯型語言和解釋型語言是兩種常見的語言型別。與解釋型語言相比,編譯型語言由於在執行之前需要先將程式碼編譯成可執行的二進位制程式碼,所以執行速度更快。這是因為編譯器能夠將原始碼最佳化為更高效的機器程式碼,從而加快程式的執行速度。

編譯器最佳化的方式有很多種,其中最常見的方式包括:

  1. 消除不必要的計算:編譯器可以在編譯程式碼時識別出不必要的計算,從而避免浪費計算資源。

  2. 迴圈展開:迴圈展開是指將迴圈體中的程式碼重複執行幾次,從而減少迴圈的次數。這樣可以提高程式的執行速度。

  3. 矩陣/向量化:矩陣/向量化是指將多個資料放入矩陣或向量中,然後一次性進行計算。這種方式可以減少迴圈的次數,從而提高程式的執行速度。

為了幫助開發者更加方便地利用編譯器最佳化程式碼,一些開源JIT編譯器如Numba被開發出來。這些編譯器可以將Python等解釋型語言程式碼轉換為可執行的機器程式碼,從而提高程式的執行速度。

2.2 基於Numba的Python加速方案

Numba是一種基於LLVM編譯器的開源JIT編譯器,可以將Python程式碼轉換為機器程式碼,從而實現程式碼加速。Numba支援多種最佳化技術,包括迴圈展開、程式碼向量化等。使用Numba可以極大地提高Python程式碼的執行速度。

下面是一個使用Numba實現程式碼加速的大致程式碼框架:

  1. 匯入numba庫
  2. 定義一個需要最佳化的函式
  3. 使用 @numba.jit 裝飾器對函式進行裝飾,生成numba最佳化後的函式
  4. 呼叫最佳化後的函式

下面以一個對比案例為例,來說明如何使用Numba實現Python程式碼加速。
假設我們有一個多重巢狀的for迴圈,用於計算一個矩陣的行列式。這是一個非常計算密集的操作,可以使用Numba進行加速。

原始Python程式碼:

def det(matrix):
    n = len(matrix)
    if n == 1:
        return matrix[0][0]
    elif n == 2:
        return matrix[0][0] * matrix[1][1] - matrix[0][1] * matrix[1][0]
    else:
        result = 0
        for j in range(n):
            sub_matrix = []
            for i in range(1, n):
                row = []
                for k in range(n):
                    if k != j:
                        row.append(matrix[i][k])
                sub_matrix.append(row)
            result += matrix[0][j] * det(sub_matrix) * (-1) ** j
        return result

可以看到,這個函式包含了多重巢狀的for迴圈,非常容易受到Python直譯器的效能限制。現在我們使用Numba進行加速。

最佳化後的Numba程式碼:

import numba

@numba.jit(nopython=True)
def det(matrix):
    n = len(matrix)
    if n == 1:
        return matrix[0][0]
    elif n == 2:
        return matrix[0][0] * matrix[1][1] - matrix[0][1] * matrix[1][0]
    else:
        result = 0
        for j in range(n):
            sub_matrix = np.zeros((n-1, n-1))
            for i in range(1, n):
                for k in range(n):
                    if k != j:
                        sub_matrix[i-1, k-(k>j)] = matrix[i, k]
            result += matrix[0][j] * det(sub_matrix) * (-1) ** j
        return result

我們使用了@numba.jit(nopython=True)裝飾器,將函式宣告為Numba可加速的函式。同時,我們使用了Numpy陣列代替了Python的列表,並使用Numpy陣列的切片操作和廣播功能,以減少迴圈和記憶體分配。

測試程式碼:

import time
import numpy as np

# Generate a random 10x10 matrix
matrix = np.random.rand(10, 10)

# Time the original Python code
start = time.time()
d = det(matrix)
end = time.time()
print(f"Python code took {end-start:.4f} seconds, result={d}")

# Time the Numba-optimized code
start = time.time()
d = det(matrix)
end = time.time()
print(f"Numba-optimized code took {end-start:.4f} seconds, result={d}")

測試結果:

Python code took 0.5960 seconds, result=-0.004127521725273144
Numba-optimized code took 0.0040 seconds, result=-0.004127521725273144

使用Numba最佳化Python程式碼的大致流程如上所述。對於需要最佳化的函式,我們可以透過新增numba.jit 裝飾器實現程式碼最佳化,該裝飾器會自動進行型別推導並將Python程式碼編譯為機器碼。

需要注意的是,Numba並不是萬能的,它能夠加速的程式碼型別是有一定限制的。例如,迴圈巢狀過深的程式碼可能不適合使用Numba進行最佳化。因此,在使用Numba時,需要仔細評估程式碼的結構和型別,以確定是否適合使用Numba進行最佳化。

此外,還需要注意的是,Numba在進行程式碼最佳化時,會將Python程式碼轉換為LLVM IR(Intermediate Representation,中間表示形式),這一過程可能會導致程式碼的可讀性降低。因此,在使用Numba進行最佳化時,需要仔細考慮程式碼可讀性和可維護性的平衡。

總結


在這篇技術部落格中,我們介紹了幾種常見的程式碼最佳化技術,包括平行計算、編譯器最佳化等。這些技術可以在不改變程式邏輯的前提下,提高程式碼的執行效率和效能。其中,我們透過對比平行計算與序列計算的執行效率差異,詳細講解了如何利用GPU加速神經網路訓練。此外,我們還介紹了編譯型語言比解釋型語言執行快的原因,以及編譯器如何最佳化程式碼實現加速的原理。最後,我們著重介紹了一個開源的JIT編譯器Numba,並透過一個實際案例演示瞭如何利用Numba最佳化多重巢狀for迴圈的Python程式碼。

值得一提的是,這篇技術部落格是在chatGPT的幫助下共同完成的。作者提供了大致的寫作思路和要點,而主要內容和其中的程式碼案例均由chatGPT構建。我們相信,本篇部落格對於想要提高程式碼效能和效率的程式設計師們會有所幫助。


附圖




相關文章