device = torch.device("cuda:1" )
model = model.to(device)
该部分参考自来自达摩院大神的讲解
PyTorch中的数据类型为Tensor,Tensor与Numpy中的ndarray类似,同样可以用于标量,向量,矩阵乃至更高维度上面的计算。PyTorch中的tensor又包括CPU上的数据类型和GPU上的数据类型,一般GPU上的Tensor是CPU上的Tensor加cuda()函数得到。系统默认的torch.Tensor是torch.FloatTensor类型。例如data = torch.Tensor(2,3)是一个2*3的张量,类型为FloatTensor; data.cuda()就将其转换为GPU的张量类型,torch.cuda.FloatTensor类型。
(1) Tensor和numpy.ndarray之间可以相互转换:
import numpy as np
import torch
a_tensor = torch.randint(0, 10, (2,3)) # Tensor声明
print(type(a_tensor))
a_numpy = a_tensor.numpy() # Tensor转Numpy
print(type(a_numpy))
b_numpy = np.array([[1,2,3], [4,5,6]]) # Numpy声明
print(type(b_numpy))
b_tensor = torch.from_numpy(b_numpy) # Numpy转Tensor
print(type(b_tensor))
# 结果打印如下:
<class 'torch.Tensor'>
<class 'numpy.ndarray'>
<class 'numpy.ndarray'>
<class 'torch.Tensor'>
numpy与Tensor最大的区别就是在对GPU的支持上。Tensor只需要调用cuda()函数就可以将其转化为能在GPU上运行的类型。
(2) CPU和GPU的Tensor之间的转换
.cuda()是为了将模型放在GPU上进行训练。