介绍
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]
例子
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，求和后维度会减少，直接做除法会触发广播，可能导致非预期结果：
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：
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 可以指定多个维度，也可以按单维度求和：
# 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 参数可以指定求和结果的类型，避免低精度溢出：
# 低精度张量求和（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)
总结
torch.sum()是张量求和的核心函数，支持全局求和（dim=None）和按指定维度求和（dim=0/1/…）。keepdim=True是关键参数：求和后保留原维度数，避免维度不匹配导致的广播错误（如 softmax 计算）。- 核心应用场景：归一化（softmax 分母）、损失计算（误差求和）、特征聚合（批量维度求和）等。