DL_7——使用GPU

1 配置环境

1.1 安装GPU版PyTorch

1.2 安装cuda

1.3 查看GPU信息

在命令行中输入nvidia-smi可以查看GPU信息

2 计算设备

注意
在使用gpu训练模型时，需确保数据和模型都在同一个设备上，否则可能会报错或者导致训练速度减慢。

import torch
from torch import nn

if __name__ == '__main__':
	torch.device('cpu')  # cpu
	torch.device('cuda:0')  # gpu
	torch.cuda.device_count()  # 统计可使用gpu数量

2.1 将Tensor放入GPU

import torch
from torch import nn

if __name__ == '__main__':
	gpu = torch.device('cuda:0')
	x = torch.tensor([1, 2, 3])
	x.device  # 输出cpu
	x2 = torch.ones(2, 3)
	x2.to(gpu)  # 也可以使用x2.cuda(0) 使用gpu
	x3 = torch.ones(2, 3, device=gpu)  # 也可以在创建的时候指定gpu

2.2 将模型放入GPU

import torch
from torch import nn

if __name__ == '__main__':
	gpu = torch.device('cuda:0')
	x = torch.ones(2, 3, device=gpu)
	net = nn.Sequential(nn.Linear(3, 1))
	net.to(gpu)  # 将模型放入gpu
	print(net(x))
	print(net[0].weight.data.device)  # 查看模型参数是否存储在gpu上