张量重塑查询

Question

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

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

我需要在 H 和 W 维度上取一个平均值，以便输出具有形状：-

N,C,1,1

我正在尝试做 for 循环，但有没有更好的方法可以使用 reshape 来做到这一点。 .

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