从 Jupyter Notebooks 启动分布式训练

本教程将指导您如何在分布式系统上，通过 Jupyter Notebook 使用 🤗 Accelerate 对计算机视觉模型进行微调。您还将学习如何设置一些必要的要求，以确保您的环境配置正确、数据准备妥当，并最终了解如何启动训练。

本教程也可作为 Jupyter Notebook 在此处获取

配置环境

在进行任何训练之前，系统中必须存在一个 Accelerate 配置文件。通常，可以在终端中运行以下命令并回答提示来完成此操作：

accelerate config

不过，如果通用默认设置即可，且您并非在 TPU 上运行，Accelerate 提供了一个实用工具，可通过 utils.write_basic_config() 快速将您的设备配置写入配置文件。

以下代码在写入配置后会重启 Jupyter，因为此操作调用了 CUDA 运行时或 XPU 运行时。

在多设备系统上，CUDA 和 XPU 无法被初始化多次。在 notebook 中进行调试并调用 CUDA/XPU 是没问题的，但为了最终进行完整训练，需要执行彻底的清理并重启。

import os
from accelerate.utils import write_basic_config

write_basic_config()  # Write a config file
os._exit(00)  # Restart the notebook

准备数据集和模型

接下来，您应该准备数据集。如前所述，在准备 DataLoaders 和模型时，务必仔细操作，以确保没有任何内容被放到任何 GPU 上。

如果确实需要，建议将该特定代码封装到一个函数中，并在稍后介绍的 notebook 启动器界面中调用该函数。

请按照此处的说明确保数据集已下载

import os, re, torch, PIL
import numpy as np

from torch.optim.lr_scheduler import OneCycleLR
from torch.utils.data import DataLoader, Dataset
from torchvision.transforms import Compose, RandomResizedCrop, Resize, ToTensor

from accelerate import Accelerator
from accelerate.utils import set_seed
from timm import create_model

首先，你需要创建一个函数，根据文件名提取类名：

import os

data_dir = "../../images"
fnames = os.listdir(data_dir)
fname = fnames[0]
print(fname)

beagle_32.jpg

在此示例中，标签为 beagle。您可以使用正则表达式从文件名中提取标签：

import re


def extract_label(fname):
    stem = fname.split(os.path.sep)[-1]
    return re.search(r"^(.*)_\d+\.jpg$", stem).groups()[0]

extract_label(fname)

你可以看到它正确地返回了我们的文件名：

"beagle"

接下来，应创建一个 Dataset 类来处理图像和标签的获取：

class PetsDataset(Dataset):
    def __init__(self, file_names, image_transform=None, label_to_id=None):
        self.file_names = file_names
        self.image_transform = image_transform
        self.label_to_id = label_to_id

    def __len__(self):
        return len(self.file_names)

    def __getitem__(self, idx):
        fname = self.file_names[idx]
        raw_image = PIL.Image.open(fname)
        image = raw_image.convert("RGB")
        if self.image_transform is not None:
            image = self.image_transform(image)
        label = extract_label(fname)
        if self.label_to_id is not None:
            label = self.label_to_id[label]
        return {"image": image, "label": label}

现在来构建数据集。在训练函数外部，你可以查找并声明所有文件名和标签，并在启动的函数内部将它们用作引用：

fnames = [os.path.join("../../images", fname) for fname in fnames if fname.endswith(".jpg")]

接下来，收集所有标签：

all_labels = [extract_label(fname) for fname in fnames]
id_to_label = list(set(all_labels))
id_to_label.sort()
label_to_id = {lbl: i for i, lbl in enumerate(id_to_label)}

接下来，您应该创建一个 get_dataloaders 函数，用于返回您构建的数据加载器。如前所述，如果在构建 DataLoaders 时数据会自动发送到 GPU 或 TPU 设备，则必须使用此方法来构建它们。

def get_dataloaders(batch_size: int = 64):
    "Builds a set of dataloaders with a batch_size"
    random_perm = np.random.permutation(len(fnames))
    cut = int(0.8 * len(fnames))
    train_split = random_perm[:cut]
    eval_split = random_perm[cut:]

    # For training a simple RandomResizedCrop will be used
    train_tfm = Compose([RandomResizedCrop((224, 224), scale=(0.5, 1.0)), ToTensor()])
    train_dataset = PetsDataset([fnames[i] for i in train_split], image_transform=train_tfm, label_to_id=label_to_id)

    # For evaluation a deterministic Resize will be used
    eval_tfm = Compose([Resize((224, 224)), ToTensor()])
    eval_dataset = PetsDataset([fnames[i] for i in eval_split], image_transform=eval_tfm, label_to_id=label_to_id)

    # Instantiate dataloaders
    train_dataloader = DataLoader(train_dataset, shuffle=True, batch_size=batch_size, num_workers=4)
    eval_dataloader = DataLoader(eval_dataset, shuffle=False, batch_size=batch_size * 2, num_workers=4)
    return train_dataloader, eval_dataloader

最后，您应该导入稍后将使用的调度器：

from torch.optim.lr_scheduler import CosineAnnealingLR

编写训练函数

现在你可以构建训练循环了。notebook_launcher() 的工作原理是传入一个将在分布式系统上运行的函数。

以下是用于动物分类问题的基本训练循环：

代码已拆分为多个部分，以便对每个部分进行说明。文末将提供可完整复制粘贴的版本

def training_loop(mixed_precision="fp16", seed: int = 42, batch_size: int = 64):
    set_seed(seed)
    accelerator = Accelerator(mixed_precision=mixed_precision)

首先，你应在训练循环中尽早设置随机种子并创建 Accelerator 对象。

如果在 TPU 上进行训练，您的训练循环应将模型作为参数传入，并且模型应在训练循环函数之外实例化。请参阅 TPU 最佳实践了解原因

接下来，你应该构建数据加载器并创建模型：

    train_dataloader, eval_dataloader = get_dataloaders(batch_size)
    model = create_model("resnet50d", pretrained=True, num_classes=len(label_to_id))

在此处构建模型，以便种子也能控制新的权重初始化

在这个示例中，由于您正在执行迁移学习，模型的编码器初始状态为冻结，以便仅对模型的头部进行初始训练：

    for param in model.parameters():
        param.requires_grad = False
    for param in model.get_classifier().parameters():
        param.requires_grad = True

对图像批次进行归一化会使训练速度略微提升：

    mean = torch.tensor(model.default_cfg["mean"])[None, :, None, None]
    std = torch.tensor(model.default_cfg["std"])[None, :, None, None]

要在当前活动设备上使用这些常量，应将其设置为 Accelerator 的设备：

    mean = mean.to(accelerator.device)
    std = std.to(accelerator.device)

接下来实例化训练中使用的其余 PyTorch 类：

    optimizer = torch.optim.Adam(params=model.parameters(), lr=3e-2 / 25)
    lr_scheduler = OneCycleLR(optimizer=optimizer, max_lr=3e-2, epochs=5, steps_per_epoch=len(train_dataloader))

在将所有内容传递给 prepare() 之前。

没有特定的顺序需要记住，您只需按照传递给 prepare 方法的顺序解包这些对象即可。

    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
    )

现在训练模型：

    for epoch in range(5):
        model.train()
        for batch in train_dataloader:
            inputs = (batch["image"] - mean) / std
            outputs = model(inputs)
            loss = torch.nn.functional.cross_entropy(outputs, batch["label"])
            accelerator.backward(loss)
            optimizer.step()
            lr_scheduler.step()
            optimizer.zero_grad()

评估循环与训练循环略有不同。每个批次中传递的元素数量以及整体总准确率将被累加到两个常量中：

        model.eval()
        accurate = 0
        num_elems = 0

接下来是其余的标准 PyTorch 循环：

        for batch in eval_dataloader:
            inputs = (batch["image"] - mean) / std
            with torch.no_grad():
                outputs = model(inputs)
            predictions = outputs.argmax(dim=-1)

最后，还有一个主要的区别。

在执行分布式评估时，预测结果和标签需要通过 gather() 进行收集，以便在当前设备上获取所有数据，从而计算出准确的指标：

            accurate_preds = accelerator.gather(predictions) == accelerator.gather(batch["label"])
            num_elems += accurate_preds.shape[0]
            accurate += accurate_preds.long().sum()

现在你只需要计算该问题的实际指标，并可以在主进程中使用 print() 将其打印出来：

        eval_metric = accurate.item() / num_elems
        accelerator.print(f"epoch {epoch}: {100 * eval_metric:.2f}")

完整的训练循环版本如下：

def training_loop(mixed_precision="fp16", seed: int = 42, batch_size: int = 64):
    set_seed(seed)
    # Initialize accelerator
    accelerator = Accelerator(mixed_precision=mixed_precision)
    # Build dataloaders
    train_dataloader, eval_dataloader = get_dataloaders(batch_size)

    # Instantiate the model (you build the model here so that the seed also controls new weight initializations)
    model = create_model("resnet50d", pretrained=True, num_classes=len(label_to_id))

    # Freeze the base model
    for param in model.parameters():
        param.requires_grad = False
    for param in model.get_classifier().parameters():
        param.requires_grad = True

    # You can normalize the batches of images to be a bit faster
    mean = torch.tensor(model.default_cfg["mean"])[None, :, None, None]
    std = torch.tensor(model.default_cfg["std"])[None, :, None, None]

    # To make these constants available on the active device, set it to the accelerator device
    mean = mean.to(accelerator.device)
    std = std.to(accelerator.device)

    # Instantiate the optimizer
    optimizer = torch.optim.Adam(params=model.parameters(), lr=3e-2 / 25)

    # Instantiate the learning rate scheduler
    lr_scheduler = OneCycleLR(optimizer=optimizer, max_lr=3e-2, epochs=5, steps_per_epoch=len(train_dataloader))

    # Prepare everything
    # There is no specific order to remember, you just need to unpack the objects in the same order you gave them to the
    # prepare method.
    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
    )

    # Now you train the model
    for epoch in range(5):
        model.train()
        for batch in train_dataloader:
            inputs = (batch["image"] - mean) / std
            outputs = model(inputs)
            loss = torch.nn.functional.cross_entropy(outputs, batch["label"])
            accelerator.backward(loss)
            optimizer.step()
            lr_scheduler.step()
            optimizer.zero_grad()

        model.eval()
        accurate = 0
        num_elems = 0
        for batch in eval_dataloader:
            inputs = (batch["image"] - mean) / std
            with torch.no_grad():
                outputs = model(inputs)
            predictions = outputs.argmax(dim=-1)
            accurate_preds = accelerator.gather(predictions) == accelerator.gather(batch["label"])
            num_elems += accurate_preds.shape[0]
            accurate += accurate_preds.long().sum()

        eval_metric = accurate.item() / num_elems
        # Use accelerator.print to print only on the main process.
        accelerator.print(f"epoch {epoch}: {100 * eval_metric:.2f}")

使用 notebook_launcher

剩下的就是使用 notebook_launcher() 了。

您传入函数、参数（以元组形式）以及用于训练的进程数。（更多信息请参阅文档）

from accelerate import notebook_launcher

args = ("fp16", 42, 64)
notebook_launcher(training_loop, args, num_processes=2)

在多个节点上运行的情况下，你需要在每个节点上设置一个 Jupyter 会话，并同时运行启动单元格。

对于包含 2 个节点（计算机）的环境，每个节点配备 8 个 GPU，且主计算机的 IP 地址为“172.31.43.8”，其配置如下所示：

notebook_launcher(training_loop, args, master_addr="172.31.43.8", node_rank=0, num_nodes=2, num_processes=8)

在另一台机器上的第二个 Jupyter 会话中：

注意 node_rank 已经发生了变化

notebook_launcher(training_loop, args, master_addr="172.31.43.8", node_rank=1, num_nodes=2, num_processes=8)

在 TPU 上运行时，其形式如下：

model = create_model("resnet50d", pretrained=True, num_classes=len(label_to_id))

args = (model, "fp16", 42, 64)
notebook_launcher(training_loop, args, num_processes=8)

要启动具有弹性的训练流程以实现容错，您可以使用 PyTorch 提供的 elastic_launch 功能。这需要设置 rdzv_backend 和 max_restarts 等额外参数。以下是使用 notebook_launcher 启用弹性功能的示例：

notebook_launcher(
    training_loop,
    args,
    num_processes=2,
    max_restarts=3
)

在运行过程中，它会打印进度，并说明你使用了多少设备。本教程是在两个 GPU 上运行的：

Launching training on 2 GPUs.
epoch 0: 88.12
epoch 1: 91.73
epoch 2: 92.58
epoch 3: 93.90
epoch 4: 94.71

就是这样！

请注意，notebook_launcher() 会忽略 Accelerate 配置文件，若要基于配置文件启动，请使用：

accelerate launch

调试

在运行 notebook_launcher 时，一个常见问题是收到 CUDA/XPU has already been initialized 错误。这通常源于 notebook 中的某个导入或先前代码调用了 PyTorch 的 torch.cuda 或 torch.xpu 子库。为了帮助缩小问题范围，您可以在环境中设置 ACCELERATE_DEBUG_MODE=yes 来启动 notebook_launcher，这样在生成进程时会进行额外的检查，以确认普通进程能够正常创建并顺利使用 CUDA/XPU。（您的 CUDA/XPU 代码之后仍然可以运行）。

结论

本笔记本演示了如何在 Jupyter Notebook 内部执行分布式训练。请记住以下要点：

请确保将任何使用 CUDA/XPU（或 CUDA/XPU 导入）的代码保存到传递给 notebook_launcher() 的函数中
将 num_processes 设置为用于训练的设备数量（例如 GPU、XPU、CPU、TPU 等的数量）
如果使用 TPU，请在训练循环函数外部声明模型

Update on GitHub