torch.square(x)是否比torch.pow(x,2)性能更优？

核心结论 #

绝大多数场景下，torch.square(x)的性能优于或等于torch.pow(x, 2)，不会出现更慢的情况。

底层机制差异 #

• torch.pow的通用性开销：torch.pow是通用幂运算函数，支持任意指数（整数、浮点数甚至张量指数）。为兼容所有场景，内部需要做分支判断、类型适配等额外逻辑，这些都会产生性能损耗。
• torch.square的针对性优化：torch.square是专门为平方操作设计的函数，PyTorch对它做了底层针对性优化。不管是CPU还是GPU后端，都会直接调用最适配平方运算的硬件指令（比如CPU的SIMD平方指令、GPU的专用平方核函数），跳过通用幂运算的冗余逻辑，计算路径更短。

特殊场景例外 #

如果指数是编译期已知的常量2，PyTorch的JIT编译器（如TorchScript）可能会自动把torch.pow(x, 2)优化成和torch.square(x)等价的实现，这时二者性能几乎无差。但在动态图（eager mode）下，这种优化通常不会触发，torch.square依然更快。

基准测试示例 #

如果自己做测试，可参考以下代码：

import torch
import time

# 生成测试张量，可切换CPU/GPU
x = torch.randn(1000, 1000).cuda()

# 测试torch.square
start = time.time()
for _ in range(1000):
    y = torch.square(x)
torch.cuda.synchronize()  # GPU环境需同步计时
print(f"torch.square 耗时: {time.time() - start:.4f}s")

# 测试torch.pow(x,2)
start = time.time()
for _ in range(1000):
    y = torch.pow(x, 2)
torch.cuda.synchronize()
print(f"torch.pow(x,2) 耗时: {time.time() - start:.4f}s")

多数硬件环境下，你会看到torch.square的耗时明显更低。

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