PyTorch可以指定用来存储和计算的设备,如使用内存的CPU或者使用显存的GPU。在默认情况下,PyTorch会将数据创建在内存,然后利用CPU来计算。
PyTorch要求计算的所有输入数据都在内存或同一块显卡的显存上。
检测是否可以使用GPU,使用一个全局变量use_gpu,便于后面操作使用
use_gpu = torch.cuda.is_available()
可以使用GPU,use_gpu的值为True,否则为False。当可以使用GPU,我们不想使用,可以直接赋值use_gpu = False
我们在进行转换时,需要把数据,网络,与损失函数转换到GPU上
1.构建网络时,把网络,与损失函数转换到GPU上
model = get_model() loss_f = t.nn.CrossEntropyLoss() if(use_gpu): model = model.cuda() loss_f = loss_f.cuda()
2.训练网络时,把数据转换到GPU上
if (use_gpu): x,y = x.cuda(),y.cuda()
3.取出数据时,需要从GPU准换到CPU上进行操作
if(use_gpu): loss = loss.cpu() acc = acc.cpu()
法一:
device=torch.device("cuda:0" if torch.cuda.is_available() else "cpu") model=model.to(device) x=x.to(device) y=y.to(device)
法二:
model=model.cuda() x=x.cuda() y=y.cuda()
list 列表 不能使用 .to(device)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = LSTM().to(device)
model.hidden_cell = (torch.zeros(1, 1, model.hidden_layer_size).to(device),
torch.zeros(1, 1, model.hidden_layer_size).to(device))
y_pred = model(seq.to(device))
single_loss = loss_function(y_pred, labels.to(device))
REF
https://discuss.pytorch.org/t/trying-to-train-lstm-on-gpu/47674
https://blog.csdn.net/yzy__zju/article/details/85014576
https://www.cnblogs.com/sdu20112013/p/12145857.html