torch.reciprocal介绍

qq_273900232025-01-06 20:22

在 PyTorch 中，torch.reciprocal 是一个用于计算张量中每个元素倒数的函数。它的作用是逐元素执行倒数计算。如果某个元素为零，则结果会是无穷大 (∞)，并产生一个警告。

函数签名

torch.reciprocal(input, *, out=None) → Tensor

参数说明

input: 输入的张量，其数据类型可以是浮点数或整数。
out (可选): 用于存储结果的张量，必须与 input 的形状一致。

返回值

返回一个新的张量，包含输入张量中每个元素的倒数。

示例用法

基本用法

import torch

# 创建张量
x = torch.tensor([1.0, 2.0, 0.5, -1.0])

# 计算倒数
y = torch.reciprocal(x)

print(y)
# 输出: tensor([ 1.0000,  0.5000,  2.0000, -1.0000])

处理包含零的张量

x = torch.tensor([1.0, 0.0, -2.0])

y = torch.reciprocal(x)

print(y)
# 输出: tensor([ 1.0000,    inf, -0.5000])
# 会产生一个警告: Division by zero encountered.

使用 out 参数

x = torch.tensor([2.0, 4.0])
out = torch.empty_like(x)

torch.reciprocal(x, out=out)
print(out)
# 输出: tensor([0.5000, 0.2500])

注意事项

零值处理:
- 当输入张量的某些元素为零时，会产生无穷大 (∞∞)，并触发警告。
- 如果需要处理零值，可以在计算前筛选或替换零值，例如使用 torch.where。
x = torch.tensor([1.0, 0.0, 2.0])
x_safe = torch.where(x == 0, torch.tensor(float('nan')), x)
y = torch.reciprocal(x_safe)
print(y)

输出: tensor([1.0000, nan, 0.5000])

整数类型的输入:

如果输入张量是整数类型，计算结果会自动转换为浮点类型。

x = torch.tensor([2, 4, 8])
y = torch.reciprocal(x)
print(y)

输出: tensor([0.5000, 0.2500, 0.1250])

常见应用场景

规范化操作: 用于计算归一化时的一部分操作。
数学运算: 计算反函数或求倒数的应用。
模型权重更新: 在优化算法中可能需要计算反比值。

torch.reciprocal 是 PyTorch 中一个简单但常用的工具函数，适用于多种场景，尤其是在需要逐元素倒数计算时。