也許你仍然在使用32位精度進(jìn)行計(jì)算，或者甚至只是在單個GPU上進(jìn)行訓(xùn)練。

然而，隨著科技的進(jìn)步，我們已經(jīng)有了更好的選擇。使用更高精度的計(jì)算，如16位浮點(diǎn)數(shù)或混合精度，可以提高訓(xùn)練速度并減少內(nèi)存消耗。同時(shí)，利用多個GPU進(jìn)行并行訓(xùn)練，可以大大加快訓(xùn)練過程。

筆者在這里總結(jié)了提升Pytorch模型訓(xùn)練速度的9個技巧，與大家分享~

這些優(yōu)化技巧可以在PyTorch-Lightning庫中找到。PyTorch-Lightning是建立在PyTorch之上的一個封裝，它提供了自動化訓(xùn)練的功能，同時(shí)允許開發(fā)者完全控制關(guān)鍵的模型組件。

這里以MNIST定義LightningModel并使用Trainer來訓(xùn)練模型為例。

#導(dǎo)入PyTorch-Lightning庫中的Trainer類，用于管理訓(xùn)練過程
from pytorch_lightning import Trainer
#創(chuàng)建LightningModule實(shí)例，作為要訓(xùn)練的模型
model = LightningModule(…)
#創(chuàng)建Trainer實(shí)例，用于配置和管理訓(xùn)練過程
trainer = Trainer()
#開始訓(xùn)練
trainer.fit(model)

trainer.fit()方法將根據(jù)LightningModule中定義的訓(xùn)練邏輯來執(zhí)行訓(xùn)練步驟，并自動處理數(shù)據(jù)加載、優(yōu)化器配置、學(xué)習(xí)率調(diào)度等細(xì)節(jié)。

使用DataLoaders

使用DataLoaders來加載數(shù)據(jù)是獲得訓(xùn)練速度提升的最簡單方法之一。通過保存h5py或numpy文件以加速數(shù)據(jù)加載的時(shí)代已經(jīng)過去了，而現(xiàn)在可以「使用PyTorch的DataLoader來輕松加載圖像數(shù)據(jù)」（對于NLP數(shù)據(jù)，請參考TorchText庫）。

在PyTorch-Lightning中，不需要顯式地編寫訓(xùn)練循環(huán)，只需要定義好DataLoaders和Trainer，PyTorch-Lightning會在需要的時(shí)候自動調(diào)用它們。

下面是一個加載MNIST數(shù)據(jù)集并使用DataLoader進(jìn)行訓(xùn)練的示例代碼：

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)
    # 其他訓(xùn)練邏輯

在這個示例中，首先創(chuàng)建了一個MNIST數(shù)據(jù)集的實(shí)例，然后使用DataLoader將其封裝成一個可迭代的數(shù)據(jù)加載器。在訓(xùn)練循環(huán)中，可以遍歷DataLoader，每次獲取一個batch的數(shù)據(jù)，并將其傳遞給模型的training_step()方法進(jìn)行訓(xùn)練。

通過使用DataLoaders，可以更高效地加載和處理大量的訓(xùn)練數(shù)據(jù)，從而提高訓(xùn)練速度。此外，DataLoader還支持?jǐn)?shù)據(jù)的隨機(jī)打亂（shuffle）、批量大?。╞atch_size）等參數(shù)設(shè)置，可以根據(jù)實(shí)際需求進(jìn)行調(diào)整。

DataLoaders中設(shè)置num_workers參數(shù)

在DataLoaders中，可以通過設(shè)置num_workers參數(shù)來允許批量并行加載數(shù)據(jù)，從而加速訓(xùn)練過程。下面是一個示例代碼：

# 慢的方式
loader = DataLoader(dataset, batch_size=32, shuffle=True)

# 快的方式（使用10個workers）
loader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=10)

在第一個示例中，創(chuàng)建了一個DataLoader，沒有指定num_workers參數(shù)，這意味著數(shù)據(jù)加載將在主進(jìn)程中進(jìn)行，而不會并行化。

在第二個示例中，通過將num_workers設(shè)置為10，啟用了批量并行加載。這意味著數(shù)據(jù)加載將在10個worker進(jìn)程中進(jìn)行，并行地加載多個batch，從而加速數(shù)據(jù)加載的過程。

通過適當(dāng)設(shè)置num_workers參數(shù)，可以根據(jù)系統(tǒng)的硬件和資源情況，選擇合適的worker數(shù)量來提高數(shù)據(jù)加載的效率。然而，需要注意的是，并不是worker越多越好，過多的worker可能會導(dǎo)致資源競爭和性能下降。

Batch size

在進(jìn)行下一個優(yōu)化步驟之前，增加批量大?。╞atch size）到CPU-RAM或GPU-RAM允許的最大范圍是一個重要的優(yōu)化策略。

增加批量大小可以帶來以下好處：

• 更高效地利用計(jì)算資源，尤其是GPU的并行計(jì)算能力。
• 減少數(shù)據(jù)加載和傳輸?shù)拇螖?shù)，提高訓(xùn)練速度。
• 可以獲得更穩(wěn)定的梯度估計(jì)，有助于模型收斂。

然而，增加批量大小也會帶來一些挑戰(zhàn)：

• 內(nèi)存占用增加：較大的批量大小需要更多的內(nèi)存空間來存儲數(shù)據(jù)和梯度。
• 學(xué)習(xí)率調(diào)整：增加批量大小后，通常需要相應(yīng)地增加學(xué)習(xí)率，以保持相似的收斂行為。

因此，在增加批量大小之前，需要確保你的硬件和資源可以支持更大的批量大小，并相應(yīng)地調(diào)整學(xué)習(xí)率。

梯度累積

梯度累積（Gradient Accumulation）是一種在計(jì)算資源有限的情況下，模擬較大批量大小的技術(shù)。通過多次執(zhí)行前向傳播、反向傳播和優(yōu)化步驟，將梯度累積起來，以獲得與較大批量大小相同的效果。

下面是一個使用梯度累積的示例代碼：

# 清除上一步的梯度
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()

# 更新權(quán)重
optimizer.step()

# 損失值現(xiàn)在按累積批次數(shù)量進(jìn)行縮放
actual_loss = scaled_loss / 16

在這個示例中，通過循環(huán)執(zhí)行16個梯度累積步驟，每個步驟進(jìn)行前向傳播、計(jì)算損失、反向傳播和梯度累積。然后調(diào)用optimizer.step()來更新權(quán)重。

在PyTorch-Lightning中，只需要設(shè)置accumulate_grad_batches參數(shù)來指定梯度累積的次數(shù)。例如：

trainer = Trainer(accumulate_grad_batches=16)
trainer.fit(model)

保留的計(jì)算圖

在記錄損失值時(shí)，為了避免撐爆內(nèi)存，只存儲損失的數(shù)值而不是整個計(jì)算圖?？梢允褂?item()方法來獲取損失的數(shù)值。

# 方式1
losses.append(loss)

# 方式2
losses.append(loss.item())

在方式1中，損失值loss會保留整個計(jì)算圖的副本，這會占用大量的內(nèi)存空間。而方式2中，使用loss.item()來獲取損失的數(shù)值，并將其存儲到列表中，這樣就只保留了數(shù)值，而不會占用過多的內(nèi)存。

PyTorch-Lightning會非常小心地確保不會保留計(jì)算圖的副本，盡量減少內(nèi)存的占用。因此，在使用PyTorch-Lightning時(shí)，可以放心地使用.item()方法來獲取損失的數(shù)值，而不必?fù)?dān)心內(nèi)存問題。

單個GPU訓(xùn)練

完成上述步驟之后，即可開始在GPU上進(jìn)行訓(xùn)練。GPU上進(jìn)行訓(xùn)練可以利用多個GPU核心之間的并行計(jì)算，從而加速訓(xùn)練過程。

在進(jìn)行GPU訓(xùn)練時(shí)，需要做兩件事情：

1. 將模型移動到GPU上；
2. 在每次數(shù)據(jù)通過時(shí)將數(shù)據(jù)放到GPU上。

下面是在PyTorch中進(jìn)行GPU訓(xùn)練的示例代碼：

# 將模型放到GPU上
model.cuda()
# 將數(shù)據(jù)放到GPU上
x = x.cuda()
# 在GPU上運(yùn)行
model(x)

如果使用PyTorch-Lightning，幾乎不需要做任何額外的工作，只需要設(shè)置Trainer的gpus參數(shù)來指定要使用的GPU數(shù)量。

# 指定訓(xùn)練的gpu id
trainer = Trainer(gpus=[0])
trainer.fit(model)

在進(jìn)行GPU訓(xùn)練時(shí)，需注意限制CPU和GPU之間的數(shù)據(jù)傳輸次數(shù)。盡量避免頻繁地在CPU和GPU之間復(fù)制數(shù)據(jù)。

此外，還要注意調(diào)用強(qiáng)制GPU同步的操作，如清空內(nèi)存緩存torch.cuda.empty_cache()。這樣的操作會阻塞所有GPU，直到它們都完成同步。

然而，如果使用PyTorch-Lightning，則通常不需要擔(dān)心這些問題。PyTorch-Lightning會小心地處理GPU同步和內(nèi)存管理，以確保高效的訓(xùn)練過程。

使用16-bit精度

使用16-bit精度是一種驚人的技術(shù)，可以將內(nèi)存占用減半。大多數(shù)模型通常使用32位精度進(jìn)行訓(xùn)練，但是研究表明，使用16位精度的模型也可以表現(xiàn)得很好?；旌暇葎t意味著在某些部分使用16位精度，但將權(quán)重等內(nèi)容保持在32位精度。

要在PyTorch中使用16位精度，可以安裝NVIDIA的apex庫，并對模型進(jìn)行如下更改：

# 在模型和優(yōu)化器上啟用16位精度
model, optimizers = amp.initialize(model, optimizers, opt_level='O2')

# 當(dāng)進(jìn)行.backward()時(shí)，讓amp處理以便它可以對損失進(jìn)行縮放
with amp.scale_loss(loss, optimizer) as scaled_loss:
    scaled_loss.backward()

apex庫會處理大部分工作，包括梯度縮放，以防止梯度爆炸或接近零。

在PyTorch-Lightning中，啟用16位精度不需要修改模型的任何內(nèi)容，也不需要執(zhí)行上述代碼。只需在創(chuàng)建Trainer時(shí)設(shè)置precision=16即可。

trainer = Trainer(amp_level='O2', use_amp=False)
trainer.fit(model)

通過這樣的設(shè)置，PyTorch-Lightning將自動啟用16位精度，并根據(jù)需要進(jìn)行梯度縮放和其他必要的操作。

移動到多個GPUs中

在多個GPU上進(jìn)行訓(xùn)練有幾種方法可以選擇。以下是其中三種常見的方法：

分批次訓(xùn)練（Batch Splitting）

分批次訓(xùn)練是指將模型復(fù)制到每個GPU上，并將每個GPU中的一部分批次數(shù)據(jù)進(jìn)行訓(xùn)練。

# 在每個GPU上復(fù)制模型，并將批次的四分之一分配給每個GPU
model = DataParallel(model, devices=[0, 1, 2, 3])

# out有4個輸出（每個GPU一個）
out = model(x.cuda(0))

在PyTorch-Lightning中，只需要增加gpus參數(shù)來指定使用的GPU數(shù)量，其他的無需更改。

trainer = Trainer(gpus=[0, 1, 2, 3])
trainer.fit(model)

模型分布訓(xùn)練（Model Parallelism）

模型可能太大無法完全放入內(nèi)存中。例如，帶有編碼器和解碼器的序列到序列模型可能需要占用大量內(nèi)存。在這種情況下，可以將編碼器和解碼器放在不同的GPU上進(jìn)行訓(xùn)練。

# 將編碼器放在GPU 0上，將解碼器放在GPU 1上
encoder_rnn.cuda(0)
decoder_rnn.cuda(1)

# 在GPU 0上運(yùn)行輸入數(shù)據(jù)通過編碼器
encoder_out = encoder_rnn(x.cuda(0))

# 在GPU 1上運(yùn)行輸出通過解碼器
out = decoder_rnn(encoder_out.cuda(1))

# 將輸出數(shù)據(jù)移回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(...)

# 更改為
# 現(xiàn)在每個RNN都基于不同的GPU設(shè)備
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))

# 注意輸入數(shù)據(jù)放在設(shè)備列表中的第一個設(shè)備上
out = self.decoder(out.cuda(4))  # <--- 這里的4

在使用多個GPU進(jìn)行訓(xùn)練時(shí)，需要考慮以下注意事項(xiàng)：

• 如果模型已經(jīng)在GPU上，model.cuda()方法不會執(zhí)行任何操作。
• 總是將輸入數(shù)據(jù)放在設(shè)備列表中的第一個設(shè)備上。
• 在設(shè)備之間傳輸數(shù)據(jù)是昂貴的，應(yīng)該將其作為最后的手段。
• 優(yōu)化器和梯度會被保存在GPU 0上，因此，GPU 0上使用的內(nèi)存可能會比其他GPU大得多。

多節(jié)點(diǎn)GPU訓(xùn)練

在分布式訓(xùn)練中，每個機(jī)器上的每個GPU都有一個模型的副本，并且每個機(jī)器都會獲得數(shù)據(jù)的一部分進(jìn)行訓(xùn)練。每個模型副本在其所在的GPU上獨(dú)立初始化，并在數(shù)據(jù)的一個分區(qū)上進(jìn)行訓(xùn)練。然后，所有模型副本會彼此同步梯度更新。

這種方式可以顯著加快訓(xùn)練速度，并且使得處理更大規(guī)模的數(shù)據(jù)集成為可能。通過將訓(xùn)練任務(wù)分布到多個機(jī)器和GPU上，可以同時(shí)進(jìn)行多個訓(xùn)練任務(wù)，從而節(jié)省了訓(xùn)練時(shí)間。

在PyTorch中，可以使用DistributedDataParallel (DDP) 模塊來實(shí)現(xiàn)這種分布式訓(xùn)練方式。它通過在每個節(jié)點(diǎn)上復(fù)制每個GPU上的模型并同步梯度，實(shí)現(xiàn)了模型的并行訓(xùn)練和梯度更新。代碼如下：

def tng_dataloader():
     d = MNIST()
     # 4: 創(chuàng)建數(shù)據(jù)加載器
     # 將訓(xùn)練數(shù)據(jù)集分發(fā)到每個機(jī)器上
     dist_sampler = DistributedSampler(dataset)
     dataloader = DataLoader(d, shuffle=False, sampler=dist_sampler)

def main_process_entrypoint(gpu_nb):
     # 2: 設(shè)置所有機(jī)器和GPU之間的連接
     world = nb_gpus * nb_nodes
     # 初始化分布式訓(xùn)練環(huán)境，并指定通信后端和當(dāng)前進(jìn)程的排名和總共的進(jìn)程數(shù)
     dist.init_process_group("nccl", rank=gpu_nb, world_size=world)

     # 3: 將模型移動到當(dāng)前GPU，并使用DistributedDataParallel將模型包裝起來
     # DistributedDataParallel會將模型的副本復(fù)制到每個GPU上，并確保在訓(xùn)練過程中同步梯度更新
     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: 生成多個進(jìn)程，每個進(jìn)程都會調(diào)用main_process_entrypoint()函數(shù)。
     # 這樣可以在每個機(jī)器上啟動多個進(jìn)程進(jìn)行分布式訓(xùn)練
     mp.spawn(main_process_entrypoint, nprocs=8)

而在Lightning中，分布式訓(xùn)練變得更加簡單。只需設(shè)置節(jié)點(diǎn)數(shù)量和GPU列表，Trainer類會自動處理剩下的細(xì)節(jié)。

# train on 1024 gpus across 128 nodes
trainer = Trainer(nb_gpu_nodes=128, gpus=[0, 1, 2, 3, 4, 5, 6, 7])

在單個節(jié)點(diǎn)上多GPU更快的訓(xùn)練

使用DistributedDataParallel (DDP) 在單個節(jié)點(diǎn)上的多個GPU上進(jìn)行訓(xùn)練通常比使用DataParallel 更快。這是因?yàn)镈DP只執(zhí)行梯度同步的通信，而不是將整個模型復(fù)制到每個GPU上。

在Lightning中，可以通過將distributed_backend參數(shù)設(shè)置為'ddp'，并指定要使用的GPU數(shù)量來輕松實(shí)現(xiàn)在單個節(jié)點(diǎn)上的多GPU訓(xùn)練。示例如下：

# 在同一臺機(jī)器上的4個GPU上進(jìn)行訓(xùn)練，使用DDP比DataParallel更快
trainer = Trainer(distributed_backend='ddp', gpus=[0, 1, 2, 3])

在這個例子中，distributed_backend參數(shù)被設(shè)置為'ddp'來啟用分布式訓(xùn)練，gpus參數(shù)指定要使用的GPU的索引。通過這種方式，Lightning將使用DDP來進(jìn)行訓(xùn)練，從而更有效地利用多個GPU。