Notebook分布式训练

! [ -e /content ] && pip install -Uqq fastai  # 在Colab上升级fastai

使用Accelerate从您的笔记本启动训练脚本

概述

在本教程中，我们将学习如何使用 Accelerate 在分布式系统上启动一个训练函数，从你的 笔记本 中进行操作！

为了简化示例，本例将跟随训练 PETs，展示只需 3 行新代码即可开始你的旅程！

设置导入和构建数据加载器

首先，通过运行以下命令确保在系统上安装了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，这是我们将调用的关键函数，以运行我们想要的内容。

隐藏
from fastai.vision.all import *
from fastai.distributed import *
from fastai.vision.models.xresnet import *

from accelerate import notebook_launcher
from accelerate.utils import write_basic_config

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

Note

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

#从accelerate.utils导入write_basic_config
#写入基本配置()

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

path = untar_data(URLs.PETS)

我们将 DataLoaders 的创建、我们的 vision_learner 和对 fine_tune 的调用封装在一个 train 函数中。

Note

重要的是不要在函数外部构建 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：

Note

为了兼容 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...
epoch train_loss valid_loss error_rate time
0 0.342019 0.228441 0.105041 00:54
epoch train_loss valid_loss error_rate time
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]))