HuggingFace Accelerate 0.12 : Tutorials : Jupyter 環境からのマルチノード訓練の起動 (翻訳/解説)

翻訳 : (株)クラスキャット セールスインフォメーション
作成日時 : 10/04/2022 (v0.12.0)

* 本ページは、HuggingFace Accelerate の以下のドキュメントを翻訳した上で適宜、補足説明したものです：

• Tutorials : Launching Multi-Node Training from a Jupyter Environment

* サンプルコードの動作確認はしておりますが、必要な場合には適宜、追加改変しています。
* ご自由にリンクを張って頂いてかまいませんが、sales-info@classcat.com までご一報いただけると嬉しいです。

 

 

 

HuggingFace Accelerate 0.12 : Tutorials : Jupyter 環境からのマルチノード訓練の起動

このチュートリアルは、分散システム上の Jupyter Notebook から 🤗 Accelerate でコンピュータビジョン・モデルを再調整する方法を教えます。また、環境が正しく構成されているか、データが正しく準備されているかを確認するために必要な幾つかの要件をセットアップし、最後に訓練を起動する方法を学習します。

This tutorial is also available as a Jupyter Notebook here

 

環境の設定

訓練が実行できる前に、🤗 Accelerate config ファイルはシステムに存在しなければなりません。通常はこれはターミナルで以下を実行してプロンプトに答えることで行なうことができます :

accelerate config

けれども、一般的なデフォルトで問題なく、TPU 上で実行していないのであれば、🤗 Accelerate は utils.write_basic_config() を通して GPU configuration を config ファイルに素早く書くためのユティリティを持ちます。

以下のコードは configuration を書いた後に Jupyter を再起動します、これを実行するために CUDA コードが呼び出されたからです。

CUDA はマルチノードシステム上では一度より多くは初期化できません。ノートブックでデバッグして CUDA への呼び出しを持つことは問題ありませんが、最終的に訓練するためには、完全なクリーンアップと再起動が実行される必要があります。

import os
from accelerate.utils import write_basic_config

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

 

データセットとモデルの準備

次にデータセットを準備する必要があります。前述のように、データローダとモデルを準備するとき、何も 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 です。regex を使用してファイル名からラベルを抽出できます :

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)

And you can see it properly returned the right name for our file:

"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")]

Next gather all the labels:

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 関数を作成する必要があります、これは構築されたデータロードを返します。前述のように、データローダを構築するときデータが自動的に 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 の device に設定する必要があります :

    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))

すべてを parepare() に渡す前に。

覚えておくべき特定の順序はありません、オブジェクトを parepare() メソッドに与えたのと同じ順序でオブジェクトをアンパックする必要があるだけです。

    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()

評価ループは訓練ループと比べて僅かに違って見えます。渡された要素の数と各バッチの全体的な合計精度が 2 つの定数に加算されます :

        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}")

A full version of this training loop is available below:

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 this constant available on the active device, set it to the accelerator device
    mean = mean.to(accelerator.device)
    std = std.to(accelerator.device)

    # Intantiate 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)

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)

それが実行されるとき、それは進捗と幾つのデバイス上で実行しているかの状態をプリントします。このチュートリアルは 2 つの 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

And that’s it!

 

終わりに

このノートブックは Jupyter Notebook の內部から分散訓練を実行する方法を示しました。幾つかの覚えておくべきキーノートです :

• notebook_launcher() に渡された関数のために CUDA を使用する (or CUDA インポート) どのようなコードも確実にセーブします。
  

• num_processes を訓練に使用される (GPU, CPU, TPU 等の数のような) デバイスの数に設定します。
  

• TPU を使用する場合、モデルを訓練ループ関数の外側で宣言します。

以上