我有一個有兩個輸入的函式：熱圖和特征圖。熱圖的形狀為(20, 14, 64, 64)，特征圖的形狀為(20, 64, 64, 64)。哪里20是批量大小，14是關鍵點的數量。熱圖和特征圖都具有空間維度，而特征圖64x64具有64通道（在第二維度上）。
現在我需要將每個熱圖乘以特征圖的每個通道。所以第一個熱圖必須乘以64特征圖的所有通道。第二個是所有頻道，依此類推。
在那之后，我應該有一個形狀張量(20, 14, 64, 64, 64)，我需要在其上應用全域最大池。
現在的問題是我不能創建一個新的張量來做到這一點，因為必須保留熱圖和特征圖的梯度。
我的實際（緩慢且非梯度保持）代碼是：
def get_keypoint_representation(self, heatmaps, features):
heatmaps = heatmaps[0]
pool = torch.nn.MaxPool2d(features.shape[2])
features = features[:, None, :, :, :]
features = features.expand(-1, 14, -1, -1, -1).clone()
for i in range(self.cfg.SINGLE_GPU_BATCH_SIZE):
for j in range(self.cfg.NUM_JOINTS):
for k in range(features.shape[2]):
features[i][j][k] = torch.matmul(heatmaps[i][j], features[i][j][k])
gmp = features.amax(dim=(-1, -2))
return gmp
任務概述：
uj5u.com熱心網友回復：
給定一個形狀的熱圖張量和一個hm形狀(b, k, h, w)的特征張量。fm(b, c, h, w)
您可以使用單個einsum運算子執行此類操作
>>> z = torch.einsum('bkhw,bchw->bkchw', hm, FM)
>>> z.shape
torch.Size([20, 14, 64, 64, 64])
然后使用 對空間維度進行最大池化操作amax：
>>> gmp = z.amax(dim=(-1,-2))
>>> gmp.shape
torch.Size([20, 14, 64])
轉載請註明出處，本文鏈接：https://www.uj5u.com/gongcheng/485285.html