Apple Silicon在pytorch的GPU測試

---

介紹

大約在暑假的時候，Pytorch就宣布Apple Silicon的GPU將可以在做深度學習中使用，而這幾天我就稍微用了我的電腦測試一下M1 GPU對於訓練一個神經網路的效果與時間

建立conda環境

• 建立與進入環境

  conda create --name pytorch-with-m1 python=3.8
  conda activate pytorch-with-m1

下載pytorch preview version

• 到pytorch官方網站找要指令輸入複製(在Preview中)

  pip3 install --pre torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/nightly/cpu

確認是否能使用Apple Silicon的GPU

• 測試是否可以執行(若安裝成功會輸出True，tensor的device會是mps:0)

  torch.backends.mps.is_available()
  torch.zeros(1).to('mps')

測試階段

• 說明
  • 我將所有random seed皆固定住，並使用辨識mnist的資料作為本次測試
  • 只需將device設定為’mps’ 或 device設定為’cpu’ 就可以分別測試使用GPU或使用CPU所需的時間
• 參數設定
  • epoch: 10
  • optimizer: adam
  • learning rate: 0.001
  • batch size: 128
  • critirion: cross entropy

    import torch
    import torch.nn as nn
    from torch.utils.data import DataLoader
    from torchvision import transforms
    from torchvision import datasets
    from qqdm import qqdm
    import time
    import os, random, numpy as np
    
    start_time = time.time()
    
    # fix random seed
    print('seeded')
    seed = 0
    os.environ["PYTHONHASHSEED"] = str(seed)
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    
    device = "mps"
    train_data = datasets.MNIST(root='./data', train=True, transform=transforms.ToTensor(), download=True)
    train_loader = DataLoader(train_data, batch_size=128)
    
    class Model(nn.Module):
        def __init__(self):
            super(Model, self).__init__()
            self.net = nn.Sequential(
                nn.Conv2d(1, 8, 3, 1, 1),
                nn.ReLU(),
                nn.MaxPool2d(2),
                nn.Conv2d(8, 16, 3, 1, 1),
                nn.ReLU(),
                nn.MaxPool2d(2),
                nn.Flatten(),
                nn.Linear(16 * 7 * 7, 10),
                nn.Softmax(dim=1),
            )
        def forward(self, x):
            return self.net(x)
    
    Net = Model().to(device)
    optimizer = torch.optim.Adam(Net.parameters(), lr=0.001)
    critirion = nn.CrossEntropyLoss()
    
    for epoch in range(10):
        train_loss = 0
        progress_bar = qqdm(train_loader)
        for (x, y) in progress_bar:
            x = x.to(device)
            y = y.to(device)
            y_pred = Net(x)
            loss = critirion(y_pred, y)
            
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            
            progress_bar.set_infos({
                'epoch': epoch,
                'loss': f'{loss:.4f}',
            })
            
    n_correct = 0
    n_sample = 0
    for i, (x, y) in enumerate(train_loader):
        x = x.to(device)
        y = y.to(device)
        out = Net(x)
        _, pred = torch.max(out, 1)
        n_correct += (pred == y).sum().item()
        n_sample += y.shape[0]
    acc = n_correct / n_sample
    
    print(f'\nAccuracy: {acc:.4f}', f'total time when using {device}: {time.time() - start_time:.4f} seconds')

結果

• CPU: Accuracy: 0.9886 total time when using cpu: 131.9462 seconds
• GPU: Accuracy: 0.9889 total time when using mps: 100.2587 seconds

分析

• 我們可以看到，對於準確率來說，兩者相異不大，會有差別的原因可能是因為CPU與GPU對於浮點運算上可能有些誤差導致而成
• 就時間上來說，GPU比CPU所花的時間少了31.7秒左右，也就是是說使用GPU快了約31.6%左右，看起來並非特別明顯，可能要等之後正式版釋出後再來看pytorch有沒有再將其優化了
• 另外，我也做了純Linear的版本，當batch size不夠大的時候(我自己是測到大小為1024)，我發現使用CPU所花的時間會比GPU的少了許多，也就是說除非能夠將batch size設得足夠大，否則在只使用Linear時使用GPU反而拖慢了訓練時間

結論

將Apple Silicon拿來做深度學習終於不是件遙不可及的事了，且得益於Apple的GPU與CPU的RAM是共用的，因此或許可以減少在訓練時Out Of Memory發生的情況。雖然說目前我認為要調整內容應該還很多，但能夠有所進步已經是一件令人非常開心的事了。

最近pytorch也宣佈將整個框架交給由AMD， AWS， Google， Meta， Microsoft，Nvidia所組成的獨立基金會來管理，或許未來AMD顯卡也適合拿來跑深度學習的日子也近了～

---