tf.reduce_sum 是 TensorFlow 库中的一个函数，用于计算张量（Tensor）沿指定维度的元素总和。这个函数非常有用，尤其是在深度学习和机器学习任务中，当需要对数据进行聚合操作时。

函数签名

tf.reduce_sum(
    input_tensor,
    axis=None,
    keepdims=None,
    name=None
)

参数说明

  • input_tensor: 必填参数，类型为 Tensor，表示要对其元素求和的输入张量。
  • axis: 可选参数，类型为整数或整数列表，表示沿着哪个维度进行求和。如果 axis 为 None（默认值），则会计算整个张量中所有元素的总和。
  • keepdims: 可选参数，默认为 False。如果设置为 True，则保持减少的维度，作为长度为 1 的维度。
  • name: 可选参数，用于指定该操作的名称。

返回值

返回一个新的 Tensor，其形状取决于 axis 和 keepdims 参数。如果 axis 为 None，返回的是一个标量 Tensor；如果指定了 axis，返回的 Tensor 将在指定的维度上进行缩减。

示例

假设我们有一个二维张量 a，我们来看几个使用 tf.reduce_sum 的例子：

1. 计算整个张量的总和

import tensorflow as tf

a = tf.constant([[1, 2], [3, 4]])
total_sum = tf.reduce_sum(a)

print(total_sum)  # 输出: 10

2. 沿着特定维度求和

import tensorflow as tf

a = tf.constant([[1, 2], [3, 4]])

# 沿着第 0 维度（行）求和
sum_along_rows = tf.reduce_sum(a, axis=0)
print(sum_along_rows)  # 输出: [4 6]

# 沿着第 1 维度（列）求和
sum_along_columns = tf.reduce_sum(a, axis=1)
print(sum_along_columns)  # 输出: [3 7]

3. 保持减少的维度

import tensorflow as tf

a = tf.constant([[1, 2], [3, 4]])

# 沿着第 1 维度（列）求和，并保持维度
sum_keepdims = tf.reduce_sum(a, axis=1, keepdims=True)
print(sum_keepdims)  # 输出: [[3]
                     #        [7]]

这些示例展示了 tf.reduce_sum 的基本用法。根据实际需求，你可以灵活选择 axis 和 keepdims 参数来实现不同的聚合效果。