张量重塑查询

Question

我有一个图像张量的形状：-

N,C,H,W = 5,512,13,13

我需要对 H 和 W 维度取平均值，以便输出的形状为 :-

N,C,1,1

我正在尝试执行 for 循环，但是否有更好的方法使用重塑来执行此操作。 .

Answer 1

 import torch

 tz = torch.rand(5, 512, 13, 13)
 tzm = tz.mean(dim=(2,3), keepdim=True)
 tzm.shape

输出

torch.Size([5, 512, 1, 1])