device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')#申明用GPU
如果有多个GPU,你可以选择具体的GPU进行使用。
import os
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"]='1'
- CPU tensor转GPU tensor:
cpu_imgs.cuda()
- GPU tensor 转CPU tensor:
gpu_imgs.cpu()
- numpy转为CPU tensor:
torch.from_numpy( imgs )
4.CPU tensor转为numpy数据:
cpu_imgs.numpy()
- note:GPU tensor不能直接转为numpy数组,必须先转到CPU tensor。
gpu tensor.cpu().numpy()
- 如果tensor是标量的话,可以直接使用 item() 函数(只能是标量)将值取出来:
print (loss_output.item())
检验模块
import torch
# 以下代码只有在PyTorch GPU版本上才会执⾏行行
if torch.cuda.is_available():
device = torch.device("cuda") # GPU
x = torch.ones(2,2)
x = x.to(device)
#x = x.cuda()
print(x.device)