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)