Diffusersの訓練時間による、TITAN VとRTX A4000比較

モチベーション

少し前に、この記事でStable Diffusionについて紹介した。また、この記事では、NVIDIA TITAN VとRTX A4000との実行時間比較を行った。この記事では、Stable Difffusionの学習時間で、TITAN VとA4000との実行時間比較を行う。

情報源

classcatさんのPyTorchのページ Diffusersによる訓練の解説（翻訳）記事。個々の情報を元に学習／手を動かしている、とても有益なページ群。
Training with Diffusers 上記記事の元記事。

ハンズオン

基本は、情報源の項番1で紹介したclasscatさんのページの内容を写経。写経しながらでも、処理の流れを知ることでき、学習になる。３箇所のセルだけ変更したので、その部分のみ、理由を含め以下に掲載する。

データセットの選択とHugging Faceへのアップ指定

以下のセルでは、次の2点を変更した。

output_dir = ‘ddmp-flowers-128’ 蝶の翅の紋様は好きではないので、花のデータセットに変更。
push_to_hub = False 「save_model_epochs」毎にHuggingFaceのページにモデルアップする仕組みとなっているのだが、アップ時にエラーとなった（TITAN VでもRTX A4000でも）ので、アップしないように変更した。

from dataclasses import dataclass

@dataclass
class TrainingConfig:
    image_size = 128  # the generated image resolution
    train_batch_size = 16
    eval_batch_size = 16  # how many images to sample during evaluation
    num_epochs = 50
    gradient_accumulation_steps = 1
    learing_rate = 1e-4
    lr_warmup_steps = 500
    save_image_epochs = 10
    save_model_epochs = 30
    mixed_precision = 'fp16'  # `no` for float32, `fp16` for automatic mixed precision
    output_dir = 'ddmp-flowers-128'  # the model name locally and on the HF Hub
    
    push_to_hub = False  # whether to upload the saved model to the HF Hub. Change to False.
    hub_private_repo = False
    overwrite_output_dir = True  # overwrite the old model when re-running the notebook
    seed = 0
    
config = TrainingConfig()

DDPMSchedulerのtensor_formatパラメータ

「tensor_format="pt”」を指定すると、エラーとなった。githubの元記事（英語版）を見ると、そちらもこのパラメータは無かった。

from diffusers import DDPMScheduler

# noise_scheduler = DDPMScheduler(num_train_timesteps=1000, tensor_format="pt")
noise_scheduler = DDPMScheduler(num_train_timesteps=1000)

evaluate関数内のpipelingの部分

evaluate関数で、imagesを得る部分の「sample」を「images」に変更。githubの元記事（英語版）はimagesであった。

from diffusers import DDPMPipeline

import math

def make_grid(images, rows, cols):
    w, h = images[0].size
    grid = Image.new('RGB', size=(cols*w, rows*h))
    for i, image in enumerate(images):
        grid.paste(image, box=(i%cols*w, i//cols*h))
    return grid

def evaluate(config, epoch, pipeline):
    # Sample some images from random noise (this is the backward diffusion process).
    # The default pipeline output type is `List[PIL.Image]`
    images = pipeline(
        batch_size = config.eval_batch_size,
        generator=torch.manual_seed(config.seed),
    ).images  # ←　sample
    
    # Make a grid out of the images
    image_grid = make_grid(images, rows=4, cols=4)
    
    # Save the images
    test_dir = os.path.join(config.output_dir, "samples")
    os.makedirs(test_dir, exist_ok=True)
    image_grid.save(f"{test_dir}/{epoch:04d}.png")

実行時間の計測

実行時間の計測は、次のセルの前後で、タイマーを開始／停止することで経過時間を測った。

from accelerate import notebook_launcher
args = (config, model, noise_scheduler, optimizer, train_dataloader, lr_scheduler)

notebook_launcher(train_loop, args, num_processes=1)

タイマーを開始

上記のセルの前に以下のセルでタイマーをスタート。

# 訓練時間を測定：開始
import time

print("*** Starting the Timer ***")
start_time = time.time()  # 実行時間計測開始

タイマーを停止

# 訓練時間測定：終了

lapse_time = time.time() - start_time
print("-" * 80)
print("実行時間 {:8.2f}秒".format(lapse_time))
print("-" * 80)

結果

TITAN VとA4000での実行時間は、次のとおり。

GPU	実行時間（秒）	実行時間（時分秒）	実行時間比
NVIDIA TITAN V	12,683	3時間32分24秒	1.00
RTX A4000	19,683	5時間28分3秒	1.54

上記の結果は、あくまで自分の環境（各 GPUの動作サーバも異なる）での結果であり、厳密な性能測定ではない。

RTX 1080でも実行したが、「CUDA out of memory」となって実行できなかった。この学習では10GB程度のGPUメモリーが必要なようだ。