[PyTorch]：加速Pytorch 模型训练的几种方法（几行代码），最快提升八倍（附实验记录）

原创

已于 2024-07-02 21:05:00 修改 · 1.8k 阅读

20 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #人工智能 #python

于 2024-07-02 21:03:46 首次发布

本篇文章转自：Some Techniques To Make Your PyTorch Models Train (Much) Faster

本篇博文概述了在不影响 PyTorch 模型准确性的情况下提高其训练性能的技术。为此，将 PyTorch 模型包装在 LightningModule 中，并使用 Trainer 类来实现各种训练优化。只需更改几行代码，就可以将单个 GPU 上的训练时间从 22.53 分钟缩短到 2.75 分钟，同时保持模型的预测准确性。

性能提升了 8 倍！

（本博文于 2023 年 3 月 17 日更新，现在使用 PyTorch 2.0 和 Lightning 2.0！）

实验汇总：

省流版（详细实验过程可参见下述实验训练记录）：

使用Pytorch与Pytorch Ligntning可以获得相同的训练时长（上图不一致在于加载checkpoint文件和log）

使用混合精度训练可将训练时间从 21.79 分钟缩短至 8.25 分钟！这几乎快了 3 倍！

测试集准确率为 93.2%——与之前的 92.6% 相比略有提高（可能是由于在不同精度模式之间切换时舍入引起的差异）。

在使用默认参数的情况下，似乎torch.compile不会在混合精度环境中提高 DistilBERT 模型的性能。

使用四块 A100 GPU，此代码运行时间为 3.07 分钟，测试准确率达到 93.1%

Trainer + 混合精度 + DeepSpeed + 多GPU运行花了 2.75 分钟

Fabric + 混合精度 + DeepSpeed + 多GPU运行花了 1.8 分钟

介绍

在本教程中，我们将微调DistilBERT 模型，这是 BERT 的精简版本，其规模缩小了 40%，但预测性能几乎相同。我们可以通过多种方式微调预训练语言模型。下图描述了三种最常见的方法。

上述三种方法都假设我们已经使用自监督学习在未标记的数据集上对模型进行了预训练（步骤 1）。然后，在步骤 2 中，当我们将模型迁移到目标任务时，我们要么

a）提取嵌入并在其上训练分类器（例如，可以是来自 scikit-learn 的支持向量机）；
b）替换/添加输出层并微调 Transformer 的最后一层（几层）；
c）替换/添加输出层并微调所有层。

方法 ac 按计算效率排序，其中 a) 通常是最快的。根据我的经验，这种排序顺序也反映了模型的预测性能，其中 c) 通常具有最高的预测准确率。

在本教程中，我们将使用方法 c) 并训练一个模型来预测总共包含 50,000 条电影评论的IMDB 大型电影评论数据集中的电影评论情绪。

1. 普通的 PyTorch 基线

作为热身练习，让我们从简单的 PyTorch 基线开始，在 IMDB 电影评论数据集上训练 DistilBERT 模型。如果您想自己运行代码，可以使用相关的 Python 库设置虚拟环境，如下所示：

conda create -n faster-blog python=3.9
conda activate faster-blog

pip install watermark transformers datasets torchmetrics lightning

作为参考，我使用的相关软件版本如下（当您在本文后面运行代码时它们将被打印到终端。）：

Python version: 3.9.15
torch         : 2.0.0+cu118
lightning     : 2.0.0
transformers  : 4.26.1

为了避免本文充斥着无聊的数据加载实用程序，我将跳过local_dataset_utilities.py文件，该文件包含用于加载数据集的代码。这里唯一相关的信息是我们将数据集划分为 35,000 个训练示例、5,000 个验证集记录和 10,000 个测试记录。

让我们来看看主要的 PyTorch 代码。除了我放在local_dataset_utilities.py文件中的数据集加载实用程序外，此代码是自包含的。在我们在下面讨论之前，请先查看 PyTorch 代码：

import os
import os.path as op
import time

from datasets import load_dataset
import torch
from torch.utils.data import DataLoader
import torchmetrics
from transformers import AutoTokenizer
from transformers import AutoModelForSequenceClassification
from watermark import watermark

from local_dataset_utilities import (
    download_dataset,
    load_dataset_into_to_dataframe,
    partition_dataset,
)
from local_dataset_utilities import IMDBDataset


def 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=