PyTorch是當(dāng)今生產(chǎn)環(huán)境中最流行的深度學(xué)習(xí)框架之一。隨著模型變得日益復(fù)雜、數(shù)據(jù)集日益龐大，優(yōu)化模型訓(xùn)練性能對(duì)于縮短訓(xùn)練時(shí)間和提高生產(chǎn)力變得至關(guān)重要。

本文將分享幾個(gè)最新的性能調(diào)優(yōu)技巧，以加速跨領(lǐng)域的機(jī)器學(xué)習(xí)模型的訓(xùn)練。這些技巧對(duì)任何想要使用PyTorch實(shí)現(xiàn)高級(jí)性能調(diào)優(yōu)的人都大有幫助。

技巧1：通過(guò)分析識(shí)別性能瓶頸

在開(kāi)始調(diào)優(yōu)之前，你應(yīng)該了解模型訓(xùn)練管道中的瓶頸。分析(Profiling)是優(yōu)化過(guò)程中的關(guān)鍵步驟，因?yàn)樗兄谧R(shí)別需要注意的內(nèi)容。你可以從PyTorch的內(nèi)置自動(dòng)求梯度分析器、TensorBoard和英偉達(dá)的Nsight系統(tǒng)中進(jìn)行選擇。下面不妨看一下三個(gè)示例。

• 代碼示例：自動(dòng)求梯度分析器

import torch.autograd.profiler as profiler
with profiler.profile(use_cuda=True) as prof:
# Run your model training code here
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))

在這個(gè)示例中，PyTorch的內(nèi)置自動(dòng)求梯度分析器識(shí)別梯度計(jì)算開(kāi)銷。use_cuda=True參數(shù)指定你想要分析CUDA內(nèi)核執(zhí)行時(shí)間。prof.key_average()函數(shù)返回一個(gè)匯總分析結(jié)果的表，按總的CUDA時(shí)間排序。

• 代碼示例：TensorBoard集成

import torch.utils.tensorboard as tensorboard
writer = tensorboard.SummaryWriter()
# Run your model training code here
writer.add_scalar('loss', loss.item(), global_step)
writer.close()

你還可以使用TensorBoard集成來(lái)顯示和分析模型訓(xùn)練。SummaryWriter類將匯總數(shù)據(jù)寫入到一個(gè)文件，該文件可以使用TensorBoard GUI加以顯示。

• 代碼示例：英偉達(dá)Nsight Systems

nsys profile -t cpu,gpu,memory python your_script.py

對(duì)于系統(tǒng)級(jí)分析，可以考慮英偉達(dá)的Nsight Systems性能分析工具。上面的命令分析了Python腳本的CPU、GPU和內(nèi)存使用情況。

技巧2：加速數(shù)據(jù)加載以提升速度和GPU利用率

數(shù)據(jù)加載是模型訓(xùn)練管道的關(guān)鍵組成部分。在典型的機(jī)器學(xué)習(xí)訓(xùn)練管道中，PyTorch的數(shù)據(jù)加載器在每個(gè)訓(xùn)練輪次開(kāi)始時(shí)從存儲(chǔ)中加載數(shù)據(jù)集。然后，數(shù)據(jù)集被傳輸?shù)紾PU實(shí)例的本地存儲(chǔ)，并在GPU內(nèi)存中進(jìn)行處理。如果數(shù)據(jù)傳輸?shù)紾PU的速度跟不上GPU的計(jì)算速度，就會(huì)導(dǎo)致GPU周期浪費(fèi)。因此，優(yōu)化數(shù)據(jù)加載對(duì)于加快訓(xùn)練速度、盡量提升GPU利用率至關(guān)重要。

為了盡量減少數(shù)據(jù)加載瓶頸，你可以考慮以下優(yōu)化：

1. 使用多個(gè)worker并行化數(shù)據(jù)加載：使用PyTorch的數(shù)據(jù)加載器與多個(gè)worker并行化數(shù)據(jù)加載。這允許CPU并行加載和處理數(shù)據(jù)，從而減少GPU空閑時(shí)間。
2. 使用緩存加速數(shù)據(jù)加載：使用Alluxio作為訓(xùn)練節(jié)點(diǎn)和存儲(chǔ)之間的緩存層，以實(shí)現(xiàn)數(shù)據(jù)按需加載，而不是將遠(yuǎn)程數(shù)據(jù)直接加載到本地存儲(chǔ)或?qū)⒂?xùn)練數(shù)據(jù)復(fù)制到本地存儲(chǔ)。

• 代碼示例：并行化數(shù)據(jù)加載

下面這個(gè)示例是使用PyTorch的數(shù)據(jù)加載器和多個(gè)worker并行化加載數(shù)據(jù)：

import torch
from torch.utils.data import DataLoader, Dataset
class MyDataset(Dataset):
def __init__(self, data_path):
self.data_path = data_path
def __getitem__(self, index):
# Load and process data for the given index
data = load_data(self.data_path, index)
data = preprocess_data(data)
return data
def __len__(self):
return len(self.data_path)
dataset = MyDataset(data_path='path/to/data')
data_loader = DataLoader(dataset, batch_size=32, num_workers=4)
for batch in data_loader:
# Process the batch on the GPU
inputs, labels = batch
outputs = model(inputs)
loss = criterion(outputs, labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()

在這個(gè)示例中，定義了自定義數(shù)據(jù)集類MyDataset。它為每個(gè)索引加載和處理數(shù)據(jù)。然后創(chuàng)建一個(gè)有多個(gè)worker(本例中有四個(gè))的數(shù)據(jù)加載器實(shí)例來(lái)并行化加載數(shù)據(jù)。

• 代碼示例：使用Alluxio緩存來(lái)加速PyTorch的數(shù)據(jù)加載

Alluxio是一個(gè)開(kāi)源分布式緩存系統(tǒng)，提供快速訪問(wèn)數(shù)據(jù)的機(jī)制。Alluxio緩存可以識(shí)別從底部存儲(chǔ)(比如Amazon S3)頻繁訪問(wèn)的數(shù)據(jù)，并在Alluxio集群的NVMe存儲(chǔ)上分布式存儲(chǔ)熱數(shù)據(jù)的多個(gè)副本。如果使用Alluxio作為緩存層，你可以顯著縮短將數(shù)據(jù)加載到訓(xùn)練節(jié)點(diǎn)所需的時(shí)間，這在處理大規(guī)模數(shù)據(jù)集或慢速存儲(chǔ)系統(tǒng)時(shí)特別有用。

下面這個(gè)示例表明了你如何使用Alluxio與PyTorch和fsspec(文件系統(tǒng)規(guī)范)來(lái)加速數(shù)據(jù)加載：

首先，安裝所需的依賴項(xiàng)：

pip install alluxiofs
pip install s3fs

接下來(lái)，創(chuàng)建一個(gè)Alluxio實(shí)例：

import fsspec
from alluxiofs import AlluxioFileSystem
# Register Alluxio to fsspec
fsspec.register_implementation("alluxiofs", AlluxioFileSystem, 
clobber=True)
# Create Alluxio instance
alluxio_fs = fsspec.filesystem("alluxiofs", etcd_hosts="localhost", 
target_protocol="s3")

然后，使用Alluxio和PyArrow在PyTorch中加載Parquet文件這個(gè)數(shù)據(jù)集：

# Example: Read a Parquet file using Pyarrow
import pyarrow.dataset as ds
dataset = ds.dataset("s3://example_bucket/datasets/example.parquet", 
filesystem=alluxio_fs)
# Get a count of the number of records in the parquet file
dataset.count_rows()
# Display the schema derived from the parquet file header record
dataset.schema
# Display the first record
dataset.take(0)

在這個(gè)示例中，創(chuàng)建了一個(gè)Alluxio實(shí)例并將其傳遞給PyArrow的dataset函數(shù)。這允許我們通過(guò)Alluxio緩存層從底層存儲(chǔ)系統(tǒng)(本例中為S3)讀取數(shù)據(jù)。

技巧3：為資源利用率優(yōu)化批任務(wù)大小

優(yōu)化GPU利用率的另一項(xiàng)重要技術(shù)是調(diào)整批任務(wù)大小，它會(huì)顯著影響GPU和內(nèi)存利用率。

• 代碼示例：批任務(wù)大小優(yōu)化

import torch
import torchvision
import torchvision.transforms as transforms
# Define the model and optimizer
model = torchvision.models.resnet50(pretrained=True)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
# Define the data loader with a batch size of 32
data_loader = torch.utils.data.DataLoader(
dataset,
batch_size=32,
shuffle=True,
num_workers=4
)
# Train the model with the optimized batch size
for epoch in range(5):
for inputs, labels in data_loader:
inputs, labels = inputs.cuda(), labels.cuda()
optimizer.zero_grad()
outputs = model(inputs)
loss = torch.nn.CrossEntropyLoss()(outputs, labels)
loss.backward()
optimizer.step()

在本例中，批任務(wù)大小定義為32。batch_size參數(shù)指定了每個(gè)批中的樣本數(shù)量。shuffle=True參數(shù)隨機(jī)化批處理的順序，num_workers=4參數(shù)指定用于加載數(shù)據(jù)的worker線程的數(shù)量。你可以嘗試不同的批任務(wù)大小，以找到在可用內(nèi)存范圍內(nèi)盡量提高GPU利用率的最佳值。

技巧4：可識(shí)別GPU的模型并行性

處理大型復(fù)雜模型時(shí)，單個(gè)GPU的限制可能會(huì)成為訓(xùn)練的瓶頸。模型并行化可以通過(guò)在多個(gè)GPU上共同分布模型以使用它們的加速能力來(lái)克服這一挑戰(zhàn)。

1.利用PyTorch的DistributedDataParallel(DDP)模塊

PyTorch提供了DistributedDataParallel(DDP)模塊，它可以通過(guò)支持多個(gè)后端來(lái)實(shí)現(xiàn)簡(jiǎn)單的模型并行化。為了盡量提高性能，使用NCCL后端，它針對(duì)英偉達(dá)GPU進(jìn)行了優(yōu)化。如果使用DDP來(lái)封裝模型，你可以跨多個(gè)GPU無(wú)縫分布模型，將訓(xùn)練擴(kuò)展到前所未有的層面。

• 代碼示例：使用DDP

import torch
from torch.nn.parallel import DistributedDataParallel as DDP
# Define your model and move it to the desired device(s)
model = MyModel()
device_ids = [0, 1, 2, 3] # Use 4 GPUs for training
model.to(device_ids[0])
model_ddp = DDP(model, device_ids=device_ids)
# Train your model as usual

2.使用PyTorch的Pipe模塊實(shí)現(xiàn)管道并行處理

對(duì)于需要順序處理的模型，比如那些具有循環(huán)或自回歸組件的模型，管道并行性可以改變游戲規(guī)則。PyTorch的Pipe允許你將模型分解為更小的部分，在單獨(dú)的GPU上處理每個(gè)部分。這使得復(fù)雜模型可以高效并行化，縮短了訓(xùn)練時(shí)間，提高了整體系統(tǒng)利用率。

3.減少通信開(kāi)銷

雖然模型并行化提供了巨大的好處，但也帶來(lái)了設(shè)備之間的通信開(kāi)銷。以下是盡量減小影響的幾個(gè)建議：

a.最小化梯度聚合：通過(guò)使用更大的批大小或在同步之前本地累積梯度，減少梯度聚合的頻次。

b.使用異步更新：使用異步更新，隱藏延遲和最大化GPU利用率。

c.啟用NCCL的分層通信：讓NCCL庫(kù)決定使用哪種分層算法：環(huán)還是樹(shù)，這可以減少特定場(chǎng)景下的通信開(kāi)銷。

d.調(diào)整NCCL的緩沖區(qū)大?。赫{(diào)整NCCL_BUFF_SIZE環(huán)境變量，為你的特定用例優(yōu)化緩沖區(qū)大小。

技巧5：混合精度訓(xùn)練

混合精度訓(xùn)練是一種強(qiáng)大的技術(shù)，可以顯著加速模型訓(xùn)練。通過(guò)利用現(xiàn)代英偉達(dá)GPU的功能，你可以減少訓(xùn)練所需的計(jì)算資源，從而加快迭代時(shí)間并提高生產(chǎn)力。

1.使用Tensor Cores加速訓(xùn)練

英偉達(dá)的Tensor Cores是專門用于加速矩陣乘法的硬件塊。這些核心可以比傳統(tǒng)的CUDA核心更快地執(zhí)行某些操作。

2.使用PyTorch的AMP簡(jiǎn)化混合精度訓(xùn)練

實(shí)現(xiàn)混合精度訓(xùn)練可能很復(fù)雜，而且容易出錯(cuò)。幸好，PyTorch提供了一個(gè)amp模塊來(lái)簡(jiǎn)化這個(gè)過(guò)程。使用自動(dòng)混合精度(AMP)，你可以針對(duì)模型的不同部分在不同精度格式(例如float32和float16)之間切換，從而優(yōu)化性能和內(nèi)存使用。

• 代碼示例：PyTorch的AMP

以下這個(gè)示例表明了如何使用PyTorch的amp模塊來(lái)實(shí)現(xiàn)混合精度訓(xùn)練：

import torch
from torch.amp import autocast
# Define your model and optimizer
model = MyModel()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
# Enable mixed precision training with AMP
with autocast(enabled=True, dtype=torch.float16):
# Train your model as usual
for epoch in range(10):
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()

3.使用低精度格式優(yōu)化內(nèi)存使用

以較低精度格式(比如float16)存儲(chǔ)模型權(quán)重可以顯著減少內(nèi)存使用。當(dāng)處理大型模型或有限的GPU資源時(shí)，這點(diǎn)尤為重要。如果使用精度較低的格式，你可以將較大的模型放入到內(nèi)存中，從而減少對(duì)昂貴內(nèi)存訪問(wèn)的需求，并提高整體訓(xùn)練性能。

記住要嘗試不同的精度格式并優(yōu)化內(nèi)存使用，以便為你的特定用例獲得最佳結(jié)果。

技巧6：新的硬件優(yōu)化：GPU和網(wǎng)絡(luò)

新的硬件技術(shù)出現(xiàn)為加速模型訓(xùn)練提供了大好機(jī)會(huì)。記得嘗試不同的硬件配置，并優(yōu)化你的工作流，以便為特定用例獲得最佳結(jié)果。

1.利用英偉達(dá)A100和H100 GPU

最新的英偉達(dá)A100和H100 GPU有先進(jìn)的性能和內(nèi)存帶寬。這些GPU為用戶提供了更多的處理能力，使用戶能夠訓(xùn)練更大的模型、處理更大的批任務(wù)，并縮短迭代時(shí)間。

2.利用NVLink和InfiniBand加速GPU-GPU通信

當(dāng)跨多個(gè)GPU訓(xùn)練大型模型時(shí)，設(shè)備之間的通信開(kāi)銷可能成為一大瓶頸。英偉達(dá)的NVLink互連技術(shù)在GPU之間提供了高帶寬低延遲的鏈路，從而實(shí)現(xiàn)更快的數(shù)據(jù)傳輸和同步。此外，InfiniBand互連技術(shù)為連接多個(gè)GPU和節(jié)點(diǎn)提供了一種易于擴(kuò)展的高性能解決方案。它有助于盡量減小通信開(kāi)銷，縮短同步梯度和加速模型訓(xùn)練所花費(fèi)的時(shí)間。

結(jié)語(yǔ)

上述這六個(gè)技巧將幫助你顯著加快模型訓(xùn)練速度。切記，獲得最佳結(jié)果的關(guān)鍵是嘗試這些技術(shù)的不同組合，為你的特定用例找到最佳配置。

原文標(biāo)題：This Is How To Optimize PyTorch for Faster Model Training，作者：Hope Wang

鏈接：https://thenewstack.io/this-is-how-to-optimize-pytorch-for-faster-model-training/。