概要

PyTorchの最適化のためのガイドラインを読み、実際にどれくらい高速化されるのかを検証しました。
その結果を共有します。

色々試したところ、次の3つが良さそうに思いました。

DataLoaderのオプションで、pin_memory=True とする
torch.compile を使う
- hugging faceとかのモデルでかなり強力に高速化できる
Automatic Mixed Precision を利用する
- nn.Module で作成したモデルに対して利用するのが良い

実験条件

PyTorchの最適化のためのガイドラインに書かれているもののうち、後述の5パターンについて検証し、処理速度を確認します。

処理速度を検証するために、以下の機械学習を実行させることとしました。

利用データセット：CIFAR10
モデル：ResNet50
学習エポック数：5（精度を出すのが目的ではないため、エポック数は少なくしました）
バッチサイズ：128
学習率：0.01

処理速度と合わせて、精度も確認し、高速化のために、精度が犠牲になっていないかを確認しました。

検証パターン

①DataLoader の pin_memory の効果
②zero_grad(set_to_none=True) の効果
③計算時の精度を変更する（AMPを利用する）ことの効果
④torch.compile によるカーネル融合の効果
⑤hugging faceモデルでのtorch.compileの効果

利用するライブラリ

"transformers",
"torch",
"accelerate",
"huggingface-hub",
"sentencepiece",
"matplotlib",
"torchvision",

実験

今回作成したプログラムは、Githubに掲載しています。

①DataLoader の `pin_memory` の効果

pin_memoryの意味と効果

Pinned memory（固定メモリ）とは、CPU 側のメモリ領域を GPU に対して直接アクセス可能にすることで、転送速度を向上させる技術です。
通常、CPU から GPU へのメモリ転送は非同期に行えませんが、pinned memory を使うと非同期転送が可能になります。
PyTorch では、DataLoader に pin_memory=True を設定すると、バッチデータが GPU にコピーされる際の転送が高速かつ非同期になります。

なぜ重要か？

GPU の計算速度に対して、データ転送がボトルネックになることが多いため、pin_memory を使うことでそのボトルネックを緩和できます。
特に大規模なデータセットや高速な GPU を使う場合、pin_memory=True にすることでトレーニングのスループットが向上します。

これを実験するためのコードのうち、実際に切り替えているのは以下です。

def setup_dataloader(is_pin: bool = False, dataset_path:str="./data") -> tuple[DataLoader, DataLoader]:
    """
    CIFAR-10 のデータセットを DataLoader として返す。
    """
    train_dataset, test_dataset = setup_dataset(dataset_path)

    train_loader = DataLoader(
        train_dataset,
        batch_size=BATCH_SIZE,
        shuffle=True,
        pin_memory=is_pin,
        num_workers=4,
    )
    test_loader = DataLoader(
        test_dataset,
        batch_size=BATCH_SIZE,
        shuffle=False,
        pin_memory=is_pin,
        num_workers=4,
    )

    return train_loader, test_loader

pin_memory=Falseの条件で実行速度を計測する

上記関数で、is_pin=Falseとした場合の実行速度は以下の通り。

train_epoch(pin=False): 53953.38 ms

train_epoch(pin=False): 52372.03 ms

train_epoch(pin=False): 52443.23 ms

train_epoch(pin=False): 52430.18 ms

train_epoch(pin=False): 52566.57 ms

loss: 2.595525830078125, acc: 0.1267
train_epoch(pin=False): n=5, mean=52753.08 ms, std=674.72 ms, min=52372.03 ms, max=53953.38 ms

1エポックの学習はざっくり平均して、52.7秒
精度は、accで0.1267

pin_memory=Trueの条件で実行速度を計測する

train_epoch(pin=True): 51903.59 ms

train_epoch(pin=True): 51913.30 ms

train_epoch(pin=True): 51931.16 ms

train_epoch(pin=True): 51883.39 ms

train_epoch(pin=True): 51928.87 ms

loss: 2.6195204555511475, acc: 0.1112
train_epoch(pin=True): n=5, mean=51912.06 ms, std=19.64 ms, min=51883.39 ms, max=51931.16 ms

1エポックの学習はざっくり平均して、51.9秒
精度は、accで0.1112

結論

今回の実験条件ではデータサイズが大きくなかったかつ量も多くなかったためか、大きな速度向上は見られなかったです。しかし、pin_memory=Trueとしただけで、1秒程度の効率化ができるなら、やっておくと良さそうですね。

②`zero_grad(set_to_none=True)` の効果

set_to_none=True は、PyTorch の optimizer.zero_grad() や model.zero_grad() に渡すオプション引数で、勾配の初期化方法を変更するために使われます。

通常の zero_grad() は、各パラメータの .grad をゼロで埋めます。

optimizer.zero_grad()  # 各 param.grad をゼロで初期化

これに対して set_to_none=True を指定すると、

optimizer.zero_grad(set_to_none=True)  # 各 param.grad を None に設定

この違いは、メモリ操作の回数と勾配の蓄積方法に影響します。

方法	メモリ操作	backward時の挙動	パフォーマンス
`zero_grad()`	メモリをゼロで埋める	勾配を加算 (`+=`)	やや遅い
`zero_grad(set_to_none=True)`	メモリを解放（None）	勾配を代入 (`=`)	より高速

メモリ効率：None にすることで、前回の勾配メモリを解放できるため、メモリ使用量が減少。
計算効率：次の backward で = による代入が使われるため、加算よりも高速。
数値的挙動：加算と代入では微妙に異なる数値誤差が生じる可能性があるため、再現性が重要な場合は注意が必要。

これを比較するための切り替えをしているのは以下。

def train_1epoch(
    model,
    criterion,
    optimizer,
    scheduler,
    train_loader,
):
    metrics = {
        "loss": AverageMetric(),
        "acc": AverageMetric(),
    }
    model.train()
    for idx, (Xs, ys) in enumerate(train_loader):
        Xs, ys = Xs.to(DEVICE), ys.to(DEVICE)
        optimizer.zero_grad()
        outputs = model(Xs)
        loss = criterion(outputs, ys)
        loss.backward()

        if hasattr(optimizer, "step"):
            optimizer.step()
        if hasattr(scheduler, "step"):
            scheduler.step()

        # Update metrics
        num_samples = ys.size(0)
        metrics["loss"].update(loss.item(), n=num_samples)
        preds = outputs.argmax(dim=1)
        metrics["acc"].update((preds == ys).float().mean().item(), n=num_samples)

    return metrics


def train_1epoch_with_nograd(
    model,
    criterion,
    optimizer,
    scheduler,
    train_loader,
):
    """zero_grad()で、オプション`set_to_none=True`を使用して勾配をNoneにする"""
    metrics = {
        "loss": AverageMetric(),
        "acc": AverageMetric(),
    }
    model.train()
    for idx, (Xs, ys) in enumerate(train_loader):
        Xs, ys = Xs.to(DEVICE), ys.to(DEVICE)
        optimizer.zero_grad(set_to_none=True) # ここで勾配をゼロにする.
        outputs = model(Xs)
        loss = criterion(outputs, ys)
        loss.backward()

        if hasattr(optimizer, "step"):
            optimizer.step()
        if hasattr(scheduler, "step"):
            scheduler.step()

        # Update metrics
        num_samples = ys.size(0)
        metrics["loss"].update(loss.item(), n=num_samples)
        preds = outputs.argmax(dim=1)
        metrics["acc"].update((preds == ys).float().mean().item(), n=num_samples)

    return metrics


def eval(
    model,
    criterion,
    test_loader,
):
    model.eval()
    metrics = {
        "loss": AverageMetric(),
        "acc": AverageMetric(),
    }
    with torch.no_grad():
        for idx, (Xs, ys) in enumerate(test_loader):
            Xs, ys = Xs.to(DEVICE), ys.to(DEVICE)
            outputs = model(Xs)
            loss = criterion(outputs, ys)

            # Update metrics
            num_samples = ys.size(0)
            metrics["loss"].update(loss.item(), n=num_samples)
            preds = outputs.argmax(dim=1)
            metrics["acc"].update((preds == ys).float().mean().item(), n=num_samples)

    return metrics

通常の学習条件で実行速度を計測する

noraml_training: 78488.78 ms

noraml_training: 76510.33 ms

noraml_training: 77398.08 ms

noraml_training: 76064.51 ms

noraml_training: 75450.02 ms

loss: 2.965869669342041, acc: 0.1271
noraml_training: n=5, mean=76782.34 ms, std=1189.06 ms, min=75450.02 ms, max=78488.78 ms

1エポックの学習はざっくり平均して、76.8秒
精度は、accで0.1271

オプション`set_to_none=True`条件で実行速度を計測する

set_to_none=True: 75386.75 ms

set_to_none=True: 76258.83 ms

set_to_none=True: 75747.22 ms

set_to_none=True: 75894.83 ms

set_to_none=True: 75559.26 ms

loss: 3.168150677108765, acc: 0.1341
set_to_none=True: n=5, mean=75769.38 ms, std=334.05 ms, min=75386.75 ms, max=76258.83 ms

1エポックの学習はざっくり平均して、75.8秒
精度は、accで0.1341

結論

約1秒高速化されました。
set_to_none=Trueというオプションを使うだけで高速化されるので、実務では特に理由がない限りは基本的に使えば良いのではないかと考えられます。

③計算時の精度を変更する（AMPを利用する）ことの効果

AMP（Automatic Mixed Precision）は、PyTorchで導入されている混合精度学習の仕組みで、主にGPU上での高速化とメモリ効率の向上を目的としています。特にVolta以降のNVIDIA GPUでTensor Coresを活用する際に非常に効果的です。

具体的には、演算の一部をFP16（16ビット浮動小数点）で行い、精度が重要な部分はFP32（32ビット浮動小数点）で維持する手法です。

AMPは、特に以下のような状況で重要な役割を果たします：

大規模モデルの学習：Transformer系やCNNなど、メモリ消費が激しいモデルにおいて、AMPはリソースの節約と高速化を両立。
限られたGPU環境：メモリが少ないGPUでも、AMPを使えばより大きなモデルやバッチサイズでの学習が可能。
実運用での効率化：推論時にもAMPを使うことで、リアルタイム性が求められるアプリケーションのレスポンス向上に貢献。

これを切り替えるための実装部分は以下です。

def train_1epoch(
    model,
    criterion,
    optimizer,
    scheduler,
    train_loader,
):
    metrics = {
        "loss": AverageMetric(),
        "acc": AverageMetric(),
    }
    model.train()
    for idx, (Xs, ys) in enumerate(train_loader):
        Xs, ys = Xs.to(DEVICE), ys.to(DEVICE)
        optimizer.zero_grad()
        outputs = model(Xs)
        loss = criterion(outputs, ys)
        loss.backward()

        if hasattr(optimizer, "step"):
            optimizer.step()
        if hasattr(scheduler, "step"):
            scheduler.step()

        # Update metrics
        num_samples = ys.size(0)
        metrics["loss"].update(loss.item(), n=num_samples)
        preds = outputs.argmax(dim=1)
        metrics["acc"].update((preds == ys).float().mean().item(), n=num_samples)

    return metrics


def train_1epoch_with_amp(
    model,
    criterion,
    optimizer,
    scheduler,
    train_loader,
):
    """"""
    metrics = {
        "loss": AverageMetric(),
        "acc": AverageMetric(),
    }
    model.train()
    for idx, (Xs, ys) in enumerate(train_loader):
        Xs, ys = Xs.to(DEVICE), ys.to(DEVICE)
        optimizer.zero_grad()
        with autocast("cuda"):
            outputs = model(Xs)
            loss = criterion(outputs, ys)
        loss.backward()

        if hasattr(optimizer, "step"):
            optimizer.step()
        if hasattr(scheduler, "step"):
            scheduler.step()

        # Update metrics
        num_samples = ys.size(0)
        metrics["loss"].update(loss.item(), n=num_samples)
        preds = outputs.argmax(dim=1)
        metrics["acc"].update((preds == ys).float().mean().item(), n=num_samples)

    return metrics


def eval(
    model,
    criterion,
    test_loader,
):
    model.eval()
    metrics = {
        "loss": AverageMetric(),
        "acc": AverageMetric(),
    }
    with torch.no_grad():
        for idx, (Xs, ys) in enumerate(test_loader):
            Xs, ys = Xs.to(DEVICE), ys.to(DEVICE)
            outputs = model(Xs)
            loss = criterion(outputs, ys)

            # Update metrics
            num_samples = ys.size(0)
            metrics["loss"].update(loss.item(), n=num_samples)
            preds = outputs.argmax(dim=1)
            metrics["acc"].update((preds == ys).float().mean().item(), n=num_samples)

    return metrics

通常の学習条件で実行速度を計測する

noraml_training: 75964.93 ms

noraml_training: 75450.22 ms

noraml_training: 75037.71 ms

noraml_training: 75579.37 ms

noraml_training: 74614.61 ms

loss: 2.700083690261841, acc: 0.1435
noraml_training: n=5, mean=75329.37 ms, std=518.89 ms, min=74614.61 ms, max=75964.93 ms

1エポックの学習はざっくり平均して、75.3秒
精度は、accで0.1435

追加実験）Automatic Mixed Precisionを使用した場合の条件で実行速度を計測する

amp: 44864.51 ms

amp: 44699.03 ms

amp: 44609.77 ms

amp: 43328.46 ms

amp: 43420.54 ms

loss: 3.6950550128936768, acc: 0.1259
amp: n=5, mean=44184.46 ms, std=745.73 ms, min=43328.46 ms, max=44864.51 ms

1エポックの学習はざっくり平均して、44.2秒
精度は、accで0.1259

結論

30秒以上の高速化することができました。
これは非常に大きな効果ですね。

④`torch.compile` によるカーネル融合の効果

PyTorch 2.0以降で導入されたモデルの高速化機能であり、従来の「eager execution」モードに対して、TorchInductorというコンパイラを用いてモデルを最適化・JITコンパイルすることで、トレーニングや推論の速度を大幅に向上させることができます。

compiled_model = torch.compile(model)

このように使うことで、PyTorchのnn.ModuleをTorchInductorを通じてコンパイルし、以下のような最適化が自動的に行われます：

カーネル融合：複数の演算（例：加算、乗算、活性化関数など）を1つのGPUカーネルにまとめて実行
メモリアクセスの削減：中間テンソルの読み書きを減らし、メモリ帯域のボトルネックを緩和
CUDA Graphsの活用：GPU上でのカーネル起動オーバーヘッドを削減
自動チューニング：最適な実行パスを探索し、ハードウェアに合わせた最適化を実施

※注意点：初回実行時はコンパイル時間がかかるため遅くなるが、2回目以降は大幅に高速化されます。

これを実装しているのは以下です。

def setup_elements(is_compile: bool=False):
    model = torchvision.models.resnet50(weights=None)

    # CIFAR-10は画像サイズが32x32と小さいため、最初の畳み込み層とプーリング層を調整
    model.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False)
    model.maxpool = nn.Identity() # MaxPool層を無効化

    # ResNet-50の最終層（全結合層）をCIFAR-10の10クラス分類用に変更
    num_ftrs = model.fc.in_features
    model.fc = nn.Linear(num_ftrs, 10)

    model = model.to(DEVICE)
    if is_compile and hasattr(torch, 'compile'):
        model = torch.compile(model, mode="default")

    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=LEARNING_RATE, momentum=0.9, weight_decay=5e-4)

    # 学習率をスケジューリング（例：10エポック毎に学習率を0.1倍にする）
    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)

    return model, criterion, optimizer, scheduler

is_compile=Falseの条件で実行速度を計測する

is_compile=False: 76552.48 ms

is_compile=False: 76481.43 ms

is_compile=False: 76627.69 ms

is_compile=False: 76244.26 ms

is_compile=False: 75998.18 ms

loss: 2.9309064430236815, acc: 0.1229
is_compile=False: n=5, mean=76380.81 ms, std=257.68 ms, min=75998.18 ms, max=76627.69 ms

1エポックの学習はざっくり平均して、76.4秒
精度は、accで0.1229

is_compile=Trueの条件で実行速度を計測する

is_compile=True: 85152.46 ms

is_compile=True: 72219.93 ms

is_compile=True: 70659.39 ms

is_compile=True: 70803.32 ms

is_compile=True: 70511.82 ms

loss: 3.0611061866760254, acc: 0.1373
is_compile=True: n=5, mean=73869.38 ms, std=6344.42 ms, min=70511.82 ms, max=85152.46 ms

1エポックの学習はざっくり平均して、73.9秒
精度は、accで0.1373

結論

初回コンパイルだけ時間がかかっているが、平均すれば、約3秒の高速化ができていました。さらに、2回目以降の実行では、70~72秒で実行できており、compileしていない場合の平均と比較して、3~6秒ほど高速化できている点に注目です。
これについても特別な理由がない限りは、基本的にtorch.compile()を利用するのが良いと考えられます。

⑤hugging faceモデルでの`torch.compile`の効果

torch.compileはhugging faceのモデルでも利用できるとのことなので、それを実験しました。
この実験だけ、モデルを学習させるわけではなく、推論速度を計測します。

#SakanaAI/TinySwallow-1.5B
model_name = "SakanaAI/TinySwallow-1.5B"
cache_directory = "./model_cache"


def load_elements(is_compile:bool=False):
    tokenizer = AutoTokenizer.from_pretrained(
        model_name,
        cache_dir=cache_directory
    )
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        torch_dtype="auto",
        device_map="auto",
        cache_dir=cache_directory
    )

    if is_compile:
        try:
            model = torch.compile(model, mode="reduce-overhead", fullgraph=True)
        except Exception as e:
            try:
                model = torch.compile(model, mode="reduce-overhead")
            except Exception as e:
                print(f"Error during torch.compile: {e}")
                print("Falling back to eager execution.")
    return model, tokenizer

compileなしで実行速度を計測

is_compile=False: 2466.67 ms

created by Alibaba Cloud. You are a helpful assistant.<|im_end|>
<|im_start|>user
Who are you?<|im_end|>
<|im_start|>assistant
Assistant! How can I help you today?<|im_end|>
is_compile=False: n=1, mean=2466.67 ms, std=0.00 ms, min=2466.67 ms, max=2466.67 ms

compileありで実行速度を計測

is_compile=True: 818.20 ms

created by Alibaba Cloud. You are a helpful assistant.<|im_end|>
<|im_start|>user
Who are you?<|im_end|>
<|im_start|>assistant
Hello! How can I help you today?<|im_end|>
is_compile=True: n=1, mean=818.20 ms, std=0.00 ms, min=818.20 ms, max=818.20 ms

結論

1プロンプトに対して、1/3に処理速度を早めることができました。
- compileなし：2.4秒
- compileあり：0.8秒
かなりの高速化効果があることを確認できました。

まとめ

CIFAR10とResNet50をベースに、5パターンの実験を行い、Pytorchの最適化によってどれくらい比較しました。
結果、次の3パターンを使うことが簡単に実装できて、高速化に効果がありそうということがわかりました。

DataLoaderのオプションで、pin_memory=True とする
torch.compile を使う
- hugging faceとかのモデルでかなり強力に高速化できる
Automatic Mixed Precision を利用する
- nn.Module で作成したモデルに対して利用するのが良い

また、これらを利用したときに、顕著に精度が落ちることは確認できなかったので、モデル精度を気にせず、高速化させることができる、と言えそうです。

今回実施できず、今後やりたいこととして、これらを組み合わせたときに、どれくらい高速化されるのかを実験したいと考えられます。