1,010 views

この記事は最終更新から 397日 が経過しています。

【1】やりたいこと

PyTorchを GPU実行 vs CPU実行で速度性能を比較してみたい。
やる前から結果は明らかだが・・・

実行環境は以下の通り。

下記の数点だけプログラムを変更した。
・GPU or CPU 使用中のデバイスを表示する。
・所要時間を測定、表示する。
・ミニバッチごとの結果表示を抑制する。
・使わないコマンドライン引数を無視する。
・その他、不要な処理を削除した。

import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from torch.optim.lr_scheduler import StepLR
import time

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout(0.25)
        self.dropout2 = nn.Dropout(0.5)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = self.dropout1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dropout2(x)
        x = self.fc2(x)
        output = F.log_softmax(x, dim=1)
        return output

def train(args, model, device, train_loader, optimizer, epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()

def test(model, device, test_loader):
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += F.nll_loss(output, target, reduction='sum').item()  # sum up batch loss
            pred = output.argmax(dim=1, keepdim=True)  # get the index of the max log-probability
            correct += pred.eq(target.view_as(pred)).sum().item()
    test_loss /= len(test_loader.dataset)
    print('Test set: Average loss: {:.4f}, Accuracy: {}/{} ({:.2f}%)'.format(
        test_loss, correct, len(test_loader.dataset), 100. * correct / len(test_loader.dataset)))

def main():
    parser = argparse.ArgumentParser(description='PyTorch MNIST Example')
    parser.add_argument('--epochs', type=int, default=3, metavar='N',    help='number of epochs to train (default: 14)')
    parser.add_argument('--no-cuda', action='store_true', default=False, help='disables CUDA training')
    args = parser.parse_args()
    use_cuda = not args.no_cuda and torch.cuda.is_available()
    if use_cuda:
        device = torch.device("cuda")
    else:
        device = torch.device("cpu")
    print("use device is {}".format(device))
    train_kwargs = {'batch_size': 64}
    test_kwargs  = {'batch_size': 1000}
    start_time = time.time()
    if use_cuda:
        cuda_kwargs = {'num_workers': 1, 'pin_memory': True, 'shuffle': True}
        train_kwargs.update(cuda_kwargs)
        test_kwargs.update(cuda_kwargs)
    transform=transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
        ])
    dataset1 = datasets.MNIST('../data', train=True, download=True, transform=transform)
    dataset2 = datasets.MNIST('../data', train=False,               transform=transform)
    train_loader = torch.utils.data.DataLoader(dataset1,**train_kwargs)
    test_loader  = torch.utils.data.DataLoader(dataset2, **test_kwargs)
    model = Net().to(device)
    optimizer = optim.Adadelta(model.parameters(), lr=1.0)
    scheduler = StepLR(optimizer, step_size=1, gamma=0.7)
    for epoch in range(1, args.epochs + 1):
        train(args, model, device, train_loader, optimizer, epoch)
        test(model, device, test_loader)
        scheduler.step()
    end_time = time.time()
    elapsed_time = round(end_time - start_time, 1)
    print(f"処理時間: {elapsed_time} 秒")

if __name__ == '__main__':
    main()

【2】やってみた

(1) GPU実行（RTX 5070ti）

(myPyTorch) $ python main.py --epochs 1
use device is cuda
Test set: Average loss: 0.0536, Accuracy: 9831/10000 (98.31%)
処理時間: 3.2 秒

(2) CPU実行（i9 13900）

同じプログラムを、コマンドライオプション --no-cuda 指定で実行すればよい。

(myPyTorch) $ python main.py --no-cuda --epochs 1
use device is cpu
Test set: Average loss: 0.0487, Accuracy: 9837/10000 (98.37%)
処理時間: 20.0 秒

GPU実行の 6倍以上の時間がかかった。

リソースモニターを見ると、使用している CPU i9 13900 の cpu coreをほぼフルに使っての結果だ。

【3】所感

・CPUの場合、メモリ容量を簡単に大きくできるというメリットがある。
　GPUのメモリは、18万円もする RTX 5070tiでさえ 16GBしかないのだ。
　ネットワーク構成が大きくなって GPUメモリが不足した場合、CPUと GPUのヘテロな組み合わせでシステムを組まなければならないかも。