pytroch中指定GPU或者CPU寫法
阿新 • • 發佈:2018-12-17
之前習慣用.cpu(),.cuda()來指定.
現在不要顯示的指定是gpu, cpu之類的. 利用.to()來執行
# at beginning of the script device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") ... # then whenever you get a new Tensor or Module # this won't copy if they are already on the desired device input = data.to(device) model = MyModule(...).to(device)
用.to()和.cuda()程式碼對比
.cuda()寫法
model = MyRNN() if use_cuda: model = model.cuda() # train total_loss = 0 for input, target in train_loader: input, target = Variable(input), Variable(target) hidden = Variable(torch.zeros(*h_shape)) # init hidden if use_cuda: input, target, hidden = input.cuda(), target.cuda(), hidden.cuda() ... # get loss and optimize total_loss += loss.data[0] # evaluate for input, target in test_loader: input = Variable(input, volatile=True) if use_cuda: ... ...
現在.to()
# torch.device object used throughout this script device = torch.device("cuda" if use_cuda else "cpu") model = MyRNN().to(device) # train total_loss = 0 for input, target in train_loader: input, target = input.to(device), target.to(device) hidden = input.new_zeros(*h_shape) # has the same device & dtype as `input` ... # get loss and optimize total_loss += loss.item() # get Python number from 1-element Tensor # evaluate with torch.no_grad(): # operations inside don't track history for input, target in test_loader: ...