[pytorch basics] several ways to expand the dimensions of the read pictures

During training, the data dimensions are generally (batch_size, c, h, w), but only one image (c, h, w) is input during testing, so the dimension needs to be expanded. There are many ways to expand the dimension, see the following code :

• Act 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])

• Act 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）

• Act 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])