Feb 23, 2023
by Sebastian Raschka

這篇文章介紹了在不影響其準(zhǔn)確性的情況下提高 PyTorch 模型訓(xùn)練性能的技術(shù)。我們將在 LightningModule 中包裝一個(gè) PyTorch 模型，并使用 Trainer 類來(lái)啟用各種訓(xùn)練優(yōu)化。 通過(guò)動(dòng)動(dòng)小手指改幾行代碼，便可在單個(gè) GPU 上的訓(xùn)練時(shí)間從 22.53 分鐘壓榨到 2.75 分鐘，關(guān)鍵是保持模型精度不垮。

這是 8 倍的性能提升！真香！

這篇博文于 03/17/2023 更新，現(xiàn)在使用 PyTorch 2.0 和 Lightning 2.0！

介紹 在本教程中，我們將微調(diào) DistilBERT 模型，它是 BERT 的精煉版本，在幾乎相同的預(yù)測(cè)性能下縮小了 40%。 我們可以通過(guò)多種方式微調(diào)預(yù)訓(xùn)練語(yǔ)言模型。 下圖描述了三種最常見(jiàn)的方法。

上述所有三種方法 (a-c) 都假設(shè)我們已經(jīng)使用自監(jiān)督學(xué)習(xí)在未標(biāo)記的數(shù)據(jù)集上對(duì)模型進(jìn)行了預(yù)訓(xùn)練。 然后，在第 2步中，當(dāng)我們將模型轉(zhuǎn)移到目標(biāo)任務(wù)時(shí)，我們要么

a) 提取嵌入并在其上訓(xùn)練分類器（例如，這可以是來(lái)自 scikit-learn 的支持向量機(jī)）； b) 替換/添加輸出層并微調(diào)transformer的最后一層； c) 替換/添加輸出層并微調(diào)所有層。 方法 a-c 按計(jì)算效率排序，其中 a) 通常是最快的。 根據(jù)經(jīng)驗(yàn)，這種排序順序也反映了模型的預(yù)測(cè)性能，其中 c) 通常會(huì)產(chǎn)生最高的預(yù)測(cè)精度。

在本文中，我們將使用 c) 訓(xùn)練一個(gè)模型來(lái)預(yù)測(cè) IMDB 電影評(píng)論數(shù)據(jù)集中的電影評(píng)論情緒，該數(shù)據(jù)集總共包含 50,000 條電影評(píng)論。

1）普通 PyTorch 基線

我們先從簡(jiǎn)單的 PyTorch 基線開(kāi)始，在 IMDB 電影評(píng)論數(shù)據(jù)集上訓(xùn)練 DistilBERT 模型。 如果你想自己運(yùn)行代碼，你可以conda一個(gè)虛擬環(huán)境，如下所示：

相關(guān)軟件版本如下：

Python version: 3.9.15torch         : 2.0.0+cu118lightning     : 2.0.0transformers  : 4.26.1

載數(shù)據(jù)集代碼的
local_dataset_utilities.py 文件。

在我們下面討論之前先看看 主要的 PyTorch 代碼：

import osimport os.path as opimport timefrom datasets import load_datasetimport torchfrom torch.utils.data import DataLoaderimport torchmetricsfrom transformers import AutoTokenizerfrom transformers import AutoModelForSequenceClassificationfrom watermark import watermarkfrom local_dataset_utilities import (   download_dataset,   load_dataset_into_to_dataframe,   partition_dataset,)from local_dataset_utilities import IMDBDatasetdef tokenize_text(batch):   return tokenizer(batch['text'], truncation=True, padding=True)def train(num_epochs, model, optimizer, train_loader, val_loader, device):   for epoch in range(num_epochs):       train_acc = torchmetrics.Accuracy(task='multiclass', num_classes=2).to(device)       for batch_idx, batch in enumerate(train_loader):           model.train()           for s in ['input_ids', 'attention_mask', 'label']:               batch[s] = batch[s].to(device)           ### FORWARD AND BACK PROP           outputs = model(               batch['input_ids'],               attention_mask=batch['attention_mask'],               labels=batch['label'],           )           optimizer.zero_grad()           outputs['loss'].backward()           ### UPDATE MODEL PARAMETERS           optimizer.step()           ### LOGGING           if not batch_idx % 300:               print(                   f'Epoch: {epoch+1:04d}/{num_epochs:04d} | Batch {batch_idx:04d}/{len(train_loader):04d} | Loss: {outputs['loss']:.4f}'               )           model.eval()           with torch.no_grad():               predicted_labels = torch.argmax(outputs['logits'], 1)               train_acc.update(predicted_labels, batch['label'])       ### MORE LOGGING       with torch.no_grad():           model.eval()           val_acc = torchmetrics.Accuracy(task='multiclass', num_classes=2).to(device)           for batch in val_loader:               for s in ['input_ids', 'attention_mask', 'label']:                   batch[s] = batch[s].to(device)               outputs = model(                   batch['input_ids'],                   attention_mask=batch['attention_mask'],                   labels=batch['label'],               )               predicted_labels = torch.argmax(outputs['logits'], 1)               val_acc.update(predicted_labels, batch['label'])           print(               f'Epoch: {epoch+1:04d}/{num_epochs:04d} | Train acc.: {train_acc.compute()*100:.2f}% | Val acc.: {val_acc.compute()*100:.2f}%'           )   print(watermark(packages='torch,lightning,transformers', python=True))   print('Torch CUDA available?', torch.cuda.is_available())   device = 'cuda:0' if torch.cuda.is_available() else 'cpu'   torch.manual_seed(123)   ##########################   ### 1 Loading the Dataset   ##########################   download_dataset()   df = load_dataset_into_to_dataframe()   if not (op.exists('train.csv') and op.exists('val.csv') and op.exists('test.csv')):       partition_dataset(df)   imdb_dataset = load_dataset(       'csv',       data_files={           'train': 'train.csv',           'validation': 'val.csv',           'test': 'test.csv',       },   )   #########################################   ### 2 Tokenization and Numericalization   #########################################   tokenizer = AutoTokenizer.from_pretrained('distilbert-base-uncased')   print('Tokenizer input max length:', tokenizer.model_max_length, flush=True)   print('Tokenizer vocabulary size:', tokenizer.vocab_size, flush=True)   print('Tokenizing ...', flush=True)   imdb_tokenized = imdb_dataset.map(tokenize_text, batched=True, batch_size=None)   del imdb_dataset   imdb_tokenized.set_format('torch', columns=['input_ids', 'attention_mask', 'label'])   os.environ['TOKENIZERS_PARALLELISM'] = 'false'   #########################################   ### 3 Set Up DataLoaders   #########################################   train_dataset = IMDBDataset(imdb_tokenized, partition_key='train')   val_dataset = IMDBDataset(imdb_tokenized, partition_key='validation')   test_dataset = IMDBDataset(imdb_tokenized, partition_key='test')   train_loader = DataLoader(       dataset=train_dataset,       batch_size=12,       shuffle=True,       num_workers=1,       drop_last=True,   )   val_loader = DataLoader(       dataset=val_dataset,       batch_size=12,       num_workers=1,       drop_last=True,   )   test_loader = DataLoader(       dataset=test_dataset,       batch_size=12,       num_workers=1,       drop_last=True,   )   #########################################   ### 4 Initializing the Model   #########################################   model = AutoModelForSequenceClassification.from_pretrained(       'distilbert-base-uncased', num_labels=2   )   model.to(device)   optimizer = torch.optim.Adam(model.parameters(), lr=5e-5)   #########################################   ### 5 Finetuning   #########################################   start = time.time()   train(       num_epochs=3,       model=model,       optimizer=optimizer,       train_loader=train_loader,       val_loader=val_loader,       device=device,   )   end = time.time()   elapsed = end - start   print(f'Time elapsed {elapsed/60:.2f} min')   with torch.no_grad():       model.eval()       test_acc = torchmetrics.Accuracy(task='multiclass', num_classes=2).to(device)       for batch in test_loader:           for s in ['input_ids', 'attention_mask', 'label']:               batch[s] = batch[s].to(device)           outputs = model(               batch['input_ids'],               attention_mask=batch['attention_mask'],               labels=batch['label'],           )           predicted_labels = torch.argmax(outputs['logits'], 1)           test_acc.update(predicted_labels, batch['label'])   print(f'Test accuracy {test_acc.compute()*100:.2f}%')

上面的代碼結(jié)構(gòu)分為兩部分，函數(shù)定義和在 if name == 'main' 下執(zhí)行的代碼。 這個(gè)結(jié)構(gòu)對(duì)于避免以后使用多個(gè) GPU 時(shí) Python 的多處理問(wèn)題是必要的。

if name == 'main' 中的的前三部分包含設(shè)置數(shù)據(jù)集加載器的代碼。 第四部分是初始化模型的地方。 第五部分運(yùn)行訓(xùn)練函數(shù)并在測(cè)試集上評(píng)估微調(diào)模型。

在 A100 GPU 上運(yùn)行代碼后，我得到了以下結(jié)果：

正如所見(jiàn)，模型從第 2 輪到第 3 輪開(kāi)始略微過(guò)度擬合，驗(yàn)證準(zhǔn)確率從 92.09% 下降到 89.88%。 最終測(cè)試準(zhǔn)確率為 89.92%，這是對(duì)模型進(jìn)行 21.33 分鐘微調(diào)后達(dá)到的。

2） 使用Trainer Class

現(xiàn)在，讓我們將 PyTorch 模型包裝在 LightningModule 中，以便我們可以使用來(lái)自 Lightning 的 Trainer 類：

import osimport os.path as opimport timefrom datasets import load_datasetimport lightning as Lfrom lightning.pytorch.callbacks import ModelCheckpointfrom lightning.pytorch.loggers import CSVLoggerimport matplotlib.pyplot as pltimport pandas as pdimport torchfrom torch.utils.data import DataLoaderimport torchmetricsfrom transformers import AutoTokenizerfrom transformers import AutoModelForSequenceClassificationfrom watermark import watermarkfrom local_dataset_utilities import (    download_dataset,    load_dataset_into_to_dataframe,    partition_dataset,)from local_dataset_utilities import IMDBDatasetdef tokenize_text(batch):    return tokenizer(batch['text'], truncation=True, padding=True)class LightningModel(L.LightningModule):    def __init__(self, model, learning_rate=5e-5):        super().__init__()        self.learning_rate = learning_rate        self.model = model        self.train_acc = torchmetrics.Accuracy(task='multiclass', num_classes=2)        self.val_acc = torchmetrics.Accuracy(task='multiclass', num_classes=2)        self.test_acc = torchmetrics.Accuracy(task='multiclass', num_classes=2)    def forward(self, input_ids, attention_mask, labels):        return self.model(input_ids, attention_mask=attention_mask, labels=labels)    def training_step(self, batch, batch_idx):        outputs = self(            batch['input_ids'],            attention_mask=batch['attention_mask'],            labels=batch['label'],        )        self.log('train_loss', outputs['loss'])        with torch.no_grad():            logits = outputs['logits']            predicted_labels = torch.argmax(logits, 1)            self.train_acc(predicted_labels, batch['label'])            self.log('train_acc', self.train_acc, on_epoch=True, on_step=False)        return outputs['loss']  # this is passed to the optimizer for training    def validation_step(self, batch, batch_idx):        outputs = self(            batch['input_ids'],            attention_mask=batch['attention_mask'],            labels=batch['label'],        )        self.log('val_loss', outputs['loss'], prog_bar=True)        logits = outputs['logits']        predicted_labels = torch.argmax(logits, 1)        self.val_acc(predicted_labels, batch['label'])        self.log('val_acc', self.val_acc, prog_bar=True)    def test_step(self, batch, batch_idx):        outputs = self(            batch['input_ids'],            attention_mask=batch['attention_mask'],            labels=batch['label'],        )        logits = outputs['logits']        predicted_labels = torch.argmax(logits, 1)        self.test_acc(predicted_labels, batch['label'])        self.log('accuracy', self.test_acc, prog_bar=True)    def configure_optimizers(self):        optimizer = torch.optim.Adam(            self.trainer.model.parameters(), lr=self.learning_rate        )        return optimizerif __name__ == '__main__':    print(watermark(packages='torch,lightning,transformers', python=True), flush=True)    print('Torch CUDA available?', torch.cuda.is_available(), flush=True)    torch.manual_seed(123)    ##########################    ### 1 Loading the Dataset    ##########################    download_dataset()    df = load_dataset_into_to_dataframe()    if not (op.exists('train.csv') and op.exists('val.csv') and op.exists('test.csv')):        partition_dataset(df)    imdb_dataset = load_dataset(        'csv',        data_files={            'train': 'train.csv',            'validation': 'val.csv',            'test': 'test.csv',        },    )    #########################################    ### 2 Tokenization and Numericalization    ########################################    tokenizer = AutoTokenizer.from_pretrained('distilbert-base-uncased')    print('Tokenizer input max length:', tokenizer.model_max_length, flush=True)    print('Tokenizer vocabulary size:', tokenizer.vocab_size, flush=True)    print('Tokenizing ...', flush=True)    imdb_tokenized = imdb_dataset.map(tokenize_text, batched=True, batch_size=None)    del imdb_dataset    imdb_tokenized.set_format('torch', columns=['input_ids', 'attention_mask', 'label'])    os.environ['TOKENIZERS_PARALLELISM'] = 'false'    #########################################    ### 3 Set Up DataLoaders    #########################################    train_dataset = IMDBDataset(imdb_tokenized, partition_key='train')    val_dataset = IMDBDataset(imdb_tokenized, partition_key='validation')    test_dataset = IMDBDataset(imdb_tokenized, partition_key='test')    train_loader = DataLoader(        dataset=train_dataset,        batch_size=12,        shuffle=True,        num_workers=1,        drop_last=True,    )    val_loader = DataLoader(        dataset=val_dataset,        batch_size=12,        num_workers=1,        drop_last=True,    )    test_loader = DataLoader(        dataset=test_dataset,        batch_size=12,        num_workers=1,        drop_last=True,    )    #########################################    ### 4 Initializing the Model    #########################################    model = AutoModelForSequenceClassification.from_pretrained(        'distilbert-base-uncased', num_labels=2    )    #########################################    ### 5 Finetuning    #########################################    lightning_model = LightningModel(model)    callbacks = [        ModelCheckpoint(save_top_k=1, mode='max', monitor='val_acc')  # save top 1 model    ]    logger = CSVLogger(save_dir='logs/', name='my-model')    trainer = L.Trainer(        max_epochs=3,        callbacks=callbacks,        accelerator='gpu',        devices=[1],        logger=logger,        log_every_n_steps=10,        deterministic=True,    )    start = time.time()    trainer.fit(        model=lightning_model,        train_dataloaders=train_loader,        val_dataloaders=val_loader,    )    end = time.time()    elapsed = end - start    print(f'Time elapsed {elapsed/60:.2f} min')    test_acc = trainer.test(lightning_model, dataloaders=test_loader, ckpt_path='best')    print(test_acc)    with open(op.join(trainer.logger.log_dir, 'outputs.txt'), 'w') as f:        f.write((f'Time elapsed {elapsed/60:.2f} min\n'))        f.write(f'Test acc: {test_acc}')

本文關(guān)注性能方面 跳過(guò) LightningModule 的細(xì)節(jié)。

簡(jiǎn)而言之，設(shè)置了一個(gè) LightningModule，它定義了如何執(zhí)行訓(xùn)練、驗(yàn)證和測(cè)試步驟。 然后，主要變化在代碼第5部分，我們?cè)谄渲形⒄{(diào)模型。 將 PyTorch 模型包裝在 LightningModel 類中，并使用 Trainer 類來(lái)擬合模型：

之前注意到驗(yàn)證準(zhǔn)確率從第 2 輪下降到第 3 輪，因此使用 ModelCheckpoint 回調(diào)加載最佳模型（基于最高驗(yàn)證準(zhǔn)確率）以在測(cè)試集上進(jìn)行模型評(píng)估。 此外，將性能記錄到 CSV 文件并將 PyTorch 行為設(shè)置為確定性。

在同一臺(tái)機(jī)器上，這個(gè)模型在 21.79 分鐘內(nèi)達(dá)到了 92.6% 的測(cè)試準(zhǔn)確率：

注意，如果禁用檢查點(diǎn)并允許 PyTorch 在非確定性模式下運(yùn)行，將獲得與普通 PyTorch 相同的運(yùn)行時(shí)間。

3）自動(dòng)混合精度訓(xùn)練

如果我們的 GPU 支持混合精度訓(xùn)練，啟用它通常是提高計(jì)算效率的主要方法之一。 特別是在訓(xùn)練期間在 32 位和 16 位浮點(diǎn)表示之間切換，而不會(huì)犧牲準(zhǔn)確性。

使用 Trainer 類，可以通過(guò)一行代碼啟用自動(dòng)混合精度訓(xùn)練：

 trainer = L.Trainer(        max_epochs=3,        callbacks=callbacks,        accelerator='gpu',        precision='16',  # <-- NEW        devices=[1],        logger=logger,        log_every_n_steps=10,        deterministic=True,    )

如下圖所示，使用混合精度訓(xùn)練可將訓(xùn)練時(shí)間從 21.79 分鐘提高到 8.25 分鐘！ 這幾乎快了3倍！

測(cè)試集準(zhǔn)確率為 93.2%——與之前的 92.6% 相比甚至略有提高（可能是由于在不同精度模式之間切換時(shí)舍入引起的差異。）

4) 使用 Torch.Compile 的靜態(tài)圖

在最近發(fā)布的 PyTorch 2.0 中，PyTorch 團(tuán)隊(duì)引入了新的 toch.compile 函數(shù)，該函數(shù)可以通過(guò)生成優(yōu)化的靜態(tài)圖來(lái)加速 PyTorch 代碼執(zhí)行。 這是一個(gè) 3 步過(guò)程，包括圖形獲取、圖形結(jié)構(gòu)降低和圖形編譯。

實(shí)現(xiàn)這一目標(biāo)的背后很復(fù)雜，在PyTorch 2.0相關(guān)介紹中有更詳細(xì)的解釋。 作為用戶，我們可以通過(guò)一個(gè)簡(jiǎn)單的命令 torch.compile 使用這一新功能。

要利用 torch.compile，可以通過(guò)添加以下一行代碼來(lái)修改我們的代碼：

不幸的是，在這種混合精度上下文中，使用默認(rèn)參數(shù)時(shí)，torch.compile 似乎不會(huì)提升 DistilBERT 模型的性能。 訓(xùn)練時(shí)間為 8.44 分鐘，而之前為 8.25 分鐘。 因此，本文中的后續(xù)基準(zhǔn)測(cè)試不會(huì)使用 torch.compile。

備注：兩個(gè)技巧:

1.將編譯放在計(jì)時(shí)開(kāi)始之前； 2。使用示例批次啟動(dòng)模型，如下所示

model.to(torch.device('cuda:0'))  model = torch.compile(model)  for batch_idx, batch in enumerate(train_loader):      model.train()      for s in ['input_ids', 'attention_mask', 'label']:          batch[s] = batch[s].to(torch.device('cuda:0'))      break  outputs = model(      batch['input_ids'],      attention_mask=batch['attention_mask'],      labels=batch['label'],  )  lightning_model = LightningModel(model)  # start timing and training below

運(yùn)行時(shí)間提高到 5.6 分鐘。 這表明初始優(yōu)化編譯步驟需要幾分鐘，但最終會(huì)加速模型訓(xùn)練。 在這種情況下，由于我們只訓(xùn)練三個(gè)時(shí)期的模型，因此由于額外的開(kāi)銷，編譯的好處不明顯。 但是，如果訓(xùn)練模型的時(shí)間更長(zhǎng)或訓(xùn)練的模型更大，那么編譯是值得的。

（注意：目前為分布式設(shè)置準(zhǔn)備模型有難搞，因?yàn)槊總€(gè)單獨(dú)的 GPU 設(shè)備都需要模型的副本。這將需要重新設(shè)計(jì)一些代碼，所以下面不會(huì)使用 torch.compile。）

5) 在4個(gè)GPU上并行訓(xùn)練分布式數(shù)據(jù)

上面添加混合精度訓(xùn)練（并嘗試添加圖形編譯）已在單個(gè) GPU 上加速我們的代碼，現(xiàn)在讓嘗試多 GPU 策略。 現(xiàn)在將在四個(gè)而不是一個(gè) GPU 上運(yùn)行相同的代碼。

請(qǐng)注意，下圖中總結(jié)了幾種不同的多 GPU 訓(xùn)練技術(shù)。

從最簡(jiǎn)單的技術(shù)開(kāi)始，通過(guò) DistributedDataParallel 實(shí)現(xiàn)數(shù)據(jù)并行。 使用Trainer，只需要修改一行代碼：

在有四個(gè) A100 GPU機(jī)器上，這段代碼運(yùn)行了 3.07 分鐘，達(dá)到了 93.1% 的測(cè)試準(zhǔn)確率。 同樣，測(cè)試集的改進(jìn)可能是由于使用數(shù)據(jù)并行時(shí)的梯度平均。

6）DeepSpeed

最后，嘗試在 Trainer 中使用的 DeepSpeed 多 GPU 策略。

但在實(shí)際嘗試之前，分享下多 GPU 使用建議。 使用哪種策略在很大程度上取決于模型、GPU 的數(shù)量和 GPU 的內(nèi)存大小。 例如，當(dāng)預(yù)訓(xùn)練模型不適合單個(gè) GPU 的大型模型時(shí)，最好從簡(jiǎn)單的“ddp_sharded”策略開(kāi)始，該策略將張量并行性添加到“ddp”。使用前面的代碼，“ddp_sharded”采用 跑完 2.58 分鐘。

或者也可以考慮更復(fù)雜的“deepspeed_stage_2”策略，它將優(yōu)化器狀態(tài)和梯度分片。 如果這不足以使模型適合 GPU 內(nèi)存，請(qǐng)嘗試“deepspeed_stage_2_offload”變體，它將優(yōu)化器和梯度狀態(tài)卸載到 CPU 內(nèi)存（以性能為代價(jià)）。 如果你想微調(diào)一個(gè)模型，計(jì)算吞吐量通常比能夠?qū)⒛Ｐ头湃胼^少數(shù)量的 GPU 的內(nèi)存中更不重要。 在這種情況下，您可以探索 deepspeed 的“stage_3”變體，它對(duì)所有內(nèi)容、優(yōu)化器、梯度和參數(shù)進(jìn)行分片，等等.

• strategy='deepspeed_stage_3'
• strategy='deepspeed_stage_3_offload' 由于 GPU 內(nèi)存不是像 DistilBERT 這樣小模型的問(wèn)題，讓我們?cè)囋嚒癲eepspeed_stage_2”：

首先，我們必須安裝 DeepSpeed Python 庫(kù)：

pip install -U deepspeed

接下來(lái)，我們只需更改一行代碼即可啟用“deepspeed_stage_2”：

在機(jī)器上運(yùn)行了 2.75 分鐘，并達(dá)到了 92.6% 的測(cè)試準(zhǔn)確率。

請(qǐng)注意，PyTorch 現(xiàn)在也有自己的 DeepSpeed 替代方案，稱為完全分片 DataParallel，我們可以通過(guò) strategy='fsdp' 使用它。

7) Fabric

隨著最近的 Lightning 2.0 發(fā)布，Lightning AI 發(fā)布了用于 PyTorch 的新 Fabric 開(kāi)源庫(kù)。 Fabric 本質(zhì)上是一種擴(kuò)展 PyTorch 代碼的替代方法，無(wú)需使用我在上面第 2 節(jié)）使用 Trainer 類中介紹的 LightningModule 和 Trainer。

Fabric只需要改幾行代碼，如下代碼所示。 - 表示已刪除的行，+ 是為將 Python 代碼轉(zhuǎn)換為使用 Fabric 而添加的行。

import osimport os.path as opimport time+ from lightning import Fabricfrom datasets import load_datasetimport matplotlib.pyplot as pltimport pandas as pdimport torchfrom torch.utils.data import DataLoaderimport torchmetricsfrom transformers import AutoTokenizerfrom transformers import AutoModelForSequenceClassificationfrom watermark import watermarkfrom local_dataset_utilities import download_dataset, load_dataset_into_to_dataframe, partition_datasetfrom local_dataset_utilities import IMDBDatasetdef tokenize_text(batch):    return tokenizer(batch['text'], truncation=True, padding=True)def plot_logs(log_dir):    metrics = pd.read_csv(op.join(log_dir, 'metrics.csv'))    aggreg_metrics = []    agg_col = 'epoch'    for i, dfg in metrics.groupby(agg_col):        agg = dict(dfg.mean())        agg[agg_col] = i        aggreg_metrics.append(agg)    df_metrics = pd.DataFrame(aggreg_metrics)    df_metrics[['train_loss', 'val_loss']].plot(        grid=True, legend=True, xlabel='Epoch', ylabel='Loss'    )    plt.savefig(op.join(log_dir, 'loss.pdf'))    df_metrics[['train_acc', 'val_acc']].plot(        grid=True, legend=True, xlabel='Epoch', ylabel='Accuracy'    )    plt.savefig(op.join(log_dir, 'acc.pdf'))- def train(num_epochs, model, optimizer, train_loader, val_loader, device):+ def train(num_epochs, model, optimizer, train_loader, val_loader, fabric):      for epoch in range(num_epochs):-         train_acc = torchmetrics.Accuracy(task='multiclass', num_classes=2).to(device)+         train_acc = torchmetrics.Accuracy(task='multiclass', num_classes=2).to(fabric.device)        model.train()        for batch_idx, batch in enumerate(train_loader):-             for s in ['input_ids', 'attention_mask', 'label']:-                 batch[s] = batch[s].to(device)            outputs = model(batch['input_ids'], attention_mask=batch['attention_mask'], labels=batch['label'])             optimizer.zero_grad()-            outputs['loss'].backward()+            fabric.backward(outputs['loss'])            ### UPDATE MODEL PARAMETERS            optimizer.step()            ### LOGGING            if not batch_idx % 300:                print(f'Epoch: {epoch+1:04d}/{num_epochs:04d} | Batch {batch_idx:04d}/{len(train_loader):04d} | Loss: {outputs['loss']:.4f}')            model.eval()            with torch.no_grad():                predicted_labels = torch.argmax(outputs['logits'], 1)                train_acc.update(predicted_labels, batch['label'])        ### MORE LOGGING        model.eval()        with torch.no_grad():-            val_acc = torchmetrics.Accuracy(task='multiclass', num_classes=2).to(device)+            val_acc = torchmetrics.Accuracy(task='multiclass', num_classes=2).to(fabric.device)            for batch in val_loader:-                for s in ['input_ids', 'attention_mask', 'label']:-                    batch[s] = batch[s].to(device)                outputs = model(batch['input_ids'], attention_mask=batch['attention_mask'], labels=batch['label'])                predicted_labels = torch.argmax(outputs['logits'], 1)                val_acc.update(predicted_labels, batch['label'])            print(f'Epoch: {epoch+1:04d}/{num_epochs:04d} | Train acc.: {train_acc.compute()*100:.2f}% | Val acc.: {val_acc.compute()*100:.2f}%')            train_acc.reset(), val_acc.reset()if __name__ == '__main__':    print(watermark(packages='torch,lightning,transformers', python=True))    print('Torch CUDA available?', torch.cuda.is_available())    -   device = 'cuda' if torch.cuda.is_available() else 'cpu'    torch.manual_seed(123)    ##########################    ### 1 Loading the Dataset    ##########################    download_dataset()    df = load_dataset_into_to_dataframe()    if not (op.exists('train.csv') and op.exists('val.csv') and op.exists('test.csv')):        partition_dataset(df)    imdb_dataset = load_dataset(        'csv',        data_files={            'train': 'train.csv',            'validation': 'val.csv',            'test': 'test.csv',        },    )    #########################################    ### 2 Tokenization and Numericalization    #########################################    tokenizer = AutoTokenizer.from_pretrained('distilbert-base-uncased')    print('Tokenizer input max length:', tokenizer.model_max_length, flush=True)    print('Tokenizer vocabulary size:', tokenizer.vocab_size, flush=True)    print('Tokenizing ...', flush=True)    imdb_tokenized = imdb_dataset.map(tokenize_text, batched=True, batch_size=None)    del imdb_dataset    imdb_tokenized.set_format('torch', columns=['input_ids', 'attention_mask', 'label'])    os.environ['TOKENIZERS_PARALLELISM'] = 'false'    #########################################    ### 3 Set Up DataLoaders    #########################################    train_dataset = IMDBDataset(imdb_tokenized, partition_key='train')    val_dataset = IMDBDataset(imdb_tokenized, partition_key='validation')    test_dataset = IMDBDataset(imdb_tokenized, partition_key='test')    train_loader = DataLoader(        dataset=train_dataset,        batch_size=12,        shuffle=True,         num_workers=2,        drop_last=True,    )    val_loader = DataLoader(        dataset=val_dataset,        batch_size=12,        num_workers=2,        drop_last=True,    )    test_loader = DataLoader(        dataset=test_dataset,        batch_size=12,        num_workers=2,        drop_last=True,    )    #########################################    ### 4 Initializing the Model    #########################################+    fabric = Fabric(accelerator='cuda', devices=4, +                    strategy='deepspeed_stage_2', precision='16-mixed')+    fabric.launch()    model = AutoModelForSequenceClassification.from_pretrained(        'distilbert-base-uncased', num_labels=2)-   model.to(device)    optimizer = torch.optim.Adam(model.parameters(), lr=5e-5)+    model, optimizer = fabric.setup(model, optimizer)+    train_loader, val_loader, test_loader = fabric.setup_dataloaders(+        train_loader, val_loader, test_loader)    #########################################    ### 5 Finetuning    #########################################    start = time.time()    train(        num_epochs=3,        model=model,        optimizer=optimizer,        train_loader=train_loader,        val_loader=val_loader,-       device=device+       fabric=fabric    )    end = time.time()    elapsed = end-start    print(f'Time elapsed {elapsed/60:.2f} min')    with torch.no_grad():        model.eval()-       test_acc = torchmetrics.Accuracy(task='multiclass', num_classes=2).to(device)+       test_acc = torchmetrics.Accuracy(task='multiclass', num_classes=2).to(fabric.device)        for batch in test_loader:-           for s in ['input_ids', 'attention_mask', 'label']:-               batch[s] = batch[s].to(device)            outputs = model(batch['input_ids'], attention_mask=batch['attention_mask'], labels=batch['label'])            predicted_labels = torch.argmax(outputs['logits'], 1)            test_acc.update(predicted_labels, batch['label'])    print(f'Test accuracy {test_acc.compute()*100:.2f}%')

正如所看到的，修改真的很輕量級(jí)！ 它運(yùn)行得如何呢？ Fabric 僅用了 1.8 分鐘就完成了微調(diào)！ Fabric 比 Trainer 更輕量級(jí)——雖然它也能夠使用回調(diào)和日志記錄，但我們沒(méi)有在這里啟用這些功能來(lái)用一個(gè)極簡(jiǎn)示例來(lái)演示 Fabric。 太快了.

何時(shí)使用 Lightning Trainer 或 Fabric 取決于個(gè)人喜好。 根據(jù)經(jīng)驗(yàn)，如果您更喜歡對(duì) PyTorch 代碼進(jìn)行輕量包裝，請(qǐng)查看 Fabric。 另一方面，如果你轉(zhuǎn)向更大的項(xiàng)目并且更喜歡 Lightning 提供的代碼組織，推薦 Trainer。

結(jié)論

在本文中，探索了各種提高 PyTorch 模型訓(xùn)練速度的技術(shù)。 如果使用 Lightning Trainer，可以用一行代碼在這些選項(xiàng)之間切換，非常方便—，尤其是當(dāng)您在調(diào)試代碼時(shí)在 CPU 和 GPU 機(jī)器之間切換時(shí)。

尚未探索的另一個(gè)方面是最大化批量大小，這可以進(jìn)一步提高我們模型的吞吐量。 未完待續(xù)。

（代碼地址github稍后見(jiàn)評(píng)論）