九個技巧讓你的PyTorch模型訓練飛快！
也許你仍然在使用32位精度進行計算，或者甚至只是在單個GPU上進行訓練。
然而，隨著科技的進步，我們已經有了更好的選擇。使用更高精度的計算，如16位浮點數或混合精度，可以提高訓練速度并減少內存消耗。同時，利用多個GPU進行并行訓練，可以大大加快訓練過程。
筆者在這里總結了提升Pytorch模型訓練速度的9個技巧，與大家分享~
這些優化技巧可以在PyTorch-Lightning庫中找到。PyTorch-Lightning是建立在PyTorch之上的一個封裝，它提供了自動化訓練的功能，同時允許開發者完全控制關鍵的模型組件。
這里以MNIST定義LightningModel并使用Trainer來訓練模型為例。
#導入PyTorch-Lightning庫中的Trainer類，用于管理訓練過程
from pytorch_lightning import Trainer
#創建LightningModule實例，作為要訓練的模型
model = LightningModule(…)
#創建Trainer實例，用于配置和管理訓練過程
trainer = Trainer()
#開始訓練
trainer.fit(model)trainer.fit()方法將根據LightningModule中定義的訓練邏輯來執行訓練步驟，并自動處理數據加載、優化器配置、學習率調度等細節。
使用DataLoaders
使用DataLoaders來加載數據是獲得訓練速度提升的最簡單方法之一。通過保存h5py或numpy文件以加速數據加載的時代已經過去了，而現在可以「使用PyTorch的DataLoader來輕松加載圖像數據」（對于NLP數據，請參考TorchText庫）。
在PyTorch-Lightning中，不需要顯式地編寫訓練循環，只需要定義好DataLoaders和Trainer，PyTorch-Lightning會在需要的時候自動調用它們。
下面是一個加載MNIST數據集并使用DataLoader進行訓練的示例代碼：
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
dataset = MNIST(root=self.hparams.data_root, train=train, download=True)
loader = DataLoader(dataset, batch_size=32, shuffle=True)
for batch in loader:
x, y = batch
model.training_step(x, y)
# 其他訓練邏輯在這個示例中，首先創建了一個MNIST數據集的實例，然后使用DataLoader將其封裝成一個可迭代的數據加載器。在訓練循環中，可以遍歷DataLoader，每次獲取一個batch的數據，并將其傳遞給模型的training_step()方法進行訓練。
通過使用DataLoaders，可以更高效地加載和處理大量的訓練數據，從而提高訓練速度。此外，DataLoader還支持數據的隨機打亂（shuffle）、批量大小（batch_size）等參數設置，可以根據實際需求進行調整。
DataLoaders中設置num_workers參數
在DataLoaders中，可以通過設置num_workers參數來允許批量并行加載數據，從而加速訓練過程。下面是一個示例代碼：
# 慢的方式
loader = DataLoader(dataset, batch_size=32, shuffle=True)
# 快的方式（使用10個workers）
loader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=10)在第一個示例中，創建了一個DataLoader，沒有指定num_workers參數，這意味著數據加載將在主進程中進行，而不會并行化。
在第二個示例中，通過將num_workers設置為10，啟用了批量并行加載。這意味著數據加載將在10個worker進程中進行，并行地加載多個batch，從而加速數據加載的過程。
通過適當設置num_workers參數，可以根據系統的硬件和資源情況，選擇合適的worker數量來提高數據加載的效率。然而，需要注意的是，并不是worker越多越好，過多的worker可能會導致資源競爭和性能下降。
Batch size
在進行下一個優化步驟之前，增加批量大小（batch size）到CPU-RAM或GPU-RAM允許的最大范圍是一個重要的優化策略。
增加批量大小可以帶來以下好處：
- 更高效地利用計算資源，尤其是GPU的并行計算能力。
- 減少數據加載和傳輸的次數，提高訓練速度。
- 可以獲得更穩定的梯度估計，有助于模型收斂。
然而，增加批量大小也會帶來一些挑戰：
- 內存占用增加：較大的批量大小需要更多的內存空間來存儲數據和梯度。
- 學習率調整：增加批量大小后，通常需要相應地增加學習率，以保持相似的收斂行為。
因此，在增加批量大小之前，需要確保你的硬件和資源可以支持更大的批量大小，并相應地調整學習率。
梯度累積
梯度累積（Gradient Accumulation）是一種在計算資源有限的情況下，模擬較大批量大小的技術。通過多次執行前向傳播、反向傳播和優化步驟，將梯度累積起來，以獲得與較大批量大小相同的效果。
下面是一個使用梯度累積的示例代碼：
# 清除上一步的梯度
optimizer.zero_grad()
# 16次梯度累積步驟
scaled_loss = 0
for accumulated_step_i in range(16):
out = model.forward()
loss = some_loss(out, y)
loss.backward()
scaled_loss += loss.item()
# 更新權重
optimizer.step()
# 損失值現在按累積批次數量進行縮放
actual_loss = scaled_loss / 16在這個示例中，通過循環執行16個梯度累積步驟，每個步驟進行前向傳播、計算損失、反向傳播和梯度累積。然后調用optimizer.step()來更新權重。
在PyTorch-Lightning中，只需要設置accumulate_grad_batches參數來指定梯度累積的次數。例如：
trainer = Trainer(accumulate_grad_batches=16)
trainer.fit(model)保留的計算圖
在記錄損失值時，為了避免撐爆內存，只存儲損失的數值而不是整個計算圖。可以使用.item()方法來獲取損失的數值。
# 方式1
losses.append(loss)
# 方式2
losses.append(loss.item())在方式1中，損失值loss會保留整個計算圖的副本，這會占用大量的內存空間。而方式2中，使用loss.item()來獲取損失的數值，并將其存儲到列表中，這樣就只保留了數值，而不會占用過多的內存。
PyTorch-Lightning會非常小心地確保不會保留計算圖的副本，盡量減少內存的占用。因此，在使用PyTorch-Lightning時，可以放心地使用.item()方法來獲取損失的數值，而不必擔心內存問題。
單個GPU訓練
完成上述步驟之后，即可開始在GPU上進行訓練。GPU上進行訓練可以利用多個GPU核心之間的并行計算，從而加速訓練過程。
在進行GPU訓練時，需要做兩件事情：
- 將模型移動到GPU上；
- 在每次數據通過時將數據放到GPU上。
下面是在PyTorch中進行GPU訓練的示例代碼：
# 將模型放到GPU上
model.cuda()
# 將數據放到GPU上
x = x.cuda()
# 在GPU上運行
model(x)如果使用PyTorch-Lightning，幾乎不需要做任何額外的工作，只需要設置Trainer的gpus參數來指定要使用的GPU數量。
# 指定訓練的gpu id
trainer = Trainer(gpus=[0])
trainer.fit(model)在進行GPU訓練時，需注意限制CPU和GPU之間的數據傳輸次數。盡量避免頻繁地在CPU和GPU之間復制數據。
此外，還要注意調用強制GPU同步的操作，如清空內存緩存torch.cuda.empty_cache()。這樣的操作會阻塞所有GPU，直到它們都完成同步。
然而，如果使用PyTorch-Lightning，則通常不需要擔心這些問題。PyTorch-Lightning會小心地處理GPU同步和內存管理，以確保高效的訓練過程。
使用16-bit精度
使用16-bit精度是一種驚人的技術，可以將內存占用減半。大多數模型通常使用32位精度進行訓練，但是研究表明，使用16位精度的模型也可以表現得很好。混合精度則意味著在某些部分使用16位精度，但將權重等內容保持在32位精度。
要在PyTorch中使用16位精度，可以安裝NVIDIA的apex庫，并對模型進行如下更改：
# 在模型和優化器上啟用16位精度
model, optimizers = amp.initialize(model, optimizers, opt_level='O2')
# 當進行.backward()時，讓amp處理以便它可以對損失進行縮放
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()apex庫會處理大部分工作，包括梯度縮放，以防止梯度爆炸或接近零。
在PyTorch-Lightning中，啟用16位精度不需要修改模型的任何內容，也不需要執行上述代碼。只需在創建Trainer時設置precision=16即可。
trainer = Trainer(amp_level='O2', use_amp=False)
trainer.fit(model)通過這樣的設置，PyTorch-Lightning將自動啟用16位精度，并根據需要進行梯度縮放和其他必要的操作。
移動到多個GPUs中
在多個GPU上進行訓練有幾種方法可以選擇。以下是其中三種常見的方法：
分批次訓練（Batch Splitting）
分批次訓練是指將模型復制到每個GPU上，并將每個GPU中的一部分批次數據進行訓練。
# 在每個GPU上復制模型，并將批次的四分之一分配給每個GPU
model = DataParallel(model, devices=[0, 1, 2, 3])
# out有4個輸出（每個GPU一個）
out = model(x.cuda(0))在PyTorch-Lightning中，只需要增加gpus參數來指定使用的GPU數量，其他的無需更改。
trainer = Trainer(gpus=[0, 1, 2, 3])
trainer.fit(model)模型分布訓練（Model Parallelism）
模型可能太大無法完全放入內存中。例如，帶有編碼器和解碼器的序列到序列模型可能需要占用大量內存。在這種情況下，可以將編碼器和解碼器放在不同的GPU上進行訓練。
# 將編碼器放在GPU 0上，將解碼器放在GPU 1上
encoder_rnn.cuda(0)
decoder_rnn.cuda(1)
# 在GPU 0上運行輸入數據通過編碼器
encoder_out = encoder_rnn(x.cuda(0))
# 在GPU 1上運行輸出通過解碼器
out = decoder_rnn(encoder_out.cuda(1))
# 將輸出數據移回GPU 0上
out = out.cuda(0)在PyTorch-Lightning中，不需要指定任何GPU，只需將模型的模塊放在正確的GPU上即可。
class MyModule(LightningModule):
def __init__():
self.encoder = RNN(...)
self.decoder = RNN(...)
def forward(x):
self.encoder.cuda(0)
self.decoder.cuda(1)
out = self.encoder(x)
out = self.decoder(out.cuda(1))
model = MyModule()
trainer = Trainer()
trainer.fit(model)混合使用（Hybrid Approach）
# 更改這些行
self.encoder = RNN(...)
self.decoder = RNN(...)
# 更改為
# 現在每個RNN都基于不同的GPU設備
self.encoder = DataParallel(self.encoder, devices=[0, 1, 2, 3])
self.decoder = DataParallel(self.encoder, devices=[4, 5, 6, 7])
# 在forward中...
out = self.encoder(x.cuda(0))
# 注意輸入數據放在設備列表中的第一個設備上
out = self.decoder(out.cuda(4)) # <--- 這里的4在使用多個GPU進行訓練時，需要考慮以下注意事項：
- 如果模型已經在GPU上，model.cuda()方法不會執行任何操作。
- 總是將輸入數據放在設備列表中的第一個設備上。
- 在設備之間傳輸數據是昂貴的，應該將其作為最后的手段。
- 優化器和梯度會被保存在GPU 0上，因此，GPU 0上使用的內存可能會比其他GPU大得多。
多節點GPU訓練
在分布式訓練中，每個機器上的每個GPU都有一個模型的副本，并且每個機器都會獲得數據的一部分進行訓練。每個模型副本在其所在的GPU上獨立初始化，并在數據的一個分區上進行訓練。然后，所有模型副本會彼此同步梯度更新。
這種方式可以顯著加快訓練速度，并且使得處理更大規模的數據集成為可能。通過將訓練任務分布到多個機器和GPU上，可以同時進行多個訓練任務，從而節省了訓練時間。
在PyTorch中，可以使用DistributedDataParallel (DDP) 模塊來實現這種分布式訓練方式。它通過在每個節點上復制每個GPU上的模型并同步梯度，實現了模型的并行訓練和梯度更新。代碼如下：
def tng_dataloader():
d = MNIST()
# 4: 創建數據加載器
# 將訓練數據集分發到每個機器上
dist_sampler = DistributedSampler(dataset)
dataloader = DataLoader(d, shuffle=False, sampler=dist_sampler)
def main_process_entrypoint(gpu_nb):
# 2: 設置所有機器和GPU之間的連接
world = nb_gpus * nb_nodes
# 初始化分布式訓練環境，并指定通信后端和當前進程的排名和總共的進程數
dist.init_process_group("nccl", rank=gpu_nb, world_size=world)
# 3: 將模型移動到當前GPU，并使用DistributedDataParallel將模型包裝起來
# DistributedDataParallel會將模型的副本復制到每個GPU上，并確保在訓練過程中同步梯度更新
torch.cuda.set_device(gpu_nb)
model.cuda(gpu_nb)
model = DistributedDataParallel(model, device_ids=[gpu_nb])
# train your model now...
if __name__ == '__main__':
# 1: 生成多個進程，每個進程都會調用main_process_entrypoint()函數。
# 這樣可以在每個機器上啟動多個進程進行分布式訓練
mp.spawn(main_process_entrypoint, nprocs=8)而在Lightning中，分布式訓練變得更加簡單。只需設置節點數量和GPU列表，Trainer類會自動處理剩下的細節。
# train on 1024 gpus across 128 nodes
trainer = Trainer(nb_gpu_nodes=128, gpus=[0, 1, 2, 3, 4, 5, 6, 7])在單個節點上多GPU更快的訓練
使用DistributedDataParallel (DDP) 在單個節點上的多個GPU上進行訓練通常比使用DataParallel 更快。這是因為DDP只執行梯度同步的通信，而不是將整個模型復制到每個GPU上。
在Lightning中，可以通過將distributed_backend參數設置為'ddp'，并指定要使用的GPU數量來輕松實現在單個節點上的多GPU訓練。示例如下：
# 在同一臺機器上的4個GPU上進行訓練，使用DDP比DataParallel更快
trainer = Trainer(distributed_backend='ddp', gpus=[0, 1, 2, 3])在這個例子中，distributed_backend參數被設置為'ddp'來啟用分布式訓練，gpus參數指定要使用的GPU的索引。通過這種方式，Lightning將使用DDP來進行訓練，從而更有效地利用多個GPU。
