從 Jupyter Notebook 啟動分散式訓練

本教程教你如何使用 🤗 Accelerate 從 Jupyter Notebook 在分散式系統上微調一個計算機視覺模型。你還將學習如何設定一些必要的前置條件，以確保你的環境配置正確、資料準備妥當，並最終啟動訓練。

本教程也提供 Jupyter Notebook 版本，請點選此處檢視。

配置環境

在進行任何訓練之前，系統中必須存在一個 Accelerate 配置檔案。通常，這可以透過在終端中執行以下命令並回答提示來完成：

accelerate config

然而，如果通用預設設定可以接受，並且你*沒有*在 TPU 上執行，Accelerate 提供了一個實用工具，可以透過 utils.write_basic_config() 快速將你的 GPU 配置寫入配置檔案。

以下程式碼將在寫入配置後重啟 Jupyter，因為執行此操作呼叫了 CUDA 程式碼。

在多 GPU 系統上，CUDA 不能被初始化超過一次。在 notebook 中除錯並呼叫 CUDA 是可以的，但為了最終進行訓練，需要執行完整的清理和重啟。

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 函式，它會為你返回構建好的 dataloader。如前所述，如果在構建 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 最佳實踐瞭解原因。

接下來，你應該構建你的 dataloader 並建立你的模型。

    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 已經初始化的錯誤。這通常源於 notebook 中的匯入或之前的程式碼呼叫了 PyTorch 的 `torch.cuda` 子庫。為了幫助縮小問題範圍，你可以在環境中設定 `ACCELERATE_DEBUG_MODE=yes` 來啟動 `notebook_launcher`，在生成程序時會進行額外檢查，確保常規程序可以無誤地建立並利用 CUDA。（你的 CUDA 程式碼之後仍然可以執行）。

結論

這個 notebook 展示瞭如何從 Jupyter Notebook 內部進行分散式訓練。以下是一些需要記住的關鍵點：

確保將任何使用 CUDA（或 CUDA 匯入）的程式碼保留在傳遞給 notebook_launcher() 的函式中。
將 num_processes 設定為用於訓練的裝置數量（例如 GPU、CPU、TPU 的數量等）。
如果使用 TPU，請在訓練迴圈函式之外宣告你的模型。

< > 在 GitHub 上更新