如何使用其他PyTorch函数替换torch.norm（针对dim=1且keepdim=True的高维张量场景）

针对你遇到的CoreML不支持torch.norm内部算子的问题，完全可以通过手动实现范数计算逻辑来替换，而且能完美匹配高维张量下dim=1且keepdim=True的场景。

直接替换代码（对应你的示例张量） #

import torch

x = torch.randn([3, 136, 64, 64])
out1 = torch.norm(x, dim=1, keepdim=True)
# 手动实现的替换逻辑
out2 = (x ** 2).sum(dim=1, keepdim=True) ** 0.5

# 验证结果一致性（浮点运算建议用allclose而非==，避免精度误差导致的误判）
print(torch.allclose(out1, out2))  # 输出: tensor(True)

逻辑解释 #

torch.norm默认计算的是Frobenius范数，本质就是"各元素平方和的平方根"。当指定dim参数时，就是沿着该维度对每个子张量独立计算这个范数：

1. x ** 2：对张量中每个元素取平方
2. .sum(dim=1, keepdim=True)：沿着第1个维度求和，keepdim=True会保持原维度结构（避免求和后维度坍缩，和torch.norm的keepdim行为完全一致）
3. ** 0.5：对平方和取平方根，得到最终的范数结果

这个逻辑完全复刻了torch.norm(x, dim=1, keepdim=True)的行为，而且全程只用到CoreML支持的基础张量运算（平方、求和、幂运算），不会触发未实现的_VF.frobenius_norm算子。

通用化扩展 #

如果需要适配其他维度，只需要修改sum的dim参数即可。比如要计算dim=2的范数：

out2 = (x ** 2).sum(dim=2, keepdim=True) ** 0.5

完全和torch.norm(x, dim=2, keepdim=True)等价。

内容的提问来源于stack exchange，提问作者tommy19970714