Pytorch 學習（十）：Pytorch 模型在 CPU 與 GPU 上的遷移

Pytorch 模型在 CPU 與 GPU 上的遷移

本方法總結自《動手學深度學習》（Pytorch版）github項目

• Pytorch計算時必須保證模型和參與當前過程的所有數據都在同一個設備(GPU 或 CPU)上

CPU 與 GPU 的相互轉化

import torch
import torch.nn as nn

x = torch.randn(1, 2)
net = nn.Sequential(
    nn.Linear(1, 1)
)

if torch.cuda.is_available():
    # 利用 to 轉換
    device = torch.device('cuda')
    x = x.to(device)
    print(x.device)
    x = x.to('cpu')
    # 直接轉換
    x = x.cuda(1)
    print(x.device)
    x = x.cpu()
    # 直接創建
    x = torch.randn(1, 1, device='cuda:0')

    net = net.cuda()
    print(net(x))

個人感覺使用 .to(device) 函數最爲方便

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
net = net.to(device)
x = x.to(device)