【Python打卡Day48】随机张量与广播机制@浙大疏锦行
在继续讲解模块消融前，先补充几个之前没提的基础概念
尤其需要搞懂张量的维度、以及计算后的维度，这对于你未来理解复杂的网络至关重要
一、 随机张量的生成
在深度学习中经常需要随机生成一些张量，比如权重的初始化，或者计算输入纬度经过模块后输出的维度，都可以用一个随机函数来实现需要的张量格式，而无需像之前一样必须加载一张真实的图片。
“张量”概念
它听起来可能有点抽象，但在数学和物理学（尤其是广义相对论、连续介质力学、电磁学、机器学习等）中非常重要。
最核心的理解：张量是向量的推广，是一种独立于坐标系存在的几何或物理实体，其分量在坐标系变换时遵循特定的变换规则。
让我们一步步拆解：
-
从标量、向量、矩阵说起：
- 标量：这是最简单的量。它只有一个分量（一个数字），并且在坐标系变换下保持不变。例如：温度、质量、密度（在特定点）。标量是零阶张量。
- 向量：在空间中，它有大小和方向。在坐标系中，它需要用多个分量（例如在三维空间是 x, y, z）来表示。关键点在于：当你旋转坐标系时，向量的分量会按照特定的规则（线性变换）改变，但向量本身（这个几何对象）没有变。向量是一阶张量。
- 矩阵：是一个二维数组的数字。它可以表示线性变换（如旋转、缩放）、两个向量之间的关系（如应力、应变）等。矩阵是二阶张量的一种表示形式。
-
张量的核心特征：
- 多重线性性：张量是“多线性映射”。简单来说，它就像一个机器，你输入若干个向量（或对偶向量），它输出一个标量，并且它对每个输入都是线性的。
- 坐标系独立性：张量本身是一个几何对象或物理量，它独立于我们描述它的坐标系而存在。就像地球上的一个点，无论你用经度纬度还是笛卡尔坐标描述，那个点本身是不变的。
- 分量变换规则：这是识别张量的关键。当我们改变描述空间的坐标系时，张量的分量必须按照特定且一致的规则进行变换：
- 协变分量：变换方式类似于基向量的变换规则。
- 逆变分量：变换方式类似于坐标微分的变换规则（与基向量变换的规则“相反”）。
- 一个真正的张量，其分量在坐标变换时必须同时满足协变和逆变分量的变换规则（取决于它的指标类型）。如果一个量在坐标变换下不满足这些规则，它就不能被称为张量（尽管它可能是个有用的数组）。
-
阶数： 张量的阶数表示它需要多少个索引来描述其分量。
- 0阶张量：标量（无索引）。
- 1阶张量：向量（1个索引，如 Vⁱ 或 Vᵢ）。ⁱ 通常表示逆变分量，ᵢ 表示协变分量。
- 2阶张量：矩阵（2个索引，如 Tⁱʲ, Tⁱⱼ, Tᵢʲ, Tᵢⱼ）。例如：
- 应力张量：描述材料内部某点的应力状态（力/面积），需要 9 个分量（在三维空间）。
- 度规张量：定义空间的距离和角度（广义相对论的核心）。
- 3阶张量：3个索引（如 Rⁱⱼₖ），例如描述晶体性质的张量。
- 更高阶张量：以此类推。在机器学习中，“Tensor” 通常指任意维度的数组（ndarray），虽然其数学定义可能不如物理/数学中严格，但灵感来源于此。
-
为什么需要张量？
- 描述复杂物理现象：许多物理量（如应力、应变、电磁场、时空曲率）的内在关系无法用简单的标量或向量描述，它们天然地需要更高阶的对象。张量提供了一种统一、严谨的数学框架来描述这些量。
- 保证物理定律的普适性：物理定律应该在任何坐标系下都成立（广义协变性）。张量方程（等式两边都是同类型张量）在坐标变换下形式保持不变，这保证了物理定律的普适性。爱因斯坦的广义相对论场方程就是一个著名的张量方程。
- 简化计算：张量运算（如缩并、内积、外积）有明确的规则，使得处理复杂的多维关系变得系统化。
-
简单例子：
- 点积（标量积）：两个（逆变）向量 Uⁱ 和 Vⁱ 的点积 UⁱVᵢ（使用了爱因斯坦求和约定）是一个标量（零阶张量）。无论坐标系怎么旋转，这两个向量的点积结果不变。
- 应力：想象材料内部的一个小平面。作用在这个面上的力不一定垂直于这个面。你需要一个量同时指定：
- 作用面的方向（法向量，需要一个索引描述）。
- 作用在该面上的力的方向（另一个索引描述）。 因此，应力自然地被描述为一个二阶张量 σᵢʲ。
总结关键点：
- 张量是几何/物理实体：它独立于坐标系存在。
- 分量表示：在特定坐标系下，我们用带索引的分量数组来表示它。
- 变换规则：当坐标系改变时，分量必须按照协变/逆变规则一致地变换，以保持张量本身不变。满足这些变换规则是张量的定义性特征。
- 阶数：由索引的数量决定（0阶=标量，1阶=向量，2阶=矩阵等）。
- 作用：描述复杂物理量，构建坐标系无关的物理定律。
可以这样形象化理解：张量本身就像是一个物理对象（比如一块应力状态下的橡皮泥）。坐标系就像是从不同角度给它拍照。照片（分量）会随着拍照角度（坐标系）改变而改变，但橡皮泥本身（张量）并没有变。张量分量变换规则就是告诉你，当你转动相机时，照片上的像素应该如何正确地变化才能准确地反映同一个橡皮泥对象。
理解张量需要一些练习，特别是分量变换规则。但把握住“独立实体+特定变换规则”这个核心思想是第一步。
作用
随机函数的种类很多，我们了解其中一种即可，毕竟目的主要就是生成，对分布要求不重要。
在深度学习中，生成随机张量扮演着极其基础和关键的角色，其作用贯穿于模型的整个生命周期（初始化、训练、推理、评估）。以下是其主要作用：
-
模型参数初始化：
- 核心作用：打破对称性。神经网络在训练开始时，其权重和偏置（统称参数）必须被初始化。如果所有参数初始化为相同的值（例如全零），那么在反向传播过程中，所有神经元将接收到相同的梯度更新，导致它们学习到完全相同的特征，极大地限制了模型的表达能力和学习效果。
- 随机初始化：使用随机张量（通常从均匀分布或正态分布中采样）来初始化参数。这确保了：
- 不同的神经元从不同的起点开始学习。
- 梯度在反向传播过程中能够多样化地更新不同的参数。
- 模型有潜力学习到复杂、多样化的特征表示。常见的初始化方法如 Xavier/Glorot 初始化、He 初始化等，都基于特定的随机分布。
-
数据增强：
- 核心作用：增加数据多样性，提高泛化能力。为了缓解过拟合并提升模型对未见数据的泛化能力，在训练过程中经常对输入数据进行随机变换（如旋转、裁剪、缩放、颜色抖动、添加噪声等）。
- 随机张量的作用：这些变换的参数通常由随机张量控制。
- 例如：随机裁剪的位置和大小由随机坐标值决定；随机旋转的角度由随机角度值决定；随机添加的高斯噪声的强度由一个随机生成的噪声张量决定。这些随机值（标量、向量或张量本身）都是随机生成的。
-
噪声注入：
- 核心作用：正则化、探索、模拟不确定性。
- 正则化：在输入数据、隐藏层激活值或网络权重中直接添加随机噪声（如高斯噪声）。这迫使模型学习对输入中的微小扰动不敏感的特征，从而提高鲁棒性和泛化能力（类似于一种正则化技术）。
- 生成模型：这是生成模型（如 GANs, VAEs, Diffusion Models）的基石。
- GANs：生成器的输入通常是一个从先验分布（如标准正态分布）中采样的随机噪声向量（一个随机张量）。生成器学习将这个随机噪声映射到逼真的数据样本。
- VAEs：在训练过程中，编码器输出的潜在变量分布参数（均值和方差）被用来采样一个随机张量（潜在向量 z）。解码器则根据这个采样的 z 来重建输入或生成新样本。推理时也需要从潜在空间采样随机噪声。
- Diffusion Models：在正向过程中逐步向数据添加随机噪声，在逆向过程中需要预测并移除噪声。随机性贯穿整个生成过程。
- 探索（如强化学习）：在强化学习的策略中，有时会加入随机噪声（如 ε - greedy 或 Gaussian noise）来鼓励智能体探索未知的状态 - 动作空间。
-
正则化技术：
- 核心作用：防止过拟合。许多正则化技术直接依赖随机生成的张量。
- Dropout：在训练过程中，对于每个训练样本和每一层，以一定的概率 p 随机将一部分神经元的激活值置零（相当于丢弃这些神经元）。实现时，会生成一个与该层激活值形状相同的随机二值掩码张量（每个元素独立地以概率 1 - p 为 1，以概率 p 为 0），然后用这个掩码张量去乘以激活值。这个随机掩码张量就是随机生成的。
- DropConnect：类似 Dropout，但随机置零的是权重连接本身，而不是激活值，也需要生成随机掩码。
- Stochastic Depth：在训练深度网络时，随机跳过某些层（将其置为恒等映射）。这需要一个随机决策（通常是基于伯努利分布）来决定是否跳过某一层，该决策可以视为一个随机标量或掩码。
-
随机优化过程：
- 核心作用：避免局部极小值，提高效率。
- 小批量随机梯度下降：虽然核心是数据，但其“随机性”体现在每次迭代随机采样一个小批量数据来计算梯度和更新参数。这种采样过程依赖于随机索引（可以看作是一种特殊的随机张量生成）。
- 随机权重平均：一些高级优化器会维护模型参数的多个随机扰动副本或探索路径。
-
评估与基准测试：
- 核心作用：可复现性，创建合成数据。
- 设置随机种子：为了确保实验结果的可复现性，通常在代码开始时设置一个随机种子。这确保每次运行生成的随机数序列（即随机张量）是相同的，从而使模型初始化、数据打乱顺序、数据增强变换、Dropout 掩码等变得确定。
- 创建合成数据集：用于快速原型设计或测试模型架构时，可能需要生成随机的合成输入数据和对应的（可能也是随机生成的）标签来进行初步验证。
总结来说，生成随机张量的核心价值在于引入“可控的随机性”：
- 打破对称性：让模型学习过程能够启动并走向多样化（初始化）。
- 引入扰动：提高模型对噪声、变换的鲁棒性，防止过拟合（数据增强、噪声注入、正则化）。
- 模拟不确定性/探索：在生成模型中创造多样性，在强化学习中鼓励探索。
- 作为生成源：直接作为生成模型的输入或中间状态（生成模型）。
- 保证可复现性：通过固定随机种子控制随机过程（评估）。
可以说，没有随机张量的生成，深度学习模型的训练和许多核心功能（尤其是生成模型）将无法有效进行。它是深度学习不可或缺的“建筑材料”之一。
1.1 torch.randn函数
在 PyTorch 中，torch.randn()是一个常用的随机张量生成函数，它可以创建一个由标准正态分布（均值为 0，标准差为 1）随机数填充的张量。这种随机张量在深度学习中非常实用，常用于初始化模型参数、生成测试数据或模拟输入特征。
torch.randn(*size, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False)
- size：必选参数，表示输出张量的形状（如(3, 4)表示 3 行 4 列的矩阵）。
- dtype：可选参数，指定张量的数据类型（如torch.float32、torch.int64等）。
- device：可选参数，指定张量存储的设备（如'cpu'或'cuda'）。
- requires_grad：可选参数，是否需要计算梯度（常用于训练模型时）。
import torch
# 生成标量（0维张量）
scalar = torch.randn(())
print(f"标量: {scalar}, 形状: {scalar.shape}") 标量: -1.2918272018432617, 形状: torch.Size([])
# 生成向量（1维张量）
vector = torch.randn(5) # 长度为5的向量
print(f"向量: {vector}, 形状: {vector.shape}") 向量: tensor([ 1.1525, 0.5698, -1.1284, -0.1001, 1.1706]), 形状: torch.Size([5])
# 生成矩阵（2维张量）
matrix = torch.randn(3, 4) # 3行4列的矩阵
print(f"矩阵：{matrix},矩阵形状: {matrix.shape}") 矩阵：tensor([[ 0.0283, 0.7692, 0.2744, -1.6120],
[ 0.3726, 1.5382, -1.0128, 0.4129],
[ 0.4898, 1.4782, 0.2019, 0.0863]]),矩阵形状: torch.Size([3, 4])
# 生成3维张量（常用于图像数据的通道、高度、宽度）
tensor_3d = torch.randn(3, 224, 224) # 3通道，高224，宽224
print(f"3维张量形状: {tensor_3d.shape}") # 输出: torch.Size([3, 224, 224])3维张量形状: torch.Size([3, 224, 224])
# 生成4维张量（常用于批量图像数据：[batch, channel, height, width]）
tensor_4d = torch.randn(2, 3, 224, 224) # 批量大小为2，3通道，高224，宽224
print(f"4维张量形状: {tensor_4d.shape}") # 输出: torch.Size([2, 3, 224, 224])4维张量形状: torch.Size([2, 3, 224, 224])
1.2 其他随机函数
除了这些随机函数还有很多，自行了解，主要是生成数据的分布不同。掌握一个即可，掌握多了参数也记不住。
torch.rand()：生成在 [0, 1) 范围内均匀分布的随机数。
x = torch.rand(3, 2) # 生成3x2的张量
print(f"均匀分布随机数: {x}, 形状: {x.shape}")均匀分布随机数: tensor([[0.3393, 0.5227],
[0.5541, 0.4965],
[0.9453, 0.1888]]), 形状: torch.Size([3, 2])
torch.randint()：生成指定范围内的随机整数
x = torch.randint(low=0, high=10, size=(3,)) # 生成3个0到9之间的整数
print(f"随机整数: {x}, 形状: {x.shape}")随机整数: tensor([1, 7, 1]), 形状: torch.Size([3])
torch.normal()：生成指定均值和标准差的正态分布随机数。
mean = torch.tensor([0.0, 0.0])
std = torch.tensor([1.0, 2.0])
x = torch.normal(mean, std) # 生成两个正态分布随机数
print(f"正态分布随机数: {x}, 形状: {x.shape}")正态分布随机数: tensor([-1.7893, 0.5234]), 形状: torch.Size([2])
# 一维张量与二维张量相加
a = torch.tensor([[1, 2, 3], [4, 5, 6]]) # 形状: (2, 3)
b = torch.tensor([10, 20, 30]) # 形状: (3,)
# 广播后：b被扩展为[[10, 20, 30], [10, 20, 30]]
result = a + b
resulttensor([[11, 22, 33],
[14, 25, 36]])
1.3 输出维度测试
import torch
import torch.nn as nn
# 生成输入张量 (批量大小, 通道数, 高度, 宽度)
input_tensor = torch.randn(1, 3, 32, 32) # 例如CIFAR-10图像
print(f"输入尺寸: {input_tensor.shape}") # 输出: [1, 3, 32, 32]
# 1. 卷积层操作
conv1 = nn.Conv2d(
in_channels=3, # 输入通道数
out_channels=16, # 输出通道数（卷积核数量）
kernel_size=3, # 卷积核大小
stride=1, # 步长
padding=1 # 填充
)
conv_output = conv1(input_tensor) # 由于 padding=1 且 stride=1，空间尺寸保持不变
print(f"卷积后尺寸: {conv_output.shape}") # 输出: [1, 16, 32, 32]卷积后尺寸: torch.Size([1, 16, 32, 32])
# 2. 池化层操作 (减小空间尺寸)
pool = nn.MaxPool2d(kernel_size=2, stride=2) # 创建一个最大池化层
pool_output = pool(conv_output)
print(f"池化后尺寸: {pool_output.shape}") # 输出: [1, 16, 16, 16]池化后尺寸: torch.Size([1, 16, 16, 16])
# 3. 将多维张量展平为向量
flattened = pool_output.view(pool_output.size(0), -1)
print(f"展平后尺寸: {flattened.shape}") # 输出: [1, 4096] (16*16*16=4096)
展平后尺寸: torch.Size([1, 4096])
# 4. 线性层操作
fc1 = nn.Linear(
in_features=4096, # 输入特征数
out_features=128 # 输出特征数
)
fc_output = fc1(flattened)
print(f"线性层后尺寸: {fc_output.shape}") # 输出: [1, 128]线性层后尺寸: torch.Size([1, 128])
# 5. 再经过一个线性层（例如分类器）
fc2 = nn.Linear(128, 10) # 假设是10分类问题
final_output = fc2(fc_output)
print(f"最终输出尺寸: {final_output.shape}") # 输出: [1, 10]
print(final_output)最终输出尺寸: torch.Size([1, 10])
tensor([[-0.3018, -0.4308, 0.3248, 0.2808, 0.5109, -0.0881, -0.0787, -0.0700,
-0.1004, -0.0580]], grad_fn=
多分类问题通常使用Softmax，二分类问题用Sigmoid
# 使用Softmax替代Sigmoid
softmax = nn.Softmax(dim=1) # 在类别维度上进行Softmax
class_probs = softmax(final_output)
print(f"Softmax输出: {class_probs}") # 总和为1的概率分布
print(f"Softmax输出总和: {class_probs.sum():.4f}")Softmax输出: tensor([[0.0712, 0.0626, 0.1332, 0.1275, 0.1605, 0.0882, 0.0890, 0.0898, 0.0871,
0.0909]], grad_fn=
Softmax输出总和: 1.0000
通过这种方法，可以很自然的看到每一层输出的shape，实际上在pycharm等非交互式环境ipynb时，可以借助断点+调试控制台不断测试维度信息，避免报错。
二、广播机制
什么叫做广播机制
PyTorch 的广播机制（Broadcasting）是一种强大的张量运算特性，允许在不同形状的张量之间进行算术运算，而无需显式地扩展张量维度或复制数据。这种机制使得代码更简洁高效，尤其在处理多维数据时非常实用。
当对两个形状不同的张量进行运算时，PyTorch 会自动调整它们的形状，使它们在维度上兼容。具体规则如下：
从右向左比较维度：PyTorch 从张量的最后一个维度开始向前比较，检查每个维度的大小是否相同或其中一个为 1。
维度扩展规则：
如果两个张量的某个维度大小相同，则继续比较下一个维度。
如果其中一个张量的某个维度大小为 1，则该维度会被扩展为另一个张量对应维度的大小。
如果两个张量的某个维度大小既不相同也不为 1，则会报错。
2.1 PyTorch 广播机制
PyTorch 的广播机制（Broadcasting）是一种高效的张量运算特性，允许在不同形状的张量之间执行元素级操作（如加法、乘法），而无需显式扩展或复制数据。这种机制通过自动调整张量维度来实现形状兼容，使代码更简洁、计算更高效。
当对两个形状不同的张量进行运算时，PyTorch 会按以下规则自动处理维度兼容性：
- 从右向左比较维度：PyTorch 从张量的最后一个维度（最右侧）开始向前逐维比较。
- 维度扩展条件：
- 相等维度：若两个张量在某一维度上大小相同，则继续比较下一维度。
- 一维扩展：若其中一个张量在某一维度上大小为 1，则该维度会被扩展为另一个张量对应维度的大小。
- 不兼容错误：若某一维度大小既不相同也不为 1，则抛出 RuntimeError。维度必须满足广播规则，否则会报错。
- 维度补全规则：若一个张量的维度少于另一个，则在其左侧补 1 直至维度数匹配。
关注2个信息
1. 广播后的尺寸变化
2. 扩展后的值变化
2.2 加法的广播机制
二维张量与一维向量相加
import torch
# 创建原始张量
a = torch.tensor([[10], [20], [30]]) # 形状: (3, 1)
b = torch.tensor([1, 2, 3]) # 形状: (3,)
result = a + b
# 广播过程
# 1. b补全维度: (3,) → (1, 3)
# 2. a扩展列: (3, 1) → (3, 3)
# 3. b扩展行: (1, 3) → (3, 3)
# 最终形状: (3, 3)
print("原始张量a:")
print(a)
print("\n原始张量b:")
print(b)
print("\n广播后a的值扩展:")
print(torch.tensor([[10, 10, 10],
[20, 20, 20],
[30, 30, 30]])) # 实际内存中未复制，仅逻辑上扩展
print("\n广播后b的值扩展:")
print(torch.tensor([[1, 2, 3],
[1, 2, 3],
[1, 2, 3]])) # 实际内存中未复制，仅逻辑上扩展
print("\n加法结果:")
print(result)原始张量a:
tensor([[10],
[20],
[30]])
原始张量b:
tensor([1, 2, 3])
广播后a的值扩展:
tensor([[10, 10, 10],
[20, 20, 20],
[30, 30, 30]])
广播后b的值扩展:
tensor([[1, 2, 3],
[1, 2, 3],
[1, 2, 3]])
加法结果:
tensor([[11, 12, 13],
[21, 22, 23],
[31, 32, 33]])
三维张量与二维张量相加
# 创建原始张量
a = torch.tensor([[[1], [2]], [[3], [4]]]) # 形状: (2, 2, 1)
b = torch.tensor([[10, 20]]) # 形状: (1, 2)
# 广播过程
# 1. b补全维度: (1, 2) → (1, 1, 2)
# 2. a扩展第三维: (2, 2, 1) → (2, 2, 2)
# 3. b扩展第一维: (1, 1, 2) → (2, 1, 2)
# 4. b扩展第二维: (2, 1, 2) → (2, 2, 2)
# 最终形状: (2, 2, 2)
result = a + b
print("原始张量a:")
print(a)
print("\n原始张量b:")
print(b)
print("\n广播后a的值扩展:")
print(torch.tensor([[[1, 1],
[2, 2]],
[[3, 3],
[4, 4]]])) # 实际内存中未复制，仅逻辑上扩展
print("\n广播后b的值扩展:")
print(torch.tensor([[[10, 20],
[10, 20]],
[[10, 20],
[10, 20]]])) # 实际内存中未复制，仅逻辑上扩展
print("\n加法结果:")
print(result)
原始张量a:
tensor([[[1],
[2]],
[[3],
[4]]])
原始张量b:
tensor([[10, 20]])
广播后a的值扩展:
tensor([[[1, 1],
[2, 2]],
[[3, 3],
[4, 4]]])
广播后b的值扩展:
tensor([[[10, 20],
[10, 20]],
[[10, 20],
[10, 20]]])
加法结果:
tensor([[[11, 21],
[12, 22]],
[[13, 23],
[14, 24]]])
二维张量与标量相加
# 创建原始张量
a = torch.tensor([[1, 2], [3, 4]]) # 形状: (2, 2)
b = 10 # 标量，形状视为 ()
# 广播过程
# 1. b补全维度: () → (1, 1)
# 2. b扩展第一维: (1, 1) → (2, 1)
# 3. b扩展第二维: (2, 1) → (2, 2)
# 最终形状: (2, 2)
result = a + b
print("原始张量a:")
print(a)
# 输出:
# tensor([[1, 2],
# [3, 4]])
print("\n标量b:")
print(b)
# 输出: 10
print("\n广播后b的值扩展:")
print(torch.tensor([[10, 10],
[10, 10]])) # 实际内存中未复制，仅逻辑上扩展
print("\n加法结果:")
print(result)
# 输出:
# tensor([[11, 12],
# [13, 14]])高维张量与低维张量相加
# 创建原始张量
a = torch.tensor([[[1, 2], [3, 4]]]) # 形状: (1, 2, 2)
b = torch.tensor([[5, 6]]) # 形状: (1, 2)
# 广播过程
# 1. b补全维度: (1, 2) → (1, 1, 2)
# 2. b扩展第二维: (1, 1, 2) → (1, 2, 2)
# 最终形状: (1, 2, 2)
result = a + b
print("原始张量a:")
print(a)
# 输出:
# tensor([[[1, 2],
# [3, 4]]])
print("\n原始张量b:")
print(b)
# 输出:
# tensor([[5, 6]])
print("\n广播后b的值扩展:")
print(torch.tensor([[[5, 6],
[5, 6]]])) # 实际内存中未复制，仅逻辑上扩展
print("\n加法结果:")
print(result)
# 输出:
# tensor([[[6, 8],
# [8, 10]]])原始张量a:
tensor([[[1, 2],
[3, 4]]])
原始张量b:
tensor([[5, 6]])
广播后b的值扩展:
tensor([[[5, 6],
[5, 6]]])
加法结果:
tensor([[[ 6, 8],
[ 8, 10]]])
关键总结
1. 尺寸变化：广播后的形状由各维度的最大值决定（示例 2 中最终形状为 (2, 2, 2)）。
2. 值扩展：维度为 1 的张量通过复制扩展值（示例 1 中 b 从 [1, 2, 3] 扩展为三行相同的值）。
3. 内存效率：扩展是逻辑上的，实际未复制数据，避免了内存浪费。
2.3 乘法的广播机制
矩阵乘法（@）的特殊规则
矩阵乘法除了遵循通用广播规则外，还需要满足矩阵乘法的维度约束：
最后两个维度必须满足：A.shape[-1] == B.shape[-2]（即 A 的列数等于 B 的行数）
其他维度（批量维度）：遵循通用广播规则
批量矩阵与单个矩阵相乘
import torch
# A: 批量大小为2，每个是3×4的矩阵
A = torch.randn(2, 3, 4) # 形状: (2, 3, 4)
# B: 单个4×5的矩阵
B = torch.randn(4, 5) # 形状: (4, 5)
# 广播过程：
# 1. B补全维度: (4, 5) → (1, 4, 5)
# 2. B扩展第一维: (1, 4, 5) → (2, 4, 5)
# 矩阵乘法: (2, 3, 4) @ (2, 4, 5) → (2, 3, 5)
result = A @ B # 结果形状: (2, 3, 5)
print("A形状:", A.shape) # 输出: torch.Size([2, 3, 4])
print("B形状:", B.shape) # 输出: torch.Size([4, 5])
print("结果形状:", result.shape) # 输出: torch.Size([2, 3, 5])A形状: torch.Size([2, 3, 4])
B形状: torch.Size([4, 5])
结果形状: torch.Size([2, 3, 5])
批量矩阵与批量矩阵相乘（部分广播）
# A: 批量大小为3，每个是2×4的矩阵
A = torch.randn(3, 2, 4) # 形状: (3, 2, 4)
# B: 批量大小为1，每个是4×5的矩阵
B = torch.randn(1, 4, 5) # 形状: (1, 4, 5)
# 广播过程：
# B扩展第一维: (1, 4, 5) → (3, 4, 5)
# 矩阵乘法: (3, 2, 4) @ (3, 4, 5) → (3, 2, 5)
result = A @ B # 结果形状: (3, 2, 5)
print("A形状:", A.shape) # 输出: torch.Size([3, 2, 4])
print("B形状:", B.shape) # 输出: torch.Size([1, 4, 5])
print("结果形状:", result.shape) # 输出: torch.Size([3, 2, 5])A形状: torch.Size([3, 2, 4])
B形状: torch.Size([1, 4, 5])
结果形状: torch.Size([3, 2, 5])
三维张量与二维张量相乘（高维广播）
# A: 批量大小为2，通道数为3，每个是4×5的矩阵
A = torch.randn(2, 3, 4, 5) # 形状: (2, 3, 4, 5)
# B: 单个5×6的矩阵
B = torch.randn(5, 6) # 形状: (5, 6)
# 广播过程：
# 1. B补全维度: (5, 6) → (1, 1, 5, 6)
# 2. B扩展第一维: (1, 1, 5, 6) → (2, 1, 5, 6)
# 3. B扩展第二维: (2, 1, 5, 6) → (2, 3, 5, 6)
# 矩阵乘法: (2, 3, 4, 5) @ (2, 3, 5, 6) → (2, 3, 4, 6)
result = A @ B # 结果形状: (2, 3, 4, 6)
print("A形状:", A.shape) # 输出: torch.Size([2, 3, 4, 5])
print("B形状:", B.shape) # 输出: torch.Size([5, 6])
print("结果形状:", result.shape) # 输出: torch.Size([2, 3, 4, 6])A形状: torch.Size([2, 3, 4, 5])
B形状: torch.Size([5, 6])
结果形状: torch.Size([2, 3, 4, 6])
@浙大疏锦行