【pytorch基础】读取的图片扩充维度的几种方式

训练时，数据维度一般都是 (batch_size, c, h, w)，而在测试时只输入一张图片(c,h,w)，所以需要扩展维度，扩展维度有多种方式，见如下代码：

• 法1

import cv2
import torch
 
image = cv2.imread(img_path)
#image = torch.tensor(image)
image = torch.from_numpy(image)
print(image.size())
 
img = image.unsqueeze(dim=0)  
print(img.size())
 
img = img.squeeze(dim=0)
print(img.size())
 
# output:
# torch.Size([(h, w, c)])
# torch.Size([1, h, w, c])
# torch.Size([h, w, c])

• 法2

import cv2
import numpy as np
 
image = cv2.imread(img_path)
print(image.shape)
img1 = image[np.newaxis, :, :, :]
img2 = image[..., np.newaxis]
print(img1.shape)
print(img2.shape)
 
# output:
# (h, w, c)
# (1, h, w, c)
#（h, w, c, 1）

• 法3

import cv2
import torch
 
image = cv2.imread(img_path)
image = torch.tensor(image)
print(image.size())
 
img = image.view(1, *image.size())
print(img.size())
 
# output:
# torch.Size([h, w, c])
# torch.Size([1, h, w, c])