はじめに

CUDAの最新バージョン12.8が公開されていたので早速飛びついてみようと思ったのですが、そこでふと気になったのが「CUDAのバージョンを上げることでPyTorchの計算速度は速くなるかどうか」です。

結論から言うとタイトルの通り、CUDAのバージョンの違いで計算速度は大して変わりませんでした。

本記事ではその測定を記録します。

今回は、現時点でリリースされている最新のCUDAが12.8であることと、PyTorchのの対応CUDAの組み合わせから4通りの場合で計算速度を計測してみます。

計測した組み合わせ

CUDA11.8 - PyTorch(CUDA11.8)
CUDA12.4 - PyTorch(CUDA12.4)
CUDA12.6 - PyTorch(CUDA12.6)
CUDA12.8 - PyTorch(CUDA12.6)

実行環境

WSL2
Ubuntu 22.04 (24.04ではCUDA12.4、11.8がうまくインストールできなかったため)
Python 3.10
PyTorch 2.5.1

＃##計測に使用したタスク

こちらのサンプルコードへ時間測定を付け加えました。
計測対象はEpochをループさせているところです。
それぞれのCUDA環境で学習を10回実行し、かかった時間を平均します。

実行コード

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import datasets, transforms

import time

#----------------------------------------------------------
# ハイパーパラメータなどの設定値
num_epochs = 10         # 学習を繰り返す回数
num_batch = 100         # 一度に処理する画像の枚数
learning_rate = 0.001   # 学習率
image_size = 28*28      # 画像の画素数(幅x高さ)

# GPU(CUDA)が使えるかどうか？
device = 'cuda' if torch.cuda.is_available() else 'cpu'

#----------------------------------------------------------
# 学習用／評価用のデータセットの作成

# 変換方法の指定
transform = transforms.Compose([
    transforms.ToTensor()
    ])

# MNISTデータの取得
# https://pytorch.org/vision/stable/generated/torchvision.datasets.MNIST.html#torchvision.datasets.MNIST
# 学習用
train_dataset = datasets.MNIST(
    './data',               # データの保存先
    train = True,           # 学習用データを取得する
    download = True,        # データが無い時にダウンロードする
    transform = transform   # テンソルへの変換など
    )
# 評価用
test_dataset = datasets.MNIST(
    './data', 
    train = False,
    transform = transform
    )

# データローダー
train_dataloader = torch.utils.data.DataLoader(
    train_dataset,
    batch_size = num_batch,
    shuffle = True)
test_dataloader = torch.utils.data.DataLoader(
    test_dataset,     
    batch_size = num_batch,
    shuffle = True)

#----------------------------------------------------------
# ニューラルネットワークモデルの定義
class Net(nn.Module):
    def __init__(self, input_size, output_size):
        super(Net, self).__init__()

        # 各クラスのインスタンス（入出力サイズなどの設定）
        self.fc1 = nn.Linear(input_size, 100)
        self.fc2 = nn.Linear(100, output_size)

    def forward(self, x):
        # 順伝播の設定（インスタンスしたクラスの特殊メソッド(__call__)を実行）
        x = self.fc1(x)
        x = torch.sigmoid(x)
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

times = []
for iter_ in range(10):
    #----------------------------------------------------------
    # ニューラルネットワークの生成
    model = Net(image_size, 10).to(device)

    #----------------------------------------------------------
    # 損失関数の設定
    criterion = nn.CrossEntropyLoss() 

    #----------------------------------------------------------
    # 最適化手法の設定
    optimizer = torch.optim.Adam(model.parameters(), lr = learning_rate) 

    #----------------------------------------------------------
    # 学習
    model.train()  # モデルを訓練モードにする

    start_time_stamp = time.time()

    for epoch in range(num_epochs): # 学習を繰り返し行う
        loss_sum = 0

        for inputs, labels in train_dataloader:

            # GPUが使えるならGPUにデータを送る
            inputs = inputs.to(device)
            labels = labels.to(device)

            # optimizerを初期化
            optimizer.zero_grad()

            # ニューラルネットワークの処理を行う
            inputs = inputs.view(-1, image_size) # 画像データ部分を一次元へ並び変える
            outputs = model(inputs)

            # 損失(出力とラベルとの誤差)の計算
            loss = criterion(outputs, labels)
            loss_sum += loss

            # 勾配の計算
            loss.backward()

            # 重みの更新
            optimizer.step()

        # 学習状況の表示
        print(f"Epoch: {epoch+1}/{num_epochs}, Loss: {loss_sum.item() / len(train_dataloader)}")

        # モデルの重みの保存
        torch.save(model.state_dict(), 'model_weights.pth')

    time_span = time.time() - start_time_stamp
    print(time_span)
    times.append(time_span)

print(sum(times) / 10)

測定結果

環境	計算時間(秒)
CUDA11.8 - PyTorch(CUDA11.8)	30.3247
CUDA12.4 - PyTorch(CUDA12.4)	31.4305
CUDA12.6 - PyTorch(CUDA12.6)	29.8244
CUDA12.8 - PyTorch(CUDA12.6)	29.5854

結論

CUDAのバージョンを上げるほど若干計算時間が短くなっているように見られますが、多分誤差の範囲かなといった感想です。
CUDAバージョンどうこうするよりも、アルゴリズムの最適化や強いGPUを買うとよいと思います。