torch.cat()

import torch

b, c, w, h = 2, 3, 4, 4

x = torch.zeros(b, c, w, h)
print(x.shape) # torch.Size([2, 3, 4, 4])
mean=[0.485, 0.456, 0.406] #rgb
std =[0.229, 0.224, 0.225]

get_channel = x[:,[0]]
print(get_channel.shape) # torch.Size([2, 1, 4, 4])


y = torch.cat([
    (x[:,[0]]-mean[0])/std[0],
    (x[:,[1]]-mean[1])/std[1],
    (x[:,[2]]-mean[2])/std[2],
], 1) # cat第几维度，第几维度就会增加
print(y.shape) # torch.Size([2, 3, 4, 4])