pytorch api文档：torch.sum()求和-马育民老师

# 介绍

`torch.sum()` 函数，这是张量运算中最基础也最常用的求和函数，既能对整个张量求和，也能按指定维度求和，是归一化、损失计算等场景的核心组件

### 应用场景

计算 softmax

# 语法

对输入张量的元素进行求和运算，支持**全局求和**或**按指定维度求和**，返回求和后的张量（或标量）。

```
torch.sum(input, dim=None, keepdim=False, dtype=None)
```

**参数：**

- `dim`：指定求和的维度（轴），`None` 表示对所有元素求和；
- `keepdim`：是否保留求和维度（保持张量维度不变），默认 `False`；
- `dtype`：指定输出张量的数据类型，默认与输入一致。

### dim 参数的理解（轴的编号）
张量的维度编号从 0 开始，比如 `shape [2,3]` 的张量：

- `dim=0`：垂直方向（行）求和；
- `dim=1`：水平方向（列）求和。

可以用下图直观理解：
```
# 张量 x:
[[1, 2, 3],  → 行0
 [4, 5, 6]]  → 行1
  ↑  ↑  ↑
  列0 列1 列2

dim=0求和 → 列0:1+4, 列1:2+5, 列2:3+6 → [5,7,9]
dim=1求和 → 行0:1+2+3, 行1:4+5+6 → [6,15]
```

# 例子

```python
import torch

# 定义基础张量
x = torch.tensor([[1, 2, 3],
                  [4, 5, 6]])  # shape: [2, 3]

# 示例1：全局求和（dim=None）
sum_all = torch.sum(x)
print("全局求和：", sum_all)  # tensor(21)（1+2+3+4+5+6），标量输出

# 示例2：按维度0求和（行→列，合并行）
sum_dim0 = torch.sum(x, dim=0)
print("按dim0求和：", sum_dim0)  # tensor([5, 7, 9])（1+4, 2+5, 3+6），shape: [3]

# 示例3：按维度1求和（列→行，合并列）
sum_dim1 = torch.sum(x, dim=1)
print("按dim1求和：", sum_dim1)  # tensor([6, 15])（1+2+3, 4+5+6），shape: [2]

# 示例4：keepdim=True（保留求和维度）
sum_dim1_keep = torch.sum(x, dim=1, keepdim=True)
print("按dim1求和（保留维度）：", sum_dim1_keep)
# tensor([[6],
#         [15]])  # shape: [2, 1]（维度数不变，仅求和维度的大小变为1）
```

### keepdim 的作用（维度匹配）

如果不设置 `keepdim=True`，求和后维度会减少，直接做除法会触发广播，可能导致非预期结果：
```python
logits = torch.tensor([[1,2],[3,4]])
exp_logits = torch.exp(logits)

# 错误：sum_dim1 是 [3,7]（shape [2]），exp_logits 是 [2,2]，广播后结果错误
sum_exp_wrong = torch.sum(exp_logits, dim=1)
wrong_softmax = exp_logits / sum_exp_wrong
print("错误softmax：\n", wrong_softmax)

# 正确：keepdim=True 后 sum_exp_right 是 [[3],[7]]（shape [2,1]），维度匹配
sum_exp_right = torch.sum(exp_logits, dim=1, keepdim=True)
right_softmax = exp_logits / sum_exp_right
print("正确softmax：\n", right_softmax)
```

# 应用场景

### 与 softmax 的结合
`torch.sum()` 是 softmax 归一化的关键步骤（计算分母的总和），结合之前的 `torch.exp()` 可以完整实现 softmax：
```python
def manual_softmax(logits, dim=1):
    # 步骤1：计算指数
    exp_logits = torch.exp(logits)
    # 步骤2：按指定维度求和（分母）
    sum_exp = torch.sum(exp_logits, dim=dim, keepdim=True)  # keepdim保证维度匹配
    # 步骤3：归一化
    return exp_logits / sum_exp

# 验证：批量分类logits（2个样本，3个类别）
logits = torch.tensor([[1.0, 2.0, 3.0],
                       [4.0, 5.0, 6.0]])
softmax_out = manual_softmax(logits, dim=1)
print("softmax输出：\n", softmax_out)
# tensor([[0.0900, 0.2447, 0.6652],
#         [0.0900, 0.2447, 0.6652]])
print("每个样本求和验证：", torch.sum(softmax_out, dim=1))  # tensor([1.0000, 1.0000])
```

# 进阶用法
### （1）高维张量求和（批量场景）
对于批量矩阵（3维及以上张量），`dim` 可以指定多个维度，也可以按单维度求和：
```python
# 3维张量：[batch, row, col] = [2, 2, 3]
x_batch = torch.tensor([[[1,2,3], [4,5,6]],
                        [[7,8,9], [10,11,12]]])

# 按batch维度（dim=0）求和
sum_batch = torch.sum(x_batch, dim=0)
print("按batch求和：\n", sum_batch)
# tensor([[ 8, 10, 12],
#         [14, 16, 18]])  # shape: [2, 3]

# 按多个维度求和（dim=(1,2)：合并行和列）
sum_row_col = torch.sum(x_batch, dim=(1,2))
print("按行+列求和：", sum_row_col)  # tensor([21, 57])（第一个样本总和21，第二个57）
```

### （3）数值类型与精度控制
通过 `dtype` 参数可以指定求和结果的类型，避免低精度溢出：
```python
# 低精度张量求和（float16）
x_fp16 = torch.tensor([1e4, 1e4], dtype=torch.float16)
sum_fp16 = torch.sum(x_fp16)
print("float16求和：", sum_fp16)  # tensor(20000., dtype=torch.float16)

# 指定dtype为float32（更高精度）
sum_fp32 = torch.sum(x_fp16, dtype=torch.float32)
print("指定float32求和：", sum_fp32)  # tensor(20000., dtype=torch.float32)
```

# 总结

1. `torch.sum()` 是张量求和的核心函数，支持**全局求和**（dim=None）和**按指定维度求和**（dim=0/1/...）。
2. `keepdim=True` 是关键参数：求和后保留原维度数，避免维度不匹配导致的广播错误（如 softmax 计算）。
3. 核心应用场景：归一化（softmax 分母）、损失计算（误差求和）、特征聚合（批量维度求和）等。

原文出处：http://malaoshi.top/show_1GW2bkUOHCyr.html