Pytorch张量在GPU上的使用方法

---

一、使用前提

1. 买张显卡（狗头）
2. 安装GPU版的pytorch，参考链接 https://blog.youkuaiyun.com/weixin_43721000/article/details/125290522

---

二、使用方法

1. 创建张量，使用第一块GPU
2. 张量移动到第一块GPU上
3. 创建神经网络，使用第一块GPU
4. 神经网络移动到第一块GPU上

import torch
from torch import nn


def try_gpu(i=0):
    '''
    获取第几块GPU
    :param i:
    :return:
    '''
    if torch.cuda.device_count() >= i+1:
        return torch.device(f'cuda:{i}')
    return torch.device('cpu')


def try_gpus():
    '''
    获取全部GPU，返回数组
    :return:
    '''
    devices = [
        torch.device(f'cuda:{i}') for i in range(torch.cuda.device_count())
    ]
    return devices if devices else torch.device('cpu')


if __name__ == '__main__':

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

    # 创建张量，使用第一块GPU ------------------
    X = torch.ones(2, 3, device=try_gpu())
    print(X)
    # tensor([[1., 1., 1.],
    #         [1., 1., 1.]], device='cuda:0')


    # 张量移动到第一块GPU上 --------------------
    # cpu上声明参数
    X = torch.ones(2, 3)
    # 移动到gpu
    X = X.cuda(device=try_gpu())
    print(X)
    # tensor([[1., 1., 1.],
    #         [1., 1., 1.]], device='cuda:0')
	
	
	# 创建神经网络，使用第一块GPU ----------------------------------
	X = torch.ones(3, device=try_gpu())
	net = nn.Sequential(nn.Linear(3, 1)).to(device=try_gpu())
	print(net(X))
	# tensor([0.9012], device='cuda:0', grad_fn=<AddBackward0>)


    # 神经网络移动到第一块GPU上 ----------------------------------------
    # cpu上声明参数和网络
    X = torch.ones(3)
    net = nn.Sequential(nn.Linear(3, 1))
    # 移动到gpu
    X = X.cuda(device=try_gpu())
    net = net.to(device=try_gpu())
    print(net(X))
    # tensor([1.4811], device='cuda:0', grad_fn=<AddBackward0>)