Notebook 分布式训练

从 Notebook 启动训练脚本：使用 Accelerate

概述

在本教程中，我们将学习如何使用 Accelerate 从您的 notebook 内部，在分布式系统上启动一个训练函数！

为简单起见，本示例将沿用 PETs 训练示例，展示只需 3 行新增代码即可实现目标！

设置导入和构建 DataLoaders

首先，通过运行以下命令确保您的系统已安装 Accelerate：

pip install accelerate -U

在您的代码中，除了正常的 from fastai.module.all import * 导入之外，还需要添加两个新的导入：

+ from fastai.distributed import *
from fastai.vision.all import *
from fastai.vision.models.xresnet import *

+ from accelerate import notebook_launcher
+ from accelerate.utils import write_basic_config

第一个导入带来了 Learner.distrib_ctx 上下文管理器。第二个导入带来了 Accelerate 的 notebook_launcher，这是我们将调用以运行所需功能的关键函数。

我们需要设置 Accelerate 以使用我们所有的 GPU。我们可以通过 write_basic_config () 快速完成此操作。

注意

由于此项会检查 torch.cuda.device_count，因此您需要重启 notebook 并跳过再次调用此项以继续。它只需运行一次！另外，如果您选择不使用此项，请在终端中运行 accelerate config 并将 mixed_precision 设置为 no。

#from accelerate.utils import write_basic_config
#write_basic_config()

接下来，让我们下载一些数据进行训练。您无需担心使用 rank0_first，因为我们在 Jupyter Notebook 中，它会像往常一样只在一个进程上运行。

path = untar_data(URLs.PETS)

我们将 DataLoaders 的创建、我们的 vision_learner 以及对 fine_tune 的调用包装在一个 train 函数内部。

注意

重要的是不要在函数外部构建 DataLoaders，因为在此之前绝对不能将任何东西加载到 CUDA 上。

def get_y(o): return o[0].isupper()
def train(path):
    dls = ImageDataLoaders.from_name_func(
        path, get_image_files(path), valid_pct=0.2,
        label_func=get_y, item_tfms=Resize(224))
    learn = vision_learner(dls, resnet34, metrics=error_rate).to_fp16()
    learn.fine_tune(1)

train 函数所需的最后一项添加是在调用 fine_tune 之前使用我们的上下文管理器，并将 in_notebook 设置为 True。

注意

在本示例中，为了与 torchvision 的 resnet34 兼容，sync_bn 已被禁用。

def train(path):
    dls = ImageDataLoaders.from_name_func(
        path, get_image_files(path), valid_pct=0.2,
        label_func=get_y, item_tfms=Resize(224))
    learn = vision_learner(dls, resnet34, metrics=error_rate).to_fp16()
    with learn.distrib_ctx(sync_bn=False, in_notebook=True):
        learn.fine_tune(1)
    learn.export("pets")

最后，只需调用 notebook_launcher，传入训练函数、任何参数（作为元组）以及要使用的 GPU（进程）数量即可。

notebook_launcher(train, (path,), num_processes=2)
Launching training on 2 GPUs.
Training Learner...
周期 训练损失 验证损失 错误率 时间
0 0.342019 0.228441 0.105041 00:54
周期 训练损失 验证损失 错误率 时间
0 0.197188 0.141764 0.062246 00:56

之后，我们可以导入导出的 Learner，进行保存，或者在 Jupyter Notebook 中进行分布式进程之外的任何其他操作。

imgs = get_image_files(path)
learn = load_learner(path/'pets')
learn.predict(imgs[0])
('False', TensorBase(0), TensorBase([0.9718, 0.0282]))