张量造好之后，绝大多数运算都可以归到两类。一类是把一堆数字压成一个或者压成一列，比如求平均、求和。另一类是对每个元素独立做同一个数学变换，比如平方、开方、取对数。前一类是神经网络里算损失、看指标、做特征压缩的主力，后一类是激活函数、概率变换的地基。两类混着用的时候，输出形状会怎么变，经常让人发懵。

这一节我们把 PyTorch 最常用的六个函数 mean、sum、pow、sqrt、exp、log 一次讲清，重点放在它们到底压掉了哪个维度、输出形状变成什么样。每个函数都拿一张固定张量把数字走一遍，再用代码验证。读之前只要知道张量是多维数组，会 import torch 就行。

1. 均值与求和

mean 算均值，sum 算总和，它们属于归约函数，意思是把某一个维度上的多个数字合并成一个，输出形状因此比输入少一维。我们拿一张固定张量走一遍：data = [[1,2,3],[4,5,6]]，形状是 2 行 3 列。先不指定 dim，data.mean() 把所有元素压成一个标量，六个数加起来是 21，除以 6 得 3.5。data.sum() 同理得到 21。

真正要弄懂的是 dim 参数。data.mean(dim=0) 的意思是沿着第 0 个维度（也就是行的方向）做归约，把两行压没，剩下每行一列，得到三个列均值：(1+4)/2 = 2.5，(2+5)/2 = 3.5，(3+6)/2 = 4.5，结果是 [2.5, 3.5, 4.5]。data.mean(dim=1) 是沿着第 1 个维度（列的方向）归约，把三列压没，得到两个行均值：(1+2+3)/3 = 2，(4+5+6)/3 = 5，结果是 [2, 5]。sum 的方向完全一致，把上面过程里的平均换成求和就行。整个对照表是这样的：

记 dim 方向有个口诀：dim 等于几，输出形状里那一维就消失。输入是 (2,3)，dim=0 之后第 0 维的 2 没了，剩 (3,)，dim=1 之后第 1 维的 3 没了，剩 (2,)。工程上还常用 keepdim=True，它把被压掉的维度保留成 1，输出形状变成 (1,3) 或 (2,1)，这样后面和别的张量做广播运算时形状能直接对上，省得手动补维度。

还有一个使用限制要知道：mean 只能吃浮点型张量，整型张量直接调 mean 会报错，所以示例代码里特意用 dtype=torch.float64 造数据。sum 没有这个限制，整型张量照算。

2. 幂与开方

和归约函数相对的是逐元素函数，它们对每个元素独立做变换，输出形状和输入一模一样，一个维度都不消失。pow 是其中最直观的：data.pow(2) 把每个元素平方，输入 [[1,2,3],[4,5,6]]，输出 [[1,4,9],[16,25,36]]。pow 的参数可以是任意数，pow(3) 就是立方，pow(0.5) 等价于开平方。

sqrt 是专门的开平方函数，对每个元素取算术平方根。照着输入逐位算：sqrt(1)=1，sqrt(2)≈1.4142，sqrt(3)≈1.7321，sqrt(4)=2，sqrt(5)≈2.2361，sqrt(6)≈2.4495，结果是 [[1, 1.4142, 1.7321],[2, 2.2361, 2.4495]]。形状完全没变，这就是逐元素函数和归约函数最本质的区别：没有 dim 参数，也不需要指定压哪个维度。

3. 指数与对数

exp 和 log 也是逐元素函数，形状同样不变。exp(x) 算 e 的 x 次方，e 是自然常数，约等于 2.71828。照着输入逐位算：e^1≈2.7183，e^2≈7.3891，e^3≈20.0855，e^4≈54.5982，e^5≈148.4132，e^6≈403.4288，结果是 [[2.7183, 7.3891, 20.0855],[54.5982, 148.4132, 403.4288]]。指数函数把小的正数压成略大于 1 的数，把大的正数快速放大，神经网络里它常用来把任意实数映射成正数概率，后面再接归一化。

log 是 exp 的反函数，data.log() 默认以 e 为底，照着输入逐位算：log(1)=0，log(2)≈0.6931，log(3)≈1.0986，log(4)≈1.3863，log(5)≈1.6094，log(6)≈1.7918。PyTorch 还提供 log2 和 log10，分别以 2 和 10 为底，结果对照如下：

用对数有一条铁律：输入必须是正数。log(0) 得到负无穷，log(负数) 得到 nan。神经网络里算交叉熵损失时，概率经过 softmax 理论上严格大于 0，但数值上可能逼近 0，工程上常给 log 前加一个极小的常数防止溢出。指数和对数在损失函数里成对出现，log(softmax(x)) 这种写法直接被封装成 log_softmax，数值稳定性更好。

4. 代码验证

上面的数字我们都用代码跑一遍确认。先拿固定张量验证手工算的结果，避免凭空推错：

import torch

data = torch.tensor([[1., 2., 3.], [4., 5., 6.]])

print('整体均值:', data.mean().item())           # 3.5
print('按列均值:', data.mean(dim=0).tolist())    # [2.5, 3.5, 4.5]
print('按行均值:', data.mean(dim=1).tolist())    # [2.0, 5.0]
print('整体求和:', data.sum().item())            # 21
print('按列求和:', data.sum(dim=0).tolist())     # [5, 7, 9]
print('按行求和:', data.sum(dim=1).tolist())      # [6, 15]
print('平方:', data.pow(2).tolist())
print('开方:', data.sqrt().round(4).tolist())
print('指数:', data.exp().round(4).tolist())
print('对数:', data.log().round(4).tolist())

下面这段是原始示例，用随机整数造一张 2 乘 3 的 float64 张量，把六个函数连同 dim 参数的各种用法依次跑出来。随机值每次都不同，但函数行为和上面的对照是一致的：

import torch


def test():

    data = torch.randint(0, 10, [2, 3], dtype=torch.float64)
    print(data)
    print('-' * 50)

    # 1. 计算均值
    # 注意: tensor 必须为 Float 或者 Double 类型
    print(data.mean())
    print(data.mean(dim=0))  # 按列计算均值
    print(data.mean(dim=1))  # 按行计算均值
    print('-' * 50)

    # 2. 计算总和
    print(data.sum())
    print(data.sum(dim=0))
    print(data.sum(dim=1))
    print('-' * 50)

    # 3. 计算平方
    print(data.pow(2))
    print('-' * 50)

    # 4. 计算平方根
    print(data.sqrt())
    print('-' * 50)

    # 5. 指数计算, e^n 次方
    print(data.exp())
    print('-' * 50)

    # 6. 对数计算
    print(data.log())  # 以 e 为底
    print(data.log2())
    print(data.log10())


if __name__ == '__main__':
    test()

5. 常见误区

第一个误区，拿整型张量直接调 mean。mean 只接受浮点型，整型张量会报错，要么造数据时指定 dtype=torch.float64，要么先 .float() 转一下。sum 没有这个限制，两者容易混。

第二个误区，把 dim 方向想反。记住口诀：dim 等于几，输出形状里那一维就消失。(2,3) 配 dim=0 剩 (3,)，配 dim=1 剩 (2,)，多在草稿纸上画两回就熟了。不确定时打一行 shape 看一眼，比硬记可靠。

第三个误区，给 log 喂 0 或负数。log(0) 是负无穷，log(负数) 是 nan，损失函数里出现 nan 基本都是这条原因，喂数据前先确认概率严格大于 0。

第四个误区，以为 pow、sqrt、exp、log 会改变形状。它们是逐元素函数，输出形状和输入完全一致，不需要也不能传 dim 参数。

小结一下。六个函数分成两拨。mean 和 sum 是归约函数，按 dim 指定的方向把对应维度压没，输出形状少一维，需要时用 keepdim 保留成 1 方便广播，mean 还要求浮点型输入。pow、sqrt、exp、log 是逐元素函数，每个元素独立变换，形状不变，log 要求输入为正。把固定张量 [[1,2,3],[4,5,6]] 在草稿纸上按 dim 推一遍，再跑一遍代码对照输出，归约和逐元素这两类函数就分清了，后面学 softmax、交叉熵也不会发懵。