pytorch 中tensor在CPU和GPU之间转换,以及numpy之间的转换

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'
  1. CPU tensor转GPU tensor:
cpu_imgs.cuda()
  1. GPU tensor 转CPU tensor:
gpu_imgs.cpu()
  1. numpy转为CPU tensor:
torch.from_numpy( imgs )

4.CPU tensor转为numpy数据:

cpu_imgs.numpy()
  1. note:GPU tensor不能直接转为numpy数组,必须先转到CPU tensor。
gpu tensor.cpu().numpy()
  1. 如果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)

你可能感兴趣的:(python,计算机视觉,计算机视觉,深度学习)