PyTorch 里所有数据都以张量（Tensor）的形式存在，它本质就是一个多维数组。模型的前向传播、反向传播，拆开看都是一连串张量运算。这些运算本身不复杂，但用错地方会悄悄改坏原数据，或者和矩阵乘法混为一谈，结果算错了还浑然不觉。
这一节我们把张量最常用的四类计算过一遍：加减乘除这类逐元素运算、阿达玛积、真正的矩阵乘法，以及把计算放到 CPU 还是 GPU 上的设备问题。每一类都拿具体张量把数字走一遍，再配代码验证。读之前只要知道张量是多维数组，会 import torch 就行。
1. 基本运算
我们先从最简单的逐元素加减乘除说起，它是所有复杂运算的地基。PyTorch 提供 add、sub、mul、div、neg 这一组函数，分别对应加、减、乘、除、取负。它们的共同点是对张量里的每个元素独立做同样的事。拿一个具体张量走一遍：data = [[1,2],[3,4]]，加上标量 10 之后，每个元素都加 10，得到 [[11,12],[13,14]]。同理，data.sub(100) 得到 [[-99,-98],[-97,-96]]，data.mul(100) 得到 [[100,200],[300,400]]，data.div(100) 得到 [[0.01,0.02],[0.03,0.04]]，data.neg() 得到 [[-1,-2],[-3,-4]]。
这一组函数还有一套带下划线的版本，add_、sub_、mul_、div_、neg_。下划线是 PyTorch 的约定，表示这个操作在原张量身上直接改，不再返回新张量。比如 data.add_(10) 执行完，data 本身就变成了 [[11,12],[13,14]]，外面不会再产生新对象。不写下划线时，函数返回一个新张量，原数据原封不动。两者的区别我们列在下面：
| 写法 | 对原数据 | 返回值 | 典型用途 |
|---|---|---|---|
| data.add(10) | 不改动 | 新张量 | 需要保留原数据时 |
| data.add_(10) | 就地改写 | 原张量本身 | 省内存，链式累加时 |
这里要提醒一句：就地改写虽然省内存，但在需要梯度的张量上乱用会出问题。一旦对 require_grad=True 的张量做了 in-place 修改，反向传播时要用的中间值已经被冲掉，PyTorch 会直接报错。所以训练循环里，除非明确知道自己在做累加器，否则优先用不带下划线的版本。下面的代码把这两类写法都跑一遍：
import numpy as np
import torch
def test():
data = torch.randint(0, 10, [2, 3])
print(data)
print('-' * 50)
# 1. 不修改原数据
new_data = data.add(10) # 等价 new_data = data + 10
print(new_data)
print('-' * 50)
# 2. 直接修改原数据
# 注意: 带下划线的函数为修改原数据本身
data.add_(10) # 等价 data += 10
print(data)
# 3. 其他函数
print(data.sub(100))
print(data.mul(100))
print(data.div(100))
print(data.neg())
if __name__ == '__main__':
test()
2. 阿达玛积
逐元素乘法我们已经见过了。两个同形状张量对应位置相乘这件事，有个专门的名字，叫阿达玛积。它要求两个张量形状完全一致，结果形状也和它们相同。拿两个 2 乘 2 张量走一遍：data1 = [[1,2],[3,4]]，data2 = [[5,6],[7,8]]，逐位相乘就是 (1×5, 2×6, 3×7, 4×8) = (5, 12, 21, 32)，结果是 [[5,12],[21,32]]。
写法上有两种：torch.mul(data1, data2) 和直接写 data1 * data2，两者完全等价。要特别注意它和下一节矩阵乘法的区别：阿达玛积是对应位置各算各的，矩阵乘法是整行乘整列再求和，两者形状规则、结果意义都不一样，新手最常把这两件事搞混。我们用代码验证一下：
import numpy as np
import torch
def test():
data1 = torch.tensor([[1, 2], [3, 4]])
data2 = torch.tensor([[5, 6], [7, 8]])
# 第一种方式
data = torch.mul(data1, data2)
print(data)
print('-' * 50)
# 第二种方式
data = data1 * data2
print(data)
print('-' * 50)
if __name__ == '__main__':
test()
3. 矩阵乘法
阿达玛积是对应位置相乘，可神经网络里真正用得多的反而是另一种，矩阵乘法。它的规则是：第一个矩阵形状 (n, m)，第二个 (m, p)，中间的 m 必须相等，结果形状是 (n, p)。我们拿一组具体数字走一遍。data1 是 3 乘 2 的 [[1,2],[3,4],[5,6]]，data2 是 2 乘 2 的 [[5,6],[7,8]]，中间维度都是 2，可以相乘。结果第一行第一列是 1×5 + 2×7 = 19，第一行第二列是 1×6 + 2×8 = 22，三行算完得到 [[19,22],[43,50],[67,78]]。
PyTorch 里做矩阵乘法有四个接口，规则各有侧重，我们先列出来对照：
| 写法 | 适用维度 | 形状要求 | 结果形状 |
|---|---|---|---|
| data1 @ data2 | 二维及以上 | 按 matmul 规则 | 按广播规则 |
| torch.mm(a, b) | 只接受二维 | (n,m) 乘 (m,p) | (n,p) |
| torch.bmm(a, b) | 只接受三维 | 批量数相同，(b,n,m) 乘 (b,m,p) | (b,n,p) |
| torch.matmul(a, b) | 不限维度 | 最后两维符合矩阵乘法，其余按广播对齐 | 按广播规则 |
@ 本质是 matmul 的运算符简写。mm 最严格，只吃二维，传三维直接报错。bmm 是批量版本，第一维是 batch size，要求两个张量的 batch 数一致，后面两维满足矩阵乘法，比如 (3,4,5) 乘 (3,5,8) 得到 (3,4,8)。matmul 最灵活，维度不同时会按广播规则自动对齐，几个典型形状跑出来是这样的：
| 输入 a 形状 | 输入 b 形状 | matmul 结果形状 |
|---|---|---|
| (3, 4, 5) | (5, 4) | (3, 4, 4) |
| (5, 4) | (3, 4, 5) | (3, 5, 5) |
| (3, 5, 4) | (3, 4, 5) | (3, 5, 5) |
第二行那个例子最能体现广播。a 是二维 (5,4)，b 是三维 (3,4,5)，a 前面被自动补成 (1,5,4)，和 b 的 batch 维 3 对齐扩成 3 份，每份算一个 (5,4) 乘 (4,5)，最后叠成 (3,5,5)。写代码时选接口的原则很简单：纯二维矩阵用 mm，确定是规整的批量三维用 bmm，维度不固定或者要写通用代码就用 matmul。下面的代码把这几个接口都跑一遍：
import numpy as np
import torch
# 1. 点积运算
def test01():
data1 = torch.tensor([[1, 2], [3, 4], [5, 6]])
data2 = torch.tensor([[5, 6], [7, 8]])
# 第一种方式
data = data1 @ data2
print(data)
print('-' * 50)
# 第二种方式
data = torch.mm(data1, data2)
print(data)
print('-' * 50)
# 第三种方式
data = torch.matmul(data1, data2)
print(data)
print('-' * 50)
# 2. torch.mm 和 torch.matmul 的区别
def test02():
# matmul 可以两个维度可以不同
# 第一个张量: (3, 4, 5)
# 第二个张量: (6, 4)
# torch.mm 不可以相乘，而 matmul 则可以相乘
print(torch.matmul(torch.randn(3, 4, 5), torch.randn(5, 4)).shape)
print(torch.matmul(torch.randn(5, 4), torch.randn(3, 4, 5)).shape)
print(torch.matmul(torch.randn(3, 5, 4), torch.randn(3, 4, 5)).shape)
# 3. torch.bmm 函数的用法
def test03():
# 批量点积运算
# 第一个维度为 batch_size
# 矩阵的二三维要满足矩阵乘法规则
data1 = torch.randn(3, 4, 5)
data2 = torch.randn(3, 5, 8)
data = torch.bmm(data1, data2)
print(data.shape)
if __name__ == '__main__':
test02()
4. 运算设备
运算本身讲完了，最后一个问题是这些计算跑在哪。PyTorch 默认把张量创建在 CPU 管理的内存里，计算也在 CPU 上做。但矩阵乘法这种大规模并行计算，放到 GPU 上能快几十倍，这也是模型训练离不开显卡的原因。把张量挪到 GPU 上，有三种做法。
第一种是用 cuda 方法把已经在 CPU 上的张量搬过去，data.cuda() 之后设备就变成 cuda:0，再用 data.cpu() 能搬回来。第二种是创建时直接指定设备，torch.tensor([10,20,30], device=’cuda:0′)，张量一生下来就在 GPU 上，省掉一次搬运。第三种是用 to 方法，data.to(‘cuda:0’)，它和 cuda 的区别是写法统一，换设备只改一个字符串，代码可移植性更好。
这里有个必须知道的硬规则：存放在不同设备上的张量不能直接运算。一个在 cuda:0，一个在 CPU，相加会直接抛 RuntimeError，提示所有张量要放在同一个设备上。所以写代码时一般先用 torch.cuda.is_available() 判断有没有显卡，有就用 cuda，没有就退回 cpu，让同一个 device 变量贯穿全程，就不会出现设备对不上的问题。下面的代码把三种搬法和设备不匹配的报错都演示了：
import torch
# 1. 使用 cuda 方法
def test01():
data = torch.tensor([10, 20 ,30])
print('存储设备:', data.device)
# 如果安装的不是 gpu 版本的 PyTorch
# 或电脑本身没有 NVIDIA 卡的计算环境
# 下面代码可能会报错
data = data.cuda()
print('存储设备:', data.device)
# 使用 cpu 函数将张量移动到 cpu 上
data = data.cpu()
print('存储设备:', data.device)
# 输出结果:
# 存储设备: cpu
# 存储设备: cuda:0
# 存储设备: cpu
# 2. 直接将张量创建在 GPU 上
def test02():
data = torch.tensor([10, 20, 30], device='cuda:0')
print('存储设备:', data.device)
# 使用 cpu 函数将张量移动到 cpu 上
data = data.cpu()
print('存储设备:', data.device)
# 输出结果:
# 存储设备: cuda:0
# 存储设备: cpu
# 3. 使用 to 方法
def test03():
data = torch.tensor([10, 20, 30])
print('存储设备:', data.device)
data = data.to('cuda:0')
print('存储设备:', data.device)
# 输出结果:
# 存储设备: cpu
# 存储设备: cuda:0
# 4. 存储在不同设备的张量不能运算
def test04():
# 判断 GPU 设备是否存在
device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
data1 = torch.tensor([10, 20, 30], device=device)
data2 = torch.tensor([10, 20, 30])
print(data1.device, data2.device)
# RuntimeError: Expected all tensors to be on the same device,
# but found at least two devices, cuda:0 and cpu!
data = data1 + data2
print(data)
if __name__ == '__main__':
test04()
5. 常见误区
第一个误区，把阿达玛积当成矩阵乘法。阿达玛积要求两个张量形状完全相同，对应位置相乘。矩阵乘法要求中间维度相等，结果形状由两侧维度决定。形状对不上时矩阵乘法直接报错，阿达玛积则要求同形状，两者的报错信息也不一样。遇到形状错误，先想清楚自己到底想做哪种乘法。
第二个误区，以为带下划线的函数会返回新张量。add_ 这类就地函数直接改写原对象，后面再引用原数据时它已经变了。需要保留中间结果做对比或者算梯度时，用不带下划线的版本。
第三个误区，拿 mm 去乘三维张量。mm 只接受二维，批量场景要么用 bmm，要么用 matmul，维度不固定时 matmul 最省心。
第四个误区，把 CPU 张量和 GPU 张量直接相加。不同设备不能运算，统一用一个 device 变量创建和搬运张量，能从根上避开这个错。
小结一下。这一节把张量的计算分成了四块。逐元素运算用 add、sub 这一组函数，带下划线的版本就地改写数据，省内存但要小心梯度。阿达玛积是同形状张量对应位置相乘，矩阵乘法是中间维度对齐后整行乘整列求和，前者用 mul 或 *，后者用 mm、bmm 或 matmul，维度越灵活越往后选。最后所有运算都要在同一台设备上做，用 torch.cuda.is_available() 统一 device，CPU 和 GPU 之间随时可以用 to、cuda、cpu 互相搬运。把上面几段代码在本地跑一遍，数字和形状都对得上，这部分就算过关了。
冀公网安备13050302001966号