如何用其他pytorch函数替换torch.norm？

Question

我想用其他 Pytorch 函数替换 torch.norm 函数。

我能够在 x 不是矩阵的情况下替换 torch.norm，如以下代码所示。

import torch
x = torch.randn(9)
out1 = torch.norm(x)
out2 = sum(abs(x)**2)**(1./2)
out1 == out2
>> tensor(True)

但是我不知道x是矩阵的时候怎么替换。特别是，我想在 dim=1 and keepdim=True.

的情况下替换它

x = torch.randn([3, 136, 64, 64])
out1 = torch.norm(x, dim=1, keepdim=True)
out2 = ???
out1 == out2

背景：

我正在将 Pytorch 模型转换为 CoreML，但 torch.norm 函数中定义的 _VF.frobenius_norm 运算符未使用 CoreMLTools 实现。（torch.norm里面的实现可以参考here。）

少数人遇到了这个问题，但 CoreMLTools 仍然不受支持（您可以查看此 issue）。所以我想在没有 torch.norm.

中使用的运算符的情况下替换它

我尝试了 torch.linalg.norm() 和 numpy.linalg.norm，但它们不受支持。

我创建了一个简单的合作实验室笔记本来复制这个。请使用以下 colab 对其进行测试。 https://colab.research.google.com/drive/11o6rTxHzEgZ_Rc7nFZHd3TvPugybB88h?usp=sharing

Answer 1

您可以尝试以下方法：

import torch
x = torch.randn([3, 136, 64, 64])
out1 = torch.norm(x, dim=1, keepdim=True)
out2 = torch.square(x).sum(dim=1, keepdim=True).sqrt()

请注意，out1 == out2 不会完全给出所有 True，因为精度有小误差。您可以检查错误的顺序是否为 1e-7 for float32.

此处，范数是使用其数学定义直接计算的。您可以从 Wolfram MathWorld 中查看 this reference 了解更多详细信息。

How to replace torch.norm with other pytorch function?