从 Jupyter Notebook 启动分布式训练
本教程将教您如何在分布式系统上的 Jupyter Notebook 中使用 🤗 Accelerate 微调计算机视觉模型。您还将学习如何设置一些确保您的环境配置正确、数据准备妥当，以及最后如何启动训练的必要条件。
本教程也以 Jupyter Notebook 形式提供 这里
配置环境
在执行任何训练之前，系统中必须存在 Accelerate 配置文件。通常可以通过在终端中运行以下命令并回答提示来完成此操作
accelerate config
但是，如果通用默认值就可以了，并且您没有在 TPU 上运行，Accelerate 拥有一个工具，可以通过 utils.write_basic_config() 快速将您的 GPU 配置写入配置文件。
以下代码将在写入配置后重新启动 Jupyter，因为已调用 CUDA 代码来执行此操作。
CUDA 在多 GPU 系统上不能初始化多次。在笔记本中进行调试并调用 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 上。
如果您确实这样做，建议将该特定代码放入函数中，并从笔记本启动器界面中调用该函数，该界面将在后面显示。
确保根据 这里 的说明下载数据集
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 initaliziations)
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_launcherargs = ("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 已初始化的错误。这通常源于笔记本中的导入或之前代码对 PyTorch torch.cuda子库的调用。为了帮助缩小错误范围，您可以在环境中使用ACCELERATE_DEBUG_MODE=yes启动notebook_launcher，在生成时会进行额外的检查，以确保可以创建一个常规进程并在没有问题的情况下使用 CUDA。（您的 CUDA 代码仍然可以在之后运行）。
结论
本笔记本展示了如何在 Jupyter 笔记本中执行分布式训练。一些关键注意事项：
- 确保将所有使用 CUDA（或 CUDA 导入）的代码保存到传递给 notebook_launcher() 的函数中。
- 将
num_processes设置为用于训练的设备数量（例如，GPU、CPU、TPU 等的数量）。 - 如果使用 TPU，请在训练循环函数之外声明您的模型。