3D 张量的 PyTorch 双随机归一化

Question

我正在尝试实现 N x N x P 张量的双随机归一化，如 Gong, CVPR 2019 中第 3.2 节所述。这可以在 N x N 情况下使用矩阵运算轻松完成，但我坚持使用 3D 张量情况。我目前拥有的是

def doubly_stochastic_normalise(E):
    """E: n x n x f"""
    E = E / torch.sum(E, dim=1, keepdim=True)  # normalised across rows
    F = E / torch.sum(E, dim=0, keepdim=True)  # normalised across cols
    E = torch.einsum('ijp,kjp->ikp', E, F)
return E

但我想知道是否有没有einsum的方法。

Answer 1

在此设置中，您始终可以退回到使用 torch.matmul（更准确地说是批处理矩阵乘法）。但是，这需要您转置轴。回忆一下两个 3D 输入的矩阵乘法，用 einsum 表示法，它给我们：

bik,bkj->bij

注意 k 维度是如何减少的。要获得此设置，我们需要转置运算符的输入。对于您的情况，我们有：

ijp ? kjp -> ikp
 ↓     ↓      ↑
pij @ pjk -> pik

这转化为：

>>> (E.permute(2,0,1) @ F.permute(2,1,0)).permute(1,2,0)
#       ijp ➝ pij          kjp ➝ pjk      pik ➝ ikp

您可以说您的方法不仅更短而且可读性更高。因此，我会坚持 torch.einsum。 einsum 运算符在这里如此有用的原因是因为您可以动态执行轴换位。

3D 张量的 PyTorch 双随机归一化

PyTorch doubly stochastic normalisation of 3D tensor

pytorch