崔和美 发表于 2025-6-25 16:23:25

PyTorch中实现开立方

技术背景

在PyTorch中,没有直接实现cbrt这一算子。这个算子是用于计算一个数的开立方,例如,最简单的-8开立方就是-2。但这里有个问题是,在PyTorch中,因为没有cbrt算子,如果直接用幂次计算去操作数字,就有可能出现问题。
代码示例

首先看一下numpy做开立方的代码示例:
In : import numpy as np

In : a = np.array(-8, np.float32)

In : a**(1/3)
<ipython-input-3-f6e83d4e282e>:1: RuntimeWarning: invalid value encountered in power
a**(1/3)
Out: np.float32(nan)

In : np.cbrt(a)
Out: np.float32(-2.0)在这个示例中,如果直接开立方,结果会是一个nan,很明显不是我们想要的一个结果。而cbrt是一个单独实现的开立方算子,可以支持负数的输入,计算结果也是正确的。在PyTorch的场景下,只能用幂次运算:
In : import torch as tc

In : a=tc.tensor(-8,dtype=tc.float32)

In : a**(1/3)
Out: tensor(nan)这样得到的结果是错误的。因此需要我们自己实现一个cbrt函数:
In : import torch as tc

In : cbrt=lambda x: tc.sign(x)*tc.abs(x)**(1/3)

In : a=tc.tensor(-8,dtype=tc.float32)

In : cbrt(a)
Out: tensor(-2.)其实逻辑也比较简单,就是先把符号提取出来,然后再转化为正数正常计算就好了。
总结概要

本文介绍了在PyTorch中直接使用幂次函数计算有可能导致的计算结果异常的问题。由于PyTorch中并未像Numpy和MindSpore一样直接支持cbrt开立方函数,因此这里也提供了一个在PyTorch中计算开立方的函数。
版权声明

本文首发链接为:https://www.cnblogs.com/dechinphy/p/cbrt.html
作者ID:DechinPhy
更多原著文章:https://www.cnblogs.com/dechinphy/
请博主喝咖啡:https://www.cnblogs.com/dechinphy/gallery/image/379634.html

来源:豆瓜网用户自行投稿发布,如果侵权,请联系站长删除
页: [1]
查看完整版本: PyTorch中实现开立方