PyTorch Ignite 0.4.8 : AI Tutorials : Intermediate – CIFAR10 の分散訓練 (翻訳/解説)

翻訳 : (株)クラスキャット セールスインフォメーション
作成日時 : 03/06/2022 (0.4.8)

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

• Tutorials : Intermediate – Distributed Training with Ignite on CIFAR10

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

 

 

PyTorch Ignite 0.4.8 : AI Tutorials : Intermediate – CIFAR10 の分散訓練

このチュートリアルは、 Ignite で 1 つ以上の CPU, GPU or TPU 上で分散訓練を行える方法の簡潔なイントロダクションです。
幾つかのヘルパー関数と Ignite の概念 (共通の訓練ハンドラのセットアップ、チェックポイントからのセーブ・ロード 等) も紹介します、これらを貴方のコードに簡単に組み込むことができます。

以下の configuration のいずれかを利用して、事前定義された CIFAR10 上の ResNet18 を訓練するために分散訓練を使用します :

• 単一ノード、1つ以上の GPU
• マルチノード、マルチ GPU
• 単一ノード、マルチ CPU
• Google Colab 上の TPU
• Jupyter Notebook 上

使用する分散訓練のタイプはデータ並列と呼称され、そこでは :

1. モデルを各 GPU にコピーする。
2. データセットを分割して異なるサブセット上でモデルを適合させる。
3. モデルの同期を維持するために各イテレーションで勾配を伝達する。

• Distributed Deep Learning 101: イントロダクション

PyTorch はこのタスクのために torch.nn.parallel.DistributedDataParallel API を提供していますが、異なるバックエンド + 構成をサポートする実装は退屈です。このサンプルでは、数行のコードだけで様々なバックエンドに対応可能なデータ分散訓練を可能にする方法を以下と共に見ます :

• 訓練と検証メトリクスの計算
• ロギングのセットアップ (そして ClearML との接続)
• 最良モデル重みのセーブ
• LR スケジューラの設定
• 自動混合精度の使用

 

Required Dependencies

!pip install pytorch-ignite

 

For parsing arguments

!pip install fire

 

For TPUs

VERSION = !curl -s https://api.github.com/repos/pytorch/xla/releases/latest | grep -Po '"tag_name": "v\K.*?(?=")'
VERSION = VERSION[0].rstrip('.0') # remove trailing zero
!pip install cloud-tpu-client==0.10 https://storage.googleapis.com/tpu-pytorch/wheels/torch_xla-{VERSION}-cp37-cp37m-linux_x86_64.whl

 

With ClearML (オプション)

以下のように実験を追跡するために ClearML によるロギングを有効にできます :

• ClearML アカウントを持っていることを確認します : https://app.community.clear.ml/
• クレデンシャルの作成 : Profile > Create new credentials > Copy to clipboard
• clearml-init を実行してクレデンシャルを貼り付けます。

!pip install clearml
!clearml-init

下の config で with_clearml=True を指定してダッシュボードで実験をモニタします。そのような実験のサンプルを見るためにはこのチュートリアルの最後を参照してください。

 

データのダウンロード

最初にデータをダウンロードしましょう、これは後でデータローダをインスタンス化するために総てのプロセスにより使用されます。次のコマンドは CIFAR10 データセットをフォルダ cifar10 にダウンロードします。

!python -c "from torchvision.datasets import CIFAR10; CIFAR10('cifar10', download=True)"

 

共通の Configuration

config 辞書を保持します、これは訓練の間に必要なパラメータをストアするために拡張したり変更したりできます。後でこれらのパラメータを使用するときにこのコードを参照できます。

config = {
    "seed": 543,
    "data_path": "cifar10",
    "output_path": "output-cifar10/",
    "model": "resnet18",
    "batch_size": 512,
    "momentum": 0.9,
    "weight_decay": 1e-4,
    "num_workers": 2,
    "num_epochs": 5,
    "learning_rate": 0.4,
    "num_warmup_epochs": 1,
    "validate_every": 3,
    "checkpoint_every": 200,
    "backend": None,
    "resume_from": None,
    "log_every_iters": 15,
    "nproc_per_node": None,
    "with_clearml": False,
    "with_amp": False,
}

 

基本的なセットアップ

インポート

from datetime import datetime
from pathlib import Path

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, models
from torchvision.transforms import (
    Compose,
    Normalize,
    Pad,
    RandomCrop,
    RandomHorizontalFlip,
    ToTensor,
)

import ignite
import ignite.distributed as idist
from ignite.contrib.engines import common
from ignite.handlers import PiecewiseLinear
from ignite.engine import Events, create_supervised_trainer, create_supervised_evaluator
from ignite.handlers import Checkpoint, global_step_from_engine
from ignite.metrics import Accuracy, Loss
from ignite.utils import manual_seed, setup_logger

次にデータローダ, モデルと optimizer を現在の configuration backend=None (非分散) あるいは nccl, gloo, and xla-tpu (分散) のようなバックエンドに自動的に適応させるために idist ( ignite.distributed ) の auto_ メソッドの助けを借ります。

auto_ メソッドを部分的に使用するか全く使用しないかは自由で、代わりに何かカスタムなものを実装できることに注意してください。

 

データローダ

次に train と test データセットを data_path からインスタンス化し、transforms をそれに適用して get_train_test_datasets() を通してそれらを返します。

def get_train_test_datasets(path):
    train_transform = Compose(
        [
            Pad(4),
            RandomCrop(32, fill=128),
            RandomHorizontalFlip(),
            ToTensor(),
            Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
        ]
    )
    test_transform = Compose(
        [
            ToTensor(),
            Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
        ]
    )

    train_ds = datasets.CIFAR10(
        root=path, train=True, download=False, transform=train_transform
    )
    test_ds = datasets.CIFAR10(
        root=path, train=False, download=False, transform=test_transform
    )

    return train_ds, test_ds

最後に、データセットを auto_dataloader() に渡します。

def get_dataflow(config):
    train_dataset, test_dataset = get_train_test_datasets(config["data_path"])

    train_loader = idist.auto_dataloader(
        train_dataset,
        batch_size=config["batch_size"],
        num_workers=config["num_workers"],
        shuffle=True,
        drop_last=True,
    )

    test_loader = idist.auto_dataloader(
        test_dataset,
        batch_size=2 * config["batch_size"],
        num_workers=config["num_workers"],
        shuffle=False,
    )
    return train_loader, test_loader

 

モデル

config で与えられたモデルが torchvision.models に存在するか確認し、最後の層を (CIFAR10 に在るように) 10 クラスを出力するように変更してそれを auto_model() に渡します、これは自動的にそれを非分散と分散構成に適応可能にします。

def get_model(config):
    model_name = config["model"]
    if model_name in models.__dict__:
        fn = models.__dict__[model_name]
    else:
        raise RuntimeError(f"Unknown model name {model_name}")

    model = idist.auto_model(fn(num_classes=10))

    return model

 

Optimizer

それから config からのハイパーパラメータを使用して optimizer をセットアップしてそれを auto_optim() に渡すことができます。

def get_optimizer(config, model):
    optimizer = optim.SGD(
        model.parameters(),
        lr=config["learning_rate"],
        momentum=config["momentum"],
        weight_decay=config["weight_decay"],
        nesterov=True,
    )
    optimizer = idist.auto_optim(optimizer)

    return optimizer

 

Criterion

損失関数をデバイスに配置します。

def get_criterion():
    return nn.CrossEntropyLoss().to(idist.device())

 

LR スケジューラ

PiecewiseLinear を使用します、これは Ignite が提供する 様々な LR スケジューラ の一つです。

def get_lr_scheduler(config, optimizer):
    milestones_values = [
        (0, 0.0),
        (config["num_iters_per_epoch"] * config["num_warmup_epochs"], config["learning_rate"]),
        (config["num_iters_per_epoch"] * config["num_epochs"], 0.0),
    ]
    lr_scheduler = PiecewiseLinear(
        optimizer, param_name="lr", milestones_values=milestones_values
    )
    return lr_scheduler

 

トレーナー

モデルのセーブ

ハンドラ (ClearML の場合) を使用するかチェックポイントファイルのパスを単純に save_handler を渡すことによりチェックポイントを作成できます。with-clearml=True を指定した場合、ClearMLSaver() を使用して ClearML の File Server でモデルをセーブします。

def get_save_handler(config):
    if config["with_clearml"]:
        from ignite.contrib.handlers.clearml_logger import ClearMLSaver

        return ClearMLSaver(dirname=config["output_path"])

    return config["output_path"]

 

チェックポイントから再開する

チェックポイントファイルパスが提供される場合、ファイルをロードすることでそこから訓練を再開できます。

def load_checkpoint(resume_from):
    checkpoint_fp = Path(resume_from)
    assert (
        checkpoint_fp.exists()
    ), f"Checkpoint '{checkpoint_fp.as_posix()}' is not found"
    checkpoint = torch.load(checkpoint_fp.as_posix(), map_location="cpu")
    return checkpoint

 

トレーナーの作成

最後に、4 つのステップでトレーナーを作成できます :

1. create_supervised_trainer() を使用して trainer オブジェクトを作成します、これは単一バッチを処理するために取られるステップを内部的に定義しています :
   • バッチを現在の分散設定で使用されるデバイスに転送します。
   • モデルを train() モードにします。
   • 入力をモデルに渡して損失を計算することにより forward パスを実行します。AMP が有効であれば、このステップは autocast が有効な状態で行なわれ、これはこのステップが混合精度で実行されることを可能にします。
   • backward パスを実行します。自動混合精度 (AMP) が有効である場合 (大きなニューラルネットワークの計算を高速化してパフォーマンスを維持しながらメモリ使用量を削減します)、backward() を呼び出す前に損失は スケール され、NaN を含むバッチを破棄しながら optimizer を step() し、そして次のイテレーションのためにスケールを update() します。
   • 損失をバッチ損失として state.output にストアします。

    
   内部的には、trainer を作成するための上のステップは以下のようなものです :

   def train_step(engine, batch):
   
         x, y = batch[0], batch[1]
         if x.device != device:
             x = x.to(device, non_blocking=True)
             y = y.to(device, non_blocking=True)
   
         model.train()
   
         with autocast(enabled=with_amp):
             y_pred = model(x)
             loss = criterion(y_pred, y)
   
         optimizer.zero_grad()
         scaler.scale(loss).backward()  # If with_amp=False, this is equivalent to loss.backward()
         scaler.step(optimizer)  # If with_amp=False, this is equivalent to optimizer.step()
         scaler.update()  # If with_amp=False, this step does nothing
   
         return {"batch loss": loss.item()}
   
   trainer = Engine(train_step)
   

2. 幾つかの一般的な Ignite 訓練ハンドラをセットアップします。これを個別に行なうか、以下と一緒に trainer とデータセットのサブセット (train_sampler) を取る setup_common_training_handlers を使用することができます。
   • チェックポイントで何をセーブするか (to_save) そしてどのくらいの頻度でセーブするか (save_every_iters) の辞書マッピング。
   • LR スケジューラ
   • train_step() の出力
   • その他のハンドラ

3. resume_from ファイルパスが提供される場合、チェックポイントからオブジェクトの状態 to_save をロードします。

   def create_trainer(
       model, optimizer, criterion, lr_scheduler, train_sampler, config, logger
   ):
   
       device = idist.device()
       amp_mode = None
       scaler = False
           
       trainer = create_supervised_trainer(
           model,
           optimizer,
           criterion,
           device=device,
           non_blocking=True,
           output_transform=lambda x, y, y_pred, loss: {"batch loss": loss.item()},
           amp_mode="amp" if config["with_amp"] else None,
           scaler=config["with_amp"],
       )
       trainer.logger = logger
   
       to_save = {
           "trainer": trainer,
           "model": model,
           "optimizer": optimizer,
           "lr_scheduler": lr_scheduler,
       }
       metric_names = [
           "batch loss",
       ]
   
       common.setup_common_training_handlers(
           trainer=trainer,
           train_sampler=train_sampler,
           to_save=to_save,
           save_every_iters=config["checkpoint_every"],
           save_handler=get_save_handler(config),
           lr_scheduler=lr_scheduler,
           output_names=metric_names if config["log_every_iters"] > 0 else None,
           with_pbars=False,
           clear_cuda_cache=False,
       )
   
       if config["resume_from"] is not None:
           checkpoint = load_checkpoint(config["resume_from"])
           Checkpoint.load_objects(to_load=to_save, checkpoint=checkpoint)
   
       return trainer
   

 

Evaluator

evaluator は create_supervised_evaluator を通して作成されます、これは内部的には :

1. モデルを eval() モードに設定する。
2. バッチを現在の分散設定で使用されるデバイスに転送します。
3. forward パスを実行する。AMP が有効であれば、autocast が有効になります。
4. メトリクスを計算するために予測値とラベルを state.output にストアする。

それはまた evaluator に渡された Ignite メトリクスも装着します。

def create_evaluator(model, metrics, config):
    device = idist.device()

    amp_mode = "amp" if config["with_amp"] else None
    evaluator = create_supervised_evaluator(
        model, metrics=metrics, device=device, non_blocking=True, amp_mode=amp_mode
    )
    
    return evaluator

 

訓練

訓練を開始する前に、マスタープロセス (rank = 0) 上で幾つかのことをセットアップしなければなりません :

• チェックポイント, ベストモデルと model_backend_rank_time 形式の tensorboard ロギングの出力をストアするためのフォルダを作成します。
• モデルをセーブするために ClearML FileServer が使用される場合には、Task が作成されなければなりません、そして config 辞書と実験の一部である特定のハイパーパラメータを渡します。

def setup_rank_zero(logger, config):
    device = idist.device()

    now = datetime.now().strftime("%Y%m%d-%H%M%S")
    output_path = config["output_path"]
    folder_name = (
        f"{config['model']}_backend-{idist.backend()}-{idist.get_world_size()}_{now}"
    )
    output_path = Path(output_path) / folder_name
    if not output_path.exists():
        output_path.mkdir(parents=True)
    config["output_path"] = output_path.as_posix()
    logger.info(f"Output path: {config['output_path']}")

    if config["with_clearml"]:
        from clearml import Task

        task = Task.init("CIFAR10-Training", task_name=output_path.stem)
        task.connect_configuration(config)
        # Log hyper parameters
        hyper_params = [
            "model",
            "batch_size",
            "momentum",
            "weight_decay",
            "num_epochs",
            "learning_rate",
            "num_warmup_epochs",
        ]
        task.connect({k: v for k, v in config.items()})

 

ロギング

このステップはオプションですが、log_basic_info() に setup_logger() オブジェクトを渡して、異なるバージョン、現在の設定、(そのローカルランクにより識別される) 現在のプロセスにより使用されるデバイスとバックエンド、そしてプロセス数 (world サイズ) のような総ての基本的な情報をログ記録することができます。idist (ignite.distributed) はこれを可能にするために get_local_rank(), backend(), get_world_size() 等のような幾つかのユティリティ関数を提供しています。

def log_basic_info(logger, config):
    logger.info(f"Train on CIFAR10")
    logger.info(f"- PyTorch version: {torch.__version__}")
    logger.info(f"- Ignite version: {ignite.__version__}")
    if torch.cuda.is_available():
        # explicitly import cudnn as torch.backends.cudnn can not be pickled with hvd spawning procs
        from torch.backends import cudnn

        logger.info(
            f"- GPU Device: {torch.cuda.get_device_name(idist.get_local_rank())}"
        )
        logger.info(f"- CUDA version: {torch.version.cuda}")
        logger.info(f"- CUDNN version: {cudnn.version()}")

    logger.info("\n")
    logger.info("Configuration:")
    for key, value in config.items():
        logger.info(f"\t{key}: {value}")
    logger.info("\n")

    if idist.get_world_size() > 1:
        logger.info("\nDistributed setting:")
        logger.info(f"\tbackend: {idist.backend()}")
        logger.info(f"\tworld size: {idist.get_world_size()}")
        logger.info("\n")

これは validate_every エポック後に train と val メトリクスをログ記録する標準的なユティリティ関数です。

def log_metrics(logger, epoch, elapsed, tag, metrics):
    metrics_output = "\n".join([f"\t{k}: {v}" for k, v in metrics.items()])
    logger.info(
        f"\nEpoch {epoch} - Evaluation time (seconds): {elapsed:.2f} - {tag} metrics:\n {metrics_output}"
    )

 

訓練開始

これは主要ロジックが在るところで i.e. ここの中から上の関数総てを呼び出します :

1. 基本的なセットアップ
   

2. manual_seed() と setup_logger() を設定してから、総ての基本情報をログ記録します。
   

3. dataloader, モデル, optimizer, criterion と lr_scheduler を初期化します。
   

4. トレーナーを作成するために上のオブジェクトを使用します。
   

5. Evaluator
   

6. Accuracy() と Loss() のような幾つかの関連 Ignite メトリクスを定義する。
   

7. train_dataloader と val_dataloader 上でメトリクスを計算するために 2 つの evaluator : train_evaluator と val_evaluator をそれぞれ作成しますが、val_evaluator は検証メトリクスに基づいて最良モデルをストアします。
   

8. 両者の dataloader 上でメトリクスを計算してそれらをログ記録するために run_validation を定義します。それから validate_every エポック後と訓練が完了した後に実行するために trainer にこの関数を装着します。
   

9. trainer と evaluator のためにマスタープロセス上で setup_tb_logging() を使用して TensorBoard ロギングをセットアップします、その結果、学習率と一緒に訓練と検証メトリクスがログ記録できます。
   

10. (メトリクスで Accracy() として定義される) 検証精度による 2 つ (n_saved) の最良モデルをストアするために Checkpoint() オブジェクトを定義して、val_evaluator が実行されるたびに実行されるようにそれを val_evaluator に装着します。
    

11. num_epochs の間 train_loader で訓練を試行します。
    

12. 訓練が完了したら TensorBoard logger をクローズします。

def training(local_rank, config):

    rank = idist.get_rank()
    manual_seed(config["seed"] + rank)

    logger = setup_logger(name="CIFAR10-Training")
    log_basic_info(logger, config)

    if rank == 0:
        setup_rank_zero(logger, config)

    train_loader, val_loader = get_dataflow(config)
    model = get_model(config)
    optimizer = get_optimizer(config, model)
    criterion = get_criterion()
    config["num_iters_per_epoch"] = len(train_loader)
    lr_scheduler = get_lr_scheduler(config, optimizer)

    trainer = create_trainer(
        model, optimizer, criterion, lr_scheduler, train_loader.sampler, config, logger
    )

    metrics = {
        "Accuracy": Accuracy(),
        "Loss": Loss(criterion),
    }

    train_evaluator = create_evaluator(model, metrics, config)
    val_evaluator = create_evaluator(model, metrics, config)

    def run_validation(engine):
        epoch = trainer.state.epoch
        state = train_evaluator.run(train_loader)
        log_metrics(logger, epoch, state.times["COMPLETED"], "train", state.metrics)
        state = val_evaluator.run(val_loader)
        log_metrics(logger, epoch, state.times["COMPLETED"], "val", state.metrics)

    trainer.add_event_handler(
        Events.EPOCH_COMPLETED(every=config["validate_every"]) | Events.COMPLETED,
        run_validation,
    )

    if rank == 0:
        evaluators = {"train": train_evaluator, "val": val_evaluator}
        tb_logger = common.setup_tb_logging(
            config["output_path"], trainer, optimizer, evaluators=evaluators
        )

    best_model_handler = Checkpoint(
        {"model": model},
        get_save_handler(config),
        filename_prefix="best",
        n_saved=2,
        global_step_transform=global_step_from_engine(trainer),
        score_name="val_accuracy",
        score_function=Checkpoint.get_default_score_fn("Accuracy"),
    )
    val_evaluator.add_event_handler(
        Events.COMPLETED,
        best_model_handler,
    )

    try:
        trainer.run(train_loader, max_epochs=config["num_epochs"])
    except Exception as e:
        logger.exception("")
        raise e

    if rank == 0:
        tb_logger.close()

 

分散コードの実行

上のコードをコンテキストマネージャ Parallel で簡単に実行できます :

with idist.Parallel(backend=backend, nproc_per_node=nproc_per_node) as parallel:
    parallel.run(training, config)

Parallel は、総てのサポートされる分散バックエンドと非分散設定に渡り同じコードをシームレスに実行することを可能にします。ここではバックエンドは分散通信フレームワークを参照します。どのバックエンドを選択するかについては ここ で更に読んでください。Parallel はバックエンドと以下のいずれかを受け取ります :

• nproc_per_node プロセスを spawn して提供されるバックエンドに従って処理グループを初期化します (スタンドアロン・スクリプトに有用)。

この方法は torch.multiprocessing.spawn を使用してプロセスを spawn するデフォルトの方法です。けれども、この方法は初期化オーバーヘッドにより遅いです。

or

• バックエンドに与えられたとき処理グループのみを初期化する ( torch.distributed.launch, horovodrun, 等のようなツールで有用)。

この方法が推奨されます、訓練が高速で複数のスクリプトに拡張することが容易だからです。

下で見るように追加情報を Parallel に spawn_kwargs として集合的に渡すことができます。

Note: 使い易さのために分散コードをスクリプトとして実行することが勧められますが、Jupyter ノートブックでプロセスを spawn することもできます (チュートリアルの最後参照)。スクリプトとしての完全なコードは ここ で見つけられます。

スクリプトを実行するために下で提示された方法の一つを選択してください。

 

単一ノード、1 つまたはそれ以上の GPU

fire を使用して run() を CLI に変換し、run() 内でパースされた引数を直接使用してスクリプトで訓練を開始します :

import fire

def run(backend=None, **spawn_kwargs):
    config["backend"] = backend
    
    with idist.Parallel(backend=config["backend"], **spawn_kwargs) as parallel:
        parallel.run(training, config)

if __name__ == "__main__":
    fire.Fire({"run": run})

そしてスクリプトを以下のように実行できます (e.g. 2 GPU のために) :

 

torch.distributed.launch で実行 (推奨)

python -u -m torch.distributed.launch --nproc_per_node=2 --use_env main.py run --backend="nccl"

 

内部の spawining で実行 ( torch.multiprocessing.spawn )

python -u main.py run --backend="nccl" --nproc_per_node=2

 

horovodrun で実行する

backend=horovod であることを確認してください。下で np はプロセス数です。

horovodrun -np=2 python -u main.py run --backend="horovod"

 

マルチノード, マルチ GPU

スクリプト内のコードは 単一ノード、1 つまたはそれ以上の GPU と同様です :

import fire

def run(backend=None, **spawn_kwargs):
    config["backend"] = backend
    
    with idist.Parallel(backend=config["backend"], **spawn_kwargs) as parallel:
        parallel.run(training, config)

if __name__ == "__main__":
    fire.Fire({"run": run})

唯一の変更はスクリプトの実行方法です。マスターノードの IP アドレスとそのポートをノード rank と一緒に提供する必要があります。
例えば、2 ノード (nnodes) と 2 GPU (nproc_per_node) のために、以下を行なうことができます :

 

torch.distributed.launch で実行 (推奨)

On ノード 0 (マスターノード):

python -u -m torch.distributed.launch \
    --nnodes=2 \
    --nproc_per_node=2 \
    --node_rank=0 \
    --master_addr=master --master_port=2222 --use_env \
    main.py run --backend="nccl"

On ノード 1 (ワーカーノード):

python -u -m torch.distributed.launch \
    --nnodes=2 \
    --nproc_per_node=2 \
    --node_rank=1 \
    --master_addr=master --master_port=2222 --use_env \
    main.py run --backend="nccl"

 

内部 spawning で実行

On ノード 0:

python -u main.py run
    --nnodes=2 \
    --nproc_per_node=2 \
    --node_rank=0 \
    --master_addr=master --master_port=2222 \
    --backend="nccl"

On node 1:

python -u main.py run
    --nnodes=2 \
    --nproc_per_node=2 \
    --node_rank=1 \
    --master_addr=master --master_port=2222 \
    --backend="nccl"

 

horovodruz で実行

下の np は nnodes x nproc_per_node で計算されます。

horovodrun -np 4 -H hostname1:2,hostname2:2 python -u main.py run --backend="horovod"

 

単一ノード、マルチ CPU

これは単一ノード、1 つまたはそれ以上の GPU と同様です。唯一の違いはスクリプトを実行する際、nccl の代わりに backend=gloo とします。

 

TPU on Google Colab

Go to Runtime > Change runtime type and select Hardware accelerator = TPU.

nproc_per_node = 8
config["backend"] = "xla-tpu"

with idist.Parallel(backend=config["backend"], nproc_per_node=nproc_per_node) as parallel:
    parallel.run(training, config)

2021-09-14 17:01:35,425 ignite.distributed.launcher.Parallel INFO: Initialized distributed launcher with backend: 'xla-tpu'
2021-09-14 17:01:35,427 ignite.distributed.launcher.Parallel INFO: - Parameters to spawn processes: 
	nproc_per_node: 8
	nnodes: 1
	node_rank: 0
2021-09-14 17:01:35,428 ignite.distributed.launcher.Parallel INFO: Spawn function '' in 8 processes
2021-09-14 17:01:47,607 CIFAR10-Training INFO: Train on CIFAR10
2021-09-14 17:01:47,639 CIFAR10-Training INFO: - PyTorch version: 1.8.2+cpu
2021-09-14 17:01:47,658 CIFAR10-Training INFO: - Ignite version: 0.4.6
2021-09-14 17:01:47,678 CIFAR10-Training INFO: 

2021-09-14 17:01:47,697 CIFAR10-Training INFO: Configuration:
2021-09-14 17:01:47,721 CIFAR10-Training INFO: 	seed: 543
2021-09-14 17:01:47,739 CIFAR10-Training INFO: 	data_path: cifar10
2021-09-14 17:01:47,765 CIFAR10-Training INFO: 	output_path: output-cifar10/
2021-09-14 17:01:47,786 CIFAR10-Training INFO: 	model: resnet18
2021-09-14 17:01:47,810 CIFAR10-Training INFO: 	batch_size: 512
2021-09-14 17:01:47,833 CIFAR10-Training INFO: 	momentum: 0.9
2021-09-14 17:01:47,854 CIFAR10-Training INFO: 	weight_decay: 0.0001
2021-09-14 17:01:47,867 CIFAR10-Training INFO: 	num_workers: 2
2021-09-14 17:01:47,887 CIFAR10-Training INFO: 	num_epochs: 5
2021-09-14 17:01:47,902 CIFAR10-Training INFO: 	learning_rate: 0.4
2021-09-14 17:01:47,922 CIFAR10-Training INFO: 	num_warmup_epochs: 1
2021-09-14 17:01:47,940 CIFAR10-Training INFO: 	validate_every: 3
2021-09-14 17:01:47,949 CIFAR10-Training INFO: 	checkpoint_every: 200
2021-09-14 17:01:47,960 CIFAR10-Training INFO: 	backend: xla-tpu
2021-09-14 17:01:47,967 CIFAR10-Training INFO: 	resume_from: None
2021-09-14 17:01:47,975 CIFAR10-Training INFO: 	log_every_iters: 15
2021-09-14 17:01:47,984 CIFAR10-Training INFO: 	nproc_per_node: None
2021-09-14 17:01:48,003 CIFAR10-Training INFO: 	with_clearml: False
2021-09-14 17:01:48,019 CIFAR10-Training INFO: 	with_amp: False
2021-09-14 17:01:48,040 CIFAR10-Training INFO: 

2021-09-14 17:01:48,059 CIFAR10-Training INFO: 
Distributed setting:
2021-09-14 17:01:48,079 CIFAR10-Training INFO: 	backend: xla-tpu
2021-09-14 17:01:48,098 CIFAR10-Training INFO: 	world size: 8
2021-09-14 17:01:48,109 CIFAR10-Training INFO: 

2021-09-14 17:01:48,130 CIFAR10-Training INFO: Output path: output-cifar10/resnet18_backend-xla-tpu-8_20210914-170148
2021-09-14 17:01:50,917 ignite.distributed.auto.auto_dataloader INFO: Use data loader kwargs for dataset 'Dataset CIFAR10': 
	{'batch_size': 64, 'num_workers': 2, 'drop_last': True, 'sampler': , 'pin_memory': False}
2021-09-14 17:01:50,950 ignite.distributed.auto.auto_dataloader INFO: DataLoader is wrapped by `MpDeviceLoader` on XLA
2021-09-14 17:01:50,975 ignite.distributed.auto.auto_dataloader INFO: Use data loader kwargs for dataset 'Dataset CIFAR10': 
	{'batch_size': 128, 'num_workers': 2, 'sampler': , 'pin_memory': False}
2021-09-14 17:01:51,000 ignite.distributed.auto.auto_dataloader INFO: DataLoader is wrapped by `MpDeviceLoader` on XLA
2021-09-14 17:01:53,866 CIFAR10-Training INFO: Engine run starting with max_epochs=5.
2021-09-14 17:02:23,913 CIFAR10-Training INFO: Epoch[1] Complete. Time taken: 00:00:30
2021-09-14 17:02:41,945 CIFAR10-Training INFO: Epoch[2] Complete. Time taken: 00:00:18
2021-09-14 17:03:13,870 CIFAR10-Training INFO: 
Epoch 3 - Evaluation time (seconds): 14.00 - train metrics:
    Accuracy: 0.32997744845360827
	Loss: 1.7080145767054606
2021-09-14 17:03:19,283 CIFAR10-Training INFO: 
Epoch 3 - Evaluation time (seconds): 5.39 - val metrics:
    Accuracy: 0.3424
	Loss: 1.691359375
2021-09-14 17:03:19,289 CIFAR10-Training INFO: Epoch[3] Complete. Time taken: 00:00:37
2021-09-14 17:03:37,535 CIFAR10-Training INFO: Epoch[4] Complete. Time taken: 00:00:18
2021-09-14 17:03:55,927 CIFAR10-Training INFO: Epoch[5] Complete. Time taken: 00:00:18
2021-09-14 17:04:07,598 CIFAR10-Training INFO: 
Epoch 5 - Evaluation time (seconds): 11.66 - train metrics:
    Accuracy: 0.42823775773195877
	Loss: 1.4969784451514174
2021-09-14 17:04:10,190 CIFAR10-Training INFO: 
Epoch 5 - Evaluation time (seconds): 2.56 - val metrics:
    Accuracy: 0.4412
	Loss: 1.47838994140625
2021-09-14 17:04:10,244 CIFAR10-Training INFO: Engine run complete. Time taken: 00:02:16
2021-09-14 17:04:10,313 ignite.distributed.launcher.Parallel INFO: End of run

 

Jupyter Notebook で実行

ノートブックでプロセスを spawn する必要があります、従ってそれを成すために内部 spawning を使用します。マルチ GPU については backend=nccl を使用して、マルチ CPU のためには backend=gloo を使用します。

spawn_kwargs = {}
spawn_kwargs["start_method"] = "fork"
spawn_kwargs["nproc_per_node"] = 2
config["backend"] = "nccl"

with idist.Parallel(backend=config["backend"], **spawn_kwargs) as parallel:
    parallel.run(training, config)

2021-09-14 19:15:15,335 ignite.distributed.launcher.Parallel INFO: Initialized distributed launcher with backend: 'nccl'
2021-09-14 19:15:15,337 ignite.distributed.launcher.Parallel INFO: - Parameters to spawn processes: 
	nproc_per_node: 2
	nnodes: 1
	node_rank: 0
	start_method: fork
2021-09-14 19:15:15,338 ignite.distributed.launcher.Parallel INFO: Spawn function '' in 2 processes
2021-09-14 19:15:18,910 CIFAR10-Training INFO: Train on CIFAR10
2021-09-14 19:15:18,911 CIFAR10-Training INFO: - PyTorch version: 1.9.0
2021-09-14 19:15:18,912 CIFAR10-Training INFO: - Ignite version: 0.4.6
2021-09-14 19:15:18,913 CIFAR10-Training INFO: - GPU Device: GeForce GTX 1080 Ti
2021-09-14 19:15:18,913 CIFAR10-Training INFO: - CUDA version: 11.1
2021-09-14 19:15:18,914 CIFAR10-Training INFO: - CUDNN version: 8005
2021-09-14 19:15:18,915 CIFAR10-Training INFO: 

2021-09-14 19:15:18,916 CIFAR10-Training INFO: Configuration:
2021-09-14 19:15:18,917 CIFAR10-Training INFO: 	seed: 543
2021-09-14 19:15:18,918 CIFAR10-Training INFO: 	data_path: cifar10
2021-09-14 19:15:18,919 CIFAR10-Training INFO: 	output_path: output-cifar10/
2021-09-14 19:15:18,920 CIFAR10-Training INFO: 	model: resnet18
2021-09-14 19:15:18,921 CIFAR10-Training INFO: 	batch_size: 512
2021-09-14 19:15:18,922 CIFAR10-Training INFO: 	momentum: 0.9
2021-09-14 19:15:18,923 CIFAR10-Training INFO: 	weight_decay: 0.0001
2021-09-14 19:15:18,924 CIFAR10-Training INFO: 	num_workers: 2
2021-09-14 19:15:18,925 CIFAR10-Training INFO: 	num_epochs: 5
2021-09-14 19:15:18,925 CIFAR10-Training INFO: 	learning_rate: 0.4
2021-09-14 19:15:18,926 CIFAR10-Training INFO: 	num_warmup_epochs: 1
2021-09-14 19:15:18,927 CIFAR10-Training INFO: 	validate_every: 3
2021-09-14 19:15:18,928 CIFAR10-Training INFO: 	checkpoint_every: 200
2021-09-14 19:15:18,929 CIFAR10-Training INFO: 	backend: nccl
2021-09-14 19:15:18,929 CIFAR10-Training INFO: 	resume_from: None
2021-09-14 19:15:18,930 CIFAR10-Training INFO: 	log_every_iters: 15
2021-09-14 19:15:18,931 CIFAR10-Training INFO: 	nproc_per_node: None
2021-09-14 19:15:18,931 CIFAR10-Training INFO: 	with_clearml: False
2021-09-14 19:15:18,932 CIFAR10-Training INFO: 	with_amp: False
2021-09-14 19:15:18,933 CIFAR10-Training INFO: 

2021-09-14 19:15:18,933 CIFAR10-Training INFO: 
Distributed setting:
2021-09-14 19:15:18,934 CIFAR10-Training INFO: 	backend: nccl
2021-09-14 19:15:18,935 CIFAR10-Training INFO: 	world size: 2
2021-09-14 19:15:18,935 CIFAR10-Training INFO: 

2021-09-14 19:15:18,936 CIFAR10-Training INFO: Output path: output-cifar10/resnet18_backend-nccl-2_20210914-191518
2021-09-14 19:15:19,725 ignite.distributed.auto.auto_dataloader INFO: Use data loader kwargs for dataset 'Dataset CIFAR10': 
	{'batch_size': 256, 'num_workers': 1, 'drop_last': True, 'sampler': , 'pin_memory': True}
2021-09-14 19:15:19,727 ignite.distributed.auto.auto_dataloader INFO: Use data loader kwargs for dataset 'Dataset CIFAR10': 
	{'batch_size': 512, 'num_workers': 1, 'sampler': , 'pin_memory': True}
2021-09-14 19:15:19,873 ignite.distributed.auto.auto_model INFO: Apply torch DistributedDataParallel on model, device id: 0
2021-09-14 19:15:20,049 CIFAR10-Training INFO: Engine run starting with max_epochs=5.
/opt/conda/lib/python3.7/site-packages/torch/nn/functional.py:718: UserWarning: Named tensors and all their associated APIs are an experimental feature and subject to change. Please do not use them for anything important until they are released as stable. (Triggered internally at  /opt/conda/conda-bld/pytorch_1623448265233/work/c10/core/TensorImpl.h:1156.)
  return torch.max_pool2d(input, kernel_size, stride, padding, dilation, ceil_mode)
/opt/conda/lib/python3.7/site-packages/torch/nn/functional.py:718: UserWarning: Named tensors and all their associated APIs are an experimental feature and subject to change. Please do not use them for anything important until they are released as stable. (Triggered internally at  /opt/conda/conda-bld/pytorch_1623448265233/work/c10/core/TensorImpl.h:1156.)
  return torch.max_pool2d(input, kernel_size, stride, padding, dilation, ceil_mode)
2021-09-14 19:15:28,800 CIFAR10-Training INFO: Epoch[1] Complete. Time taken: 00:00:09
2021-09-14 19:15:37,474 CIFAR10-Training INFO: Epoch[2] Complete. Time taken: 00:00:09
2021-09-14 19:15:54,675 CIFAR10-Training INFO: 
Epoch 3 - Evaluation time (seconds): 8.50 - train metrics:
    Accuracy: 0.5533988402061856
	Loss: 1.2227583423103254
2021-09-14 19:15:56,077 CIFAR10-Training INFO: 
Epoch 3 - Evaluation time (seconds): 1.36 - val metrics:
    Accuracy: 0.5699
	Loss: 1.1869916015625
2021-09-14 19:15:56,079 CIFAR10-Training INFO: Epoch[3] Complete. Time taken: 00:00:19
2021-09-14 19:16:04,686 CIFAR10-Training INFO: Epoch[4] Complete. Time taken: 00:00:09
2021-09-14 19:16:13,347 CIFAR10-Training INFO: Epoch[5] Complete. Time taken: 00:00:09
2021-09-14 19:16:21,857 CIFAR10-Training INFO: 
Epoch 5 - Evaluation time (seconds): 8.46 - train metrics:
    Accuracy: 0.6584246134020618
	Loss: 0.9565292830319748
2021-09-14 19:16:23,269 CIFAR10-Training INFO: 
Epoch 5 - Evaluation time (seconds): 1.38 - val metrics:
    Accuracy: 0.6588
	Loss: 0.9517111328125
2021-09-14 19:16:23,271 CIFAR10-Training INFO: Engine run complete. Time taken: 00:01:03
2021-09-14 19:16:23,547 ignite.distributed.launcher.Parallel INFO: End of run

 

重要なリンク

1. 完全なコードは ここ で見つかります。
   

2. Example of the logs of a ClearML experiment run on this code:
   • With torch.distributed.launch
   • With default internal spawning
   • On Jupyter
   • On Colab with XLA

以上