Zenn
😺

PyTorchをM1 MacBook のGPU(MPS)で動かす．実行時間の検証もしたよ

2022/07/31に公開
PyTorch
MacBook
GPU
mps
tech

最近，PyTorchがM1 MacBookのGPUに対応したとのことで，そのインストール方法を説明します．また，簡単に計算時間を検証してみたので，その結果を共有します．
この記事は，ある程度PyTorchを使い慣れている人向けの記事です．
さっくりとM1 MacBookのGPUを使って深層学習（ディープラーニング）を動かしてみたい人にちょうど良いかと思います．

M1 MacBookのGPUを使うためのPyTorchのインストール方法 (2023/03/30 更新)

※ 本記事初公開時はStable版がなかったのですが，2023/03/30時点でStable版(ver2.0.0)があったのでインストール方法を修正しています．

筆者はインストールのために以下のページを参考にしました．
https://pytorch.org/get-started/locally/

pipもしくはcondaでStable版のインストールが可能です．
以下のいずれかのコマンドを実行することで，インストールできます．

  • pipでインストールする場合
$ pip install torch torchvision
  • condaでインストールする場合
$ conda install pytorch torchvision -c pytorch

参考ページではtorchaudio をインストールしていましたが，今回は必要ないのでpytorchとtorchvisionだけインストールしています．なお，2023/03/30時点では，python3.11ではcondaでのインストールに失敗しましたが，python3.10.10ではうまくインストールできました．

インストールが完了したので，バージョンを確認してみます．

  • pipでインストールした場合
$ pip list | grep torch
torch              2.0.0
torchvision        0.15.1
  • condaでインストールした場合
conda list | grep torch
pytorch                   2.0.0                  py3.10_0    pytorch
torchvision               0.15.0                py310_cpu    pytorch

Pythonのインタプリタを開いて，以下のように実行できれば無事に動いています．

>>> import torch
>>> torch.backends.mps.is_available()
True

検証用のコード

それでは，早速GPUで学習してみましょう．今回は以下のようなコードを用意しました．

pytorch_m1_macbook.py
# -*- coding: utf-8 -*-
import torch
from torch.nn import CrossEntropyLoss
from torch.optim import SGD
from torch.utils.data import DataLoader
from torchvision.datasets import CIFAR10
from torchvision import transforms as tt
from torchvision.models import resnet18

import os
from argparse import ArgumentParser
import time

def main(device):
    # ResNetのハイパーパラメータ
    n_epoch = 5            # エポック数
    batch_size = 512       # ミニバッチサイズ
    momentum = 0.9         # SGDのmomentum
    lr = 0.01              # 学習率
    weight_decay = 0.00005 # weight decay

    # 訓練データとテストデータを用意
    mean = (0.491, 0.482, 0.446)
    std = (0.247, 0.243, 0.261)
    train_transform = tt.Compose([
        tt.RandomHorizontalFlip(p=0.5),
        tt.RandomCrop(size=32, padding=4, padding_mode='reflect'),
        tt.ToTensor(),
        tt.Normalize(mean=mean, std=std)
    ])
    test_transform = tt.Compose([tt.ToTensor(), tt.Normalize(mean, std)])
    root = os.path.dirname(os.path.abspath(__file__))
    train_set = CIFAR10(root=root, train=True,
                        download=True, transform=train_transform)
    train_loader = DataLoader(train_set, batch_size=batch_size,
                              shuffle=True, num_workers=8)

    # ResNetの準備
    resnet = resnet18()
    resnet.fc = torch.nn.Linear(resnet.fc.in_features, 10)

    # 訓練
    criterion = CrossEntropyLoss()
    optimizer = SGD(resnet.parameters(), lr=lr,
                    momentum=momentum, weight_decay=weight_decay)
    train_start_time = time.time()
    resnet.to(device)
    resnet.train()
    for epoch in range(1, n_epoch+1):
        train_loss = 0.0
        for inputs, labels in train_loader:
            inputs = inputs.to(device)
            optimizer.zero_grad()
            outputs = resnet(inputs)
            labels = labels.to(device)
            loss = criterion(outputs, labels)
            loss.backward()
            train_loss += loss.item()
            del loss  # メモリ節約のため
            optimizer.step()
        print('Epoch {} / {}: time = {}[s], loss = {:.2f}'.format(
            epoch, n_epoch, time.time() - train_start_time, train_loss))
    print('Train time on {}: {:.2f}[s] (Train loss = {:.2f})'.format(
        device, time.time() - train_start_time, train_loss))

    # 評価
    test_set = CIFAR10(root=root, train=False, download=True,
                       transform=test_transform)
    test_loader = DataLoader(test_set, batch_size=batch_size,
                             shuffle=False, num_workers=8)
    test_loss = 0.0
    test_start_time = time.time()
    resnet.eval()
    for inputs, labels in test_loader:
        inputs = inputs.to(device)
        outputs = resnet(inputs)
        labels = labels.to(device)
        loss = criterion(outputs, labels)
        test_loss += loss.item()
    print('Test time on {}: {:.2f}[s](Test loss = {:.2f})'.format(
        device, time.time() - test_start_time, test_loss))


if __name__ == '__main__':
    parser = ArgumentParser()
    parser.add_argument('--device', type=str, default='mps',
                        choices=['cpu', 'mps'])
    args = parser.parse_args()
    device = torch.device(args.device)
    main(device)

ResNet18を，CIFAR10で5エポックだけ学習し，テスト用データを評価するだけのシンプルなコードです．訓練時間とテスト時間の両方を，計測して表示します．

main関数内には特に変わった点はありません．一点，M1 MacBookで起動させる際に気をつける点は，以下のようにdeviceを取得する必要がある点です．実際に筆者が用意したコードとは若干違うことにご注意ください．

device = torch.device('mps')

ここで，mpsとは，大分雑な説明をすると，NVIDIAのGPUでいうCUDAに相当します．
mpsは，Metal Perfomance Shadersの略称です．
CUDAで使い慣れた人であれば，CUDAが使える環境下で以下のコードを実行したことに相当すると理解すれば良いでしょう．

device = torch.device('cuda:0')

簡単ですね．cuda:0がmpsに置き換わっただけです．

今回用意したコードは，cpuとmpsをpythonの実行時に切り替えられるようにArgumentParserを使ってるので，以下のようにして実行しています．

parser = ArgumentParser()
parser.add_argument('--device', type=str, default='mps',
                    choices=['cpu', 'mps'])
args = parser.parse_args()
device = torch.device(args.device)
main(device)

実行結果

それでは，実行して訓練時間・テスト時間を確かめてみましょう．
今回実行する環境は，2021年モデルのMacBook Pro M1 Maxです．
メインメモリは64GB,コア数は10です．当時最も計算パワーのあるオプションの構成です．
OSは，Monterey 12.4です．PyTorchの公式サイトでは，OSのバージョンは12.3以上必要なようなので，必要に応じてアップデートしてください．

まずはCPUで実行してみます．

$ python pytorch_m1_macbook.py --device cpu
Epoch 1 / 5: time = 427.44[s], loss = 168.22
Epoch 2 / 5: time = 859.18[s], loss = 134.87
Epoch 3 / 5: time = 1284.35[s], loss = 121.36
Epoch 4 / 5: time = 1716.77[s], loss = 112.03
Epoch 5 / 5: time = 2148.09[s], loss = 105.66
Train time on cpu: 2148.09[s] (Train loss = 105.66)
Test time on cpu: 90.19[s](Test loss = 20.52)

訓練時間は2148秒, テスト時間は90秒という結果となりました．
ちなみに筆者としては，これでもM1 MacBookのCPUはかなり高速に実行できていると思います．
最終的な訓練時のロスは105で，テストロスは20.52でした．

続いて，GPU(MPS)で実行してみます．

$ python pytorch_m1_macbook.py --device mps
Epoch 1 / 5: time = 57.19[s], loss = 170.57
Epoch 2 / 5: time = 114.71[s], loss = 138.20
Epoch 3 / 5: time = 172.46[s], loss = 124.46
Epoch 4 / 5: time = 229.30[s], loss = 113.49
Epoch 5 / 5: time = 286.53[s], loss = 104.97
Train time on mps: 286.53[s] (Train loss = 104.97)
Test time on mps: 45.05[s](Test loss = 20.24)

訓練時間は286秒，テスト時間は45秒ですね．CPUでは訓練時間は2148秒だったので，286 / 2148≒13.3%くらいまで実行時間を短縮できています．

訓練のロスの経過も表示していますが，約170で始まり，約105で終わっています．CPUで実行した場合は105だったので，ほとんど変わっていません．CPU版にバグがないのであれば，GPU(MPS)で訓練した場合も問題なさそうに見えます．

この計算速度が実用に耐えるかは使用者次第だと思います．筆者は，ローカルでのデバッグ用くらいには使えそうなので，ギリギリ実用的なレベルまで達しているかと考えています．

コードのダウンロード

本記事のコードは以下のページ(Github)からダウンロードできます．なお，MITライセンスで公開しております．改変・公開等ご自由にお使いください．

https://github.com/HidetoshiKawaguchi/tech-blog-codes/tree/main/20220731_pytorch-m1-macbook-gpu

おわりに

今回は，2022年7月31日の段階で最新のPyTorchを試してみました．CIFAR10やResNetといった，教科書的なCNNの訓練であれば，そこそこ使えるという結果になりました．早くStable版にも反映されてほしいですね．この記事がPyTorch使いの方の役に立ったのであれば幸いです．

2023/03/30追記

Stable版(ver2.0.0)でも動作確認ができました．実行時間にも大きな差はありませんでした．

おまけ

この記事で使用したコードを少しだけ改造して，NVIDIAのGPUで同じように訓練時間とテスト時間を測定してみました．これもあくまで参考値ですが，いかに記載しておきます．

  • Tesla V100 16GB:
    • 訓練時間: 約53秒
    • テスト時間: 約 1.4秒
  • Geforce GTX1080Ti:
    • 訓練時間: 約55秒
    • テスト時間: 約1.3秒
  • RTX A6000:
    • 訓練時間: 約17秒
    • テスト時間: 約0.5秒

やはり，まだまだNVIDIAのGPUの方が速いですね．ただ筆者としては，それでもノートパソコンであるM1 MacBookでGPUを使ってそれなりに早く計算できるようになったのはポジティブに捉えています．

Discussion