【pytorch函数笔记】torch.sum()、torch.unsqueeze()

1.torch.sum

torch.sum(imgs, dim=0) # 按列求和

torch.sum(imgs, dim=1)  # 按行求和

imgs = torch.Tensor([i for i in zip( range(10), range(10))])
print(imgs)
s1=torch.sum(imgs, dim=0) # 按列求和
s2=torch.sum(imgs, dim=1)  # 按行求和
print(s1)
print(s2)

2.torch.unsqueeze

将每个图像张量的维度扩展，即在每个图像张量的最前面添加一个额外维度，以匹配模型的输入形状要求。

 torch.unsqueeze(i, dim=0)  # 按dim=0的维度都变为tensor类型。

imgs = torch.Tensor([i for i in zip( range(10), range(10))])
print(imgs)
print(imgs.shape)
imgs = [torch.unsqueeze(i, dim=0) for i in imgs]
print(imgs)
print(imgs[0].shape)