PyTorch中的一些常见数据类型转换方法，与list和np.ndarray的转换方法

目录

• • • 0.导包
    • 1.tensor类型转换
    • 2.tensor和list转换
    • 3.tensor和np.ndarray转换
    • 4.图像tensor和图像ndarray转换，tensor -> cv2
    • 5.cv2(ndarray)与PIL之间的转换

0.导包

import torch
import numpy as np
import PIL

1.tensor类型转换

# 设置Torch默认类型
torch.set_default_tensor_type(torch.FloatTensor)
#        CPU                     GPU            #  数据类型
# torch.FloatTensor	    torch.cuda.FloatTensor  # 32位浮点数
# torch.DoubleTensor	torch.cuda.DoubleTensor # 64位浮点数
#                       torch.cuda.HalfTensor   # 16位浮点数

# 类型转换
tensor = torch.rand(1,3,224,224)
tensor = tensor.float()     	# 转为float类型，其他:half() int() double() char() byte() short() long()
tensor = tensor.cuda()      	# 转为GPU对象
tensor = tensor.cpu()       	# 转为CPU对象
# type_as()方法
tensor1 = torch.rand(2,3)
tensor2 = torch.IntTensor(2,3)
tensor3 = tensor1.type_as(tensor2)  # 将tensor1转为tensor2的类型
# print(tensor1,tensor2,tensor3)

2.tensor和list转换

'''tensor和list转换'''
tensor1 = torch.ones([1,5])     # 创建一个1x5的单位张量
list1 = tensor1.tolist()        # tensor -> list
tensor1 = torch.tensor(list1)   # list -> tensor

3.tensor和np.ndarray转换

'''tensor和np.ndarray转换'''
ndarray = tensor.numpy()            # 注意GPU对象要先转为CPU对象
tensor = torch.from_numpy(ndarray)  # 注意一般转为浮点数对象

4.图像tensor和图像ndarray转换，tensor -> cv2

'''图像tensor和图像ndarray转换，tensor -> cv2'''
# tensor的图像是维度信息是[N, C, H, W]
ndarray = np.random.random((3,224,224))
tensor = torch.tensor(ndarray).permute(2,0,1)       #ndarray转tensor图像要注意维度是HWC->CHW
img = tensor.permute(1,2,0).cpu().numpy()           #反之同理

5.cv2(ndarray)与PIL之间的转换

'''PIL.Image和np.ndarray转换，也可是cv2与PIL之间的转换'''
path = '1.jpg'    #这里传入自己的图像地址
ndarray = np.asarray(PIL.Image.open(path))          #Image转ndarray
img = PIL.Image.fromarray(ndarray.astype(np.uint8)) #ndarray转Image