【Bug已解决】How to multiply a matrix by a vector in PyTorch 解决方案问题描述在 PyTorch 中进行矩阵与向量的乘法时很多开发者会遇到维度不匹配或使用了错误函数的问题。常见的困惑包括torch.mm、torch.matmul、torch.bmm、运算符之间的区别以及如何正确处理批量矩阵乘法。典型的问题场景import torch # 场景1矩阵乘向量 matrix torch.randn(3, 4) # 3x4 矩阵 vector torch.randn(4) # 长度为4的向量 # 期望结果3x4 4 - 长度为3的向量 # 尝试使用 torch.mm try: result torch.mm(matrix, vector) except RuntimeError as e: print(fError: {e}) # [3, 4] [4] - RuntimeError: 1-dimensional tensor expected# 场景2批量矩阵乘向量 batch_matrix torch.randn(10, 3, 4) # 10个3x4矩阵 vector torch.randn(4) # 长度为4的向量 # 期望结果10个长度为3的向量 try: result torch.mm(batch_matrix, vector) except RuntimeError as e: print(fError: {e})# 场景3使用 运算符 matrix torch.randn(3, 4) vector torch.randn(4) result matrix vector # 这实际上是正确的 print(fResult shape: {result.shape}) # [3]错误复现import torch # 复现1torch.mm 不支持向量 print( * 50) print(复现1torch.mm 需要2维输入) print( * 50) matrix torch.randn(3, 4) vector torch.randn(4) try: result torch.mm(matrix, vector) except RuntimeError as e: print(fError: {e}) # 2-dimensional tensor expected, got 1-dimensional tensor # torch.mm 只支持 2D x 2D matrix2 torch.randn(3, 4) result torch.mm(matrix2, matrix2.T) print(f2D x 2D: {result.shape}) # [3, 3]# 复现2torch.bmm 需要3维输入 print(\n * 50) print(复现2torch.bmm 需要3维输入) print( * 50) batch_matrix torch.randn(10, 3, 4) batch_matrix2 torch.randn(10, 4, 5) result torch.bmm(batch_matrix, batch_matrix2) print(f3D x 3D: {result.shape}) # [10, 3, 5] # 但不能用于向量 vector torch.randn(4) try: result torch.bmm(batch_matrix, vector) except RuntimeError as e: print(fError: {e})# 复现3维度顺序错误 print(\n * 50) print(复现3维度顺序错误) print( * 50) matrix torch.randn(3, 4) vector torch.randn(3) # 注意长度为3不是4 try: result matrix vector except RuntimeError as e: print(fError: {e}) # size mismatch: matrix [3, 4] and vector [3] cannot be multiplied根因分析1. PyTorch 矩阵乘法函数对比PyTorch 提供了多个矩阵乘法函数它们有不同的适用场景函数输入维度说明torch.mm2D x 2D矩阵乘法不支持广播torch.bmm3D x 3D批量矩阵乘法torch.matmul1D-ND通用矩阵乘法支持广播运算符1D-ND等同于torch.matmultorch.mul/*任意逐元素乘法不是矩阵乘法import torch # torch.mm: 只支持 2D x 2D A torch.randn(3, 4) B torch.randn(4, 5) print(fmm: {torch.mm(A, B).shape}) # [3, 5] # torch.bmm: 只支持 3D x 3D A_batch torch.randn(10, 3, 4) B_batch torch.randn(10, 4, 5) print(fbmm: {torch.bmm(A_batch, B_batch).shape}) # [10, 3, 5] # torch.matmul: 支持多种维度 print(fmatmul 2D1D: {torch.matmul(A, torch.randn(4)).shape}) # [3] print(fmatmul 2D2D: {torch.matmul(A, B).shape}) # [3, 5] print(fmatmul 3D3D: {torch.matmul(A_batch, B_batch).shape}) # [10, 3, 5]2. matmul 的广播规则torch.matmul的广播规则比较复杂1D x 1D点积返回标量2D x 1D矩阵乘向量返回1D1D x 2D向量乘矩阵返回1D2D x 2D矩阵乘法返回2DND x ND批量矩阵乘法前面的维度会广播# 1D x 1D: 点积 a torch.randn(4) b torch.randn(4) print(f1D x 1D: {torch.matmul(a, b).shape}) # [] (标量) # 2D x 1D: 矩阵乘向量 A torch.randn(3, 4) v torch.randn(4) print(f2D x 1D: {torch.matmul(A, v).shape}) # [3] # 1D x 2D: 向量乘矩阵 v torch.randn(3) B torch.randn(3, 4) print(f1D x 2D: {torch.matmul(v, B).shape}) # [4] # 2D x 2D: 矩阵乘法 A torch.randn(3, 4) B torch.randn(4, 5) print(f2D x 2D: {torch.matmul(A, B).shape}) # [3, 5] # 3D x 3D: 批量矩阵乘法 A torch.randn(2, 3, 4) B torch.randn(2, 4, 5) print(f3D x 3D: {torch.matmul(A, B).shape}) # [2, 3, 5] # 广播: 3D x 2D A torch.randn(2, 3, 4) B torch.randn(4, 5) print(f3D x 2D: {torch.matmul(A, B).shape}) # [2, 3, 5]3. 矩阵乘向量的特殊情况当使用matmul进行矩阵乘向量时向量会被自动扩展# 矩阵 [3, 4] 乘向量 [4] - [3] A torch.randn(3, 4) v torch.randn(4) result torch.matmul(A, v) print(fMatrix Vector: {A.shape} {v.shape} - {result.shape}) # 批量矩阵 [10, 3, 4] 乘向量 [4] # 需要将向量扩展为 [10, 4, 1] 或使用 unsqueeze A_batch torch.randn(10, 3, 4) v torch.randn(4) # 方法1使用 matmul自动广播 result torch.matmul(A_batch, v) print(fBatch Matrix Vector: {A_batch.shape} {v.shape} - {result.shape}) # [10, 3]解决方案方案一使用 torch.matmul推荐torch.matmul是最通用的矩阵乘法函数支持各种维度的输入。import torch # 矩阵乘向量 matrix torch.randn(3, 4) vector torch.randn(4) result torch.matmul(matrix, vector) print(fMatrix Vector: {result.shape}) # [3] # 等价于使用 运算符 result2 matrix vector print(fEqual: {torch.equal(result, result2)}) # True # 批量矩阵乘向量 batch_matrix torch.randn(10, 3, 4) vector torch.randn(4) result torch.matmul(batch_matrix, vector) print(fBatch Vector: {result.shape}) # [10, 3] # 矩阵乘矩阵 A torch.randn(3, 4) B torch.randn(4, 5) result torch.matmul(A, B) print(fMatrix Matrix: {result.shape}) # [3, 5]方案二使用 unsqueeze mm/bmm如果需要更精确的控制可以使用unsqueeze将向量扩展为矩阵然后使用mm或bmm。import torch # 矩阵乘向量使用 mm matrix torch.randn(3, 4) vector torch.randn(4) # 将向量扩展为列向量 [4, 1] vector_col vector.unsqueeze(1) # [4, 1] result torch.mm(matrix, vector_col) # [3, 1] result result.squeeze(1) # [3] print(fmm result: {result.shape}) # 批量矩阵乘向量使用 bmm batch_matrix torch.randn(10, 3, 4) vector torch.randn(4) # 将向量扩展为 [10, 4, 1] vector_batch vector.unsqueeze(0).unsqueeze(2).expand(10, -1, -1) # 或 vector_batch vector.unsqueeze(1).unsqueeze(0).expand(10, -1, -1) result torch.bmm(batch_matrix, vector_batch) # [10, 3, 1] result result.squeeze(2) # [10, 3] print(fbmm result: {result.shape})方案三使用 einsum复杂操作对于复杂的张量运算torch.einsum提供了更灵活的方式。import torch # 矩阵乘向量 matrix torch.randn(3, 4) vector torch.randn(4) result torch.einsum(ij,j-i, matrix, vector) print(feinsum ij,j-i: {result.shape}) # [3] # 批量矩阵乘向量 batch_matrix torch.randn(10, 3, 4) vector torch.randn(4) result torch.einsum(bij,j-bi, batch_matrix, vector) print(feinsum bij,j-bi: {result.shape}) # [10, 3] # 批量矩阵乘批量向量 batch_matrix torch.randn(10, 3, 4) batch_vector torch.randn(10, 4) result torch.einsum(bij,bj-bi, batch_matrix, batch_vector) print(feinsum bij,bj-bi: {result.shape}) # [10, 3] # 更复杂的操作注意力机制 Q torch.randn(2, 8, 4, 64) # [batch, heads, seq, dim] K torch.randn(2, 8, 4, 64) V torch.randn(2, 8, 4, 64) # Q K^T / sqrt(d) attn torch.einsum(bhid,bhjd-bhij, Q, K) / (64 ** 0.5) attn torch.softmax(attn, dim-1) # attn V output torch.einsum(bhij,bhjd-bhid, attn, V) print(fAttention output: {output.shape}) # [2, 8, 4, 64]方案四使用 F.linear神经网络场景在神经网络中F.linear是专门用于线性变换的函数。import torch import torch.nn.functional as F # F.linear: y x W^T b x torch.randn(10, 4) # [batch, in_features] W torch.randn(3, 4) # [out_features, in_features] b torch.randn(3) # [out_features] result F.linear(x, W, b) print(fF.linear: {result.shape}) # [10, 3] # 批量场景 x torch.randn(5, 10, 4) # [batch, seq, in_features] result F.linear(x, W, b) print(fBatch F.linear: {result.shape}) # [5, 10, 3]完整修复代码 完整修复代码PyTorch 矩阵与向量乘法 实现各种矩阵乘法场景的完整解决方案 import torch import torch.nn as nn import torch.nn.functional as F from typing import Optional class MatrixVectorOps: 矩阵向量运算工具类 staticmethod def mat_vec(matrix: torch.Tensor, vector: torch.Tensor) - torch.Tensor: 矩阵乘向量 [M, N] [N] - [M] assert matrix.dim() 2, fMatrix must be 2D, got {matrix.dim()}D assert vector.dim() 1, fVector must be 1D, got {vector.dim()}D assert matrix.size(1) vector.size(0), \ fSize mismatch: matrix {matrix.shape} vector {vector.shape} return torch.matmul(matrix, vector) staticmethod def batch_mat_vec(batch_matrix: torch.Tensor, vector: torch.Tensor) - torch.Tensor: 批量矩阵乘向量 [B, M, N] [N] - [B, M] assert batch_matrix.dim() 3, fBatch matrix must be 3D assert vector.dim() 1, fVector must be 1D assert batch_matrix.size(2) vector.size(0), Size mismatch return torch.matmul(batch_matrix, vector) staticmethod def batch_mat_batch_vec(batch_matrix: torch.Tensor, batch_vector: torch.Tensor) - torch.Tensor: 批量矩阵乘批量向量 [B, M, N] [B, N] - [B, M] assert batch_matrix.dim() 3 assert batch_vector.dim() 2 assert batch_matrix.size(0) batch_vector.size(0), Batch size mismatch assert batch_matrix.size(2) batch_vector.size(1), Feature size mismatch return torch.einsum(bmn,bn-bm, batch_matrix, batch_vector) staticmethod def batch_matmul(A: torch.Tensor, B: torch.Tensor) - torch.Tensor: 批量矩阵乘法支持广播 return torch.matmul(A, B) staticmethod def attention(Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor, mask: Optional[torch.Tensor] None) - torch.Tensor: 缩放点积注意力 d_k Q.size(-1) scores torch.matmul(Q, K.transpose(-2, -1)) / (d_k ** 0.5) if mask is not None: scores scores.masked_fill(mask 0, float(-inf)) attn_weights F.softmax(scores, dim-1) output torch.matmul(attn_weights, V) return output class LinearLayer(nn.Module): 自定义线性层展示矩阵乘法的应用 def __init__(self, in_features: int, out_features: int, bias: bool True): super().__init__() self.in_features in_features self.out_features out_features self.weight nn.Parameter(torch.randn(out_features, in_features)) if bias: self.bias nn.Parameter(torch.zeros(out_features)) else: self.register_parameter(bias, None) def forward(self, x: torch.Tensor) - torch.Tensor: # x W^T b # 使用 matmul 实现线性变换 output torch.matmul(x, self.weight.t()) if self.bias is not None: output output self.bias return output def test_operations(): 测试各种矩阵向量运算 print( * 60) print(Testing Matrix-Vector Operations) print( * 60) ops MatrixVectorOps() # 测试1矩阵乘向量 print(\n--- Test 1: Matrix Vector ---) matrix torch.randn(3, 4) vector torch.randn(4) result ops.mat_vec(matrix, vector) print(fMatrix {matrix.shape} Vector {vector.shape} - {result.shape}) # 测试2批量矩阵乘向量 print(\n--- Test 2: Batch Matrix Vector ---) batch_matrix torch.randn(10, 3, 4) vector torch.randn(4) result ops.batch_mat_vec(batch_matrix, vector) print(fBatch Matrix {batch_matrix.shape} Vector {vector.shape} - {result.shape}) # 测试3批量矩阵乘批量向量 print(\n--- Test 3: Batch Matrix Batch Vector ---) batch_matrix torch.randn(10, 3, 4) batch_vector torch.randn(10, 4) result ops.batch_mat_batch_vec(batch_matrix, batch_vector) print(fBatch Matrix {batch_matrix.shape} Batch Vector {batch_vector.shape} - {result.shape}) # 测试4注意力机制 print(\n--- Test 4: Attention Mechanism ---) Q torch.randn(2, 8, 4, 64) K torch.randn(2, 8, 4, 64) V torch.randn(2, 8, 4, 64) output ops.attention(Q, K, V) print(fAttention output: {output.shape}) # 测试5自定义线性层 print(\n--- Test 5: Custom Linear Layer ---) layer LinearLayer(10, 5) x torch.randn(4, 10) output layer(x) print(fLinear output: {output.shape}) # 测试6各种乘法函数对比 print(\n--- Test 6: Function Comparison ---) A torch.randn(3, 4) B torch.randn(4, 5) r1 torch.mm(A, B) r2 torch.matmul(A, B) r3 A B print(fmm vs matmul vs : {torch.allclose(r1, r2) and torch.allclose(r2, r3)}) # 测试7einsum 对比 print(\n--- Test 7: Einsum Comparison ---) matrix torch.randn(3, 4) vector torch.randn(4) r1 torch.matmul(matrix, vector) r2 torch.einsum(ij,j-i, matrix, vector) print(fmatmul vs einsum: {torch.allclose(r1, r2)}) # 测试8广播矩阵乘法 print(\n--- Test 8: Broadcasting ---) A torch.randn(2, 3, 4) # 批量矩阵 B torch.randn(4, 5) # 普通矩阵 result torch.matmul(A, B) print(fBroadcast [2,3,4] [4,5] - {result.shape}) print(\n * 60) print(All tests passed!) print( * 60) if __name__ __main__: test_operations()常见陷阱与注意事项1. mm vs matmul vs bmmA torch.randn(3, 4) B torch.randn(4, 5) # mm: 只支持 2D x 2D torch.mm(A, B) # OK # matmul: 支持各种维度 torch.matmul(A, B) # OK # bmm: 只支持 3D x 3D try: torch.bmm(A, B) # Error: 3D expected except RuntimeError as e: print(fError: {e})2. mul vs matmulA torch.randn(3, 4) B torch.randn(3, 4) # mul: 逐元素乘法 result torch.mul(A, B) # 或 A * B print(fmul: {result.shape}) # [3, 4] # matmul: 矩阵乘法 B2 torch.randn(4, 5) result torch.matmul(A, B2) print(fmatmul: {result.shape}) # [3, 5]3. 转置操作A torch.randn(3, 4) # 注意矩阵乘法需要维度匹配 # A A 会报错4 ! 3 try: torch.matmul(A, A) except RuntimeError as e: print(fError: {e}) # 正确A A^T 或 A^T A print(fA A^T: {torch.matmul(A, A.T).shape}) # [3, 3] print(fA^T A: {torch.matmul(A.T, A).shape}) # [4, 4]4. 向量的 unsqueezev torch.randn(4) # 1D 向量 print(f1D: {v.shape}) # [4] # 列向量 [4, 1] v_col v.unsqueeze(1) print(fColumn: {v_col.shape}) # [4, 1] # 行向量 [1, 4] v_row v.unsqueeze(0) print(fRow: {v_row.shape}) # [1, 4]5. 梯度传播A torch.randn(3, 4, requires_gradTrue) v torch.randn(4, requires_gradTrue) # matmul 支持自动微分 result torch.matmul(A, v) loss result.sum() loss.backward() print(fA.grad: {A.grad.shape}) # [3, 4] print(fv.grad: {v.grad.shape}) # [4]总结在 PyTorch 中进行矩阵与向量的乘法关键在于选择正确的函数和理解广播规则torch.matmul/最通用的矩阵乘法支持 1D 到 ND 的各种维度组合推荐优先使用。torch.mm只支持 2D x 2D 的矩阵乘法不支持广播。torch.bmm只支持 3D x 3D 的批量矩阵乘法。torch.mul/*逐元素乘法不是矩阵乘法。torch.einsum使用 Einstein 求和约定适合复杂的张量运算。F.linear神经网络专用的线性变换y xW^T b。最佳实践优先使用torch.matmul或运算符需要精确控制时使用mm或bmm复杂运算使用einsum注意维度匹配和转置操作神经网络中使用F.linear或nn.Linear通过本文的详细分析和完整代码示例你应该能够正确地在 PyTorch 中进行各种矩阵与向量的乘法运算避免常见的维度错误和函数误用。
【Bug已解决】How to multiply a matrix by a vector in PyTorch 解决方案
【Bug已解决】How to multiply a matrix by a vector in PyTorch 解决方案问题描述在 PyTorch 中进行矩阵与向量的乘法时很多开发者会遇到维度不匹配或使用了错误函数的问题。常见的困惑包括torch.mm、torch.matmul、torch.bmm、运算符之间的区别以及如何正确处理批量矩阵乘法。典型的问题场景import torch # 场景1矩阵乘向量 matrix torch.randn(3, 4) # 3x4 矩阵 vector torch.randn(4) # 长度为4的向量 # 期望结果3x4 4 - 长度为3的向量 # 尝试使用 torch.mm try: result torch.mm(matrix, vector) except RuntimeError as e: print(fError: {e}) # [3, 4] [4] - RuntimeError: 1-dimensional tensor expected# 场景2批量矩阵乘向量 batch_matrix torch.randn(10, 3, 4) # 10个3x4矩阵 vector torch.randn(4) # 长度为4的向量 # 期望结果10个长度为3的向量 try: result torch.mm(batch_matrix, vector) except RuntimeError as e: print(fError: {e})# 场景3使用 运算符 matrix torch.randn(3, 4) vector torch.randn(4) result matrix vector # 这实际上是正确的 print(fResult shape: {result.shape}) # [3]错误复现import torch # 复现1torch.mm 不支持向量 print( * 50) print(复现1torch.mm 需要2维输入) print( * 50) matrix torch.randn(3, 4) vector torch.randn(4) try: result torch.mm(matrix, vector) except RuntimeError as e: print(fError: {e}) # 2-dimensional tensor expected, got 1-dimensional tensor # torch.mm 只支持 2D x 2D matrix2 torch.randn(3, 4) result torch.mm(matrix2, matrix2.T) print(f2D x 2D: {result.shape}) # [3, 3]# 复现2torch.bmm 需要3维输入 print(\n * 50) print(复现2torch.bmm 需要3维输入) print( * 50) batch_matrix torch.randn(10, 3, 4) batch_matrix2 torch.randn(10, 4, 5) result torch.bmm(batch_matrix, batch_matrix2) print(f3D x 3D: {result.shape}) # [10, 3, 5] # 但不能用于向量 vector torch.randn(4) try: result torch.bmm(batch_matrix, vector) except RuntimeError as e: print(fError: {e})# 复现3维度顺序错误 print(\n * 50) print(复现3维度顺序错误) print( * 50) matrix torch.randn(3, 4) vector torch.randn(3) # 注意长度为3不是4 try: result matrix vector except RuntimeError as e: print(fError: {e}) # size mismatch: matrix [3, 4] and vector [3] cannot be multiplied根因分析1. PyTorch 矩阵乘法函数对比PyTorch 提供了多个矩阵乘法函数它们有不同的适用场景函数输入维度说明torch.mm2D x 2D矩阵乘法不支持广播torch.bmm3D x 3D批量矩阵乘法torch.matmul1D-ND通用矩阵乘法支持广播运算符1D-ND等同于torch.matmultorch.mul/*任意逐元素乘法不是矩阵乘法import torch # torch.mm: 只支持 2D x 2D A torch.randn(3, 4) B torch.randn(4, 5) print(fmm: {torch.mm(A, B).shape}) # [3, 5] # torch.bmm: 只支持 3D x 3D A_batch torch.randn(10, 3, 4) B_batch torch.randn(10, 4, 5) print(fbmm: {torch.bmm(A_batch, B_batch).shape}) # [10, 3, 5] # torch.matmul: 支持多种维度 print(fmatmul 2D1D: {torch.matmul(A, torch.randn(4)).shape}) # [3] print(fmatmul 2D2D: {torch.matmul(A, B).shape}) # [3, 5] print(fmatmul 3D3D: {torch.matmul(A_batch, B_batch).shape}) # [10, 3, 5]2. matmul 的广播规则torch.matmul的广播规则比较复杂1D x 1D点积返回标量2D x 1D矩阵乘向量返回1D1D x 2D向量乘矩阵返回1D2D x 2D矩阵乘法返回2DND x ND批量矩阵乘法前面的维度会广播# 1D x 1D: 点积 a torch.randn(4) b torch.randn(4) print(f1D x 1D: {torch.matmul(a, b).shape}) # [] (标量) # 2D x 1D: 矩阵乘向量 A torch.randn(3, 4) v torch.randn(4) print(f2D x 1D: {torch.matmul(A, v).shape}) # [3] # 1D x 2D: 向量乘矩阵 v torch.randn(3) B torch.randn(3, 4) print(f1D x 2D: {torch.matmul(v, B).shape}) # [4] # 2D x 2D: 矩阵乘法 A torch.randn(3, 4) B torch.randn(4, 5) print(f2D x 2D: {torch.matmul(A, B).shape}) # [3, 5] # 3D x 3D: 批量矩阵乘法 A torch.randn(2, 3, 4) B torch.randn(2, 4, 5) print(f3D x 3D: {torch.matmul(A, B).shape}) # [2, 3, 5] # 广播: 3D x 2D A torch.randn(2, 3, 4) B torch.randn(4, 5) print(f3D x 2D: {torch.matmul(A, B).shape}) # [2, 3, 5]3. 矩阵乘向量的特殊情况当使用matmul进行矩阵乘向量时向量会被自动扩展# 矩阵 [3, 4] 乘向量 [4] - [3] A torch.randn(3, 4) v torch.randn(4) result torch.matmul(A, v) print(fMatrix Vector: {A.shape} {v.shape} - {result.shape}) # 批量矩阵 [10, 3, 4] 乘向量 [4] # 需要将向量扩展为 [10, 4, 1] 或使用 unsqueeze A_batch torch.randn(10, 3, 4) v torch.randn(4) # 方法1使用 matmul自动广播 result torch.matmul(A_batch, v) print(fBatch Matrix Vector: {A_batch.shape} {v.shape} - {result.shape}) # [10, 3]解决方案方案一使用 torch.matmul推荐torch.matmul是最通用的矩阵乘法函数支持各种维度的输入。import torch # 矩阵乘向量 matrix torch.randn(3, 4) vector torch.randn(4) result torch.matmul(matrix, vector) print(fMatrix Vector: {result.shape}) # [3] # 等价于使用 运算符 result2 matrix vector print(fEqual: {torch.equal(result, result2)}) # True # 批量矩阵乘向量 batch_matrix torch.randn(10, 3, 4) vector torch.randn(4) result torch.matmul(batch_matrix, vector) print(fBatch Vector: {result.shape}) # [10, 3] # 矩阵乘矩阵 A torch.randn(3, 4) B torch.randn(4, 5) result torch.matmul(A, B) print(fMatrix Matrix: {result.shape}) # [3, 5]方案二使用 unsqueeze mm/bmm如果需要更精确的控制可以使用unsqueeze将向量扩展为矩阵然后使用mm或bmm。import torch # 矩阵乘向量使用 mm matrix torch.randn(3, 4) vector torch.randn(4) # 将向量扩展为列向量 [4, 1] vector_col vector.unsqueeze(1) # [4, 1] result torch.mm(matrix, vector_col) # [3, 1] result result.squeeze(1) # [3] print(fmm result: {result.shape}) # 批量矩阵乘向量使用 bmm batch_matrix torch.randn(10, 3, 4) vector torch.randn(4) # 将向量扩展为 [10, 4, 1] vector_batch vector.unsqueeze(0).unsqueeze(2).expand(10, -1, -1) # 或 vector_batch vector.unsqueeze(1).unsqueeze(0).expand(10, -1, -1) result torch.bmm(batch_matrix, vector_batch) # [10, 3, 1] result result.squeeze(2) # [10, 3] print(fbmm result: {result.shape})方案三使用 einsum复杂操作对于复杂的张量运算torch.einsum提供了更灵活的方式。import torch # 矩阵乘向量 matrix torch.randn(3, 4) vector torch.randn(4) result torch.einsum(ij,j-i, matrix, vector) print(feinsum ij,j-i: {result.shape}) # [3] # 批量矩阵乘向量 batch_matrix torch.randn(10, 3, 4) vector torch.randn(4) result torch.einsum(bij,j-bi, batch_matrix, vector) print(feinsum bij,j-bi: {result.shape}) # [10, 3] # 批量矩阵乘批量向量 batch_matrix torch.randn(10, 3, 4) batch_vector torch.randn(10, 4) result torch.einsum(bij,bj-bi, batch_matrix, batch_vector) print(feinsum bij,bj-bi: {result.shape}) # [10, 3] # 更复杂的操作注意力机制 Q torch.randn(2, 8, 4, 64) # [batch, heads, seq, dim] K torch.randn(2, 8, 4, 64) V torch.randn(2, 8, 4, 64) # Q K^T / sqrt(d) attn torch.einsum(bhid,bhjd-bhij, Q, K) / (64 ** 0.5) attn torch.softmax(attn, dim-1) # attn V output torch.einsum(bhij,bhjd-bhid, attn, V) print(fAttention output: {output.shape}) # [2, 8, 4, 64]方案四使用 F.linear神经网络场景在神经网络中F.linear是专门用于线性变换的函数。import torch import torch.nn.functional as F # F.linear: y x W^T b x torch.randn(10, 4) # [batch, in_features] W torch.randn(3, 4) # [out_features, in_features] b torch.randn(3) # [out_features] result F.linear(x, W, b) print(fF.linear: {result.shape}) # [10, 3] # 批量场景 x torch.randn(5, 10, 4) # [batch, seq, in_features] result F.linear(x, W, b) print(fBatch F.linear: {result.shape}) # [5, 10, 3]完整修复代码 完整修复代码PyTorch 矩阵与向量乘法 实现各种矩阵乘法场景的完整解决方案 import torch import torch.nn as nn import torch.nn.functional as F from typing import Optional class MatrixVectorOps: 矩阵向量运算工具类 staticmethod def mat_vec(matrix: torch.Tensor, vector: torch.Tensor) - torch.Tensor: 矩阵乘向量 [M, N] [N] - [M] assert matrix.dim() 2, fMatrix must be 2D, got {matrix.dim()}D assert vector.dim() 1, fVector must be 1D, got {vector.dim()}D assert matrix.size(1) vector.size(0), \ fSize mismatch: matrix {matrix.shape} vector {vector.shape} return torch.matmul(matrix, vector) staticmethod def batch_mat_vec(batch_matrix: torch.Tensor, vector: torch.Tensor) - torch.Tensor: 批量矩阵乘向量 [B, M, N] [N] - [B, M] assert batch_matrix.dim() 3, fBatch matrix must be 3D assert vector.dim() 1, fVector must be 1D assert batch_matrix.size(2) vector.size(0), Size mismatch return torch.matmul(batch_matrix, vector) staticmethod def batch_mat_batch_vec(batch_matrix: torch.Tensor, batch_vector: torch.Tensor) - torch.Tensor: 批量矩阵乘批量向量 [B, M, N] [B, N] - [B, M] assert batch_matrix.dim() 3 assert batch_vector.dim() 2 assert batch_matrix.size(0) batch_vector.size(0), Batch size mismatch assert batch_matrix.size(2) batch_vector.size(1), Feature size mismatch return torch.einsum(bmn,bn-bm, batch_matrix, batch_vector) staticmethod def batch_matmul(A: torch.Tensor, B: torch.Tensor) - torch.Tensor: 批量矩阵乘法支持广播 return torch.matmul(A, B) staticmethod def attention(Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor, mask: Optional[torch.Tensor] None) - torch.Tensor: 缩放点积注意力 d_k Q.size(-1) scores torch.matmul(Q, K.transpose(-2, -1)) / (d_k ** 0.5) if mask is not None: scores scores.masked_fill(mask 0, float(-inf)) attn_weights F.softmax(scores, dim-1) output torch.matmul(attn_weights, V) return output class LinearLayer(nn.Module): 自定义线性层展示矩阵乘法的应用 def __init__(self, in_features: int, out_features: int, bias: bool True): super().__init__() self.in_features in_features self.out_features out_features self.weight nn.Parameter(torch.randn(out_features, in_features)) if bias: self.bias nn.Parameter(torch.zeros(out_features)) else: self.register_parameter(bias, None) def forward(self, x: torch.Tensor) - torch.Tensor: # x W^T b # 使用 matmul 实现线性变换 output torch.matmul(x, self.weight.t()) if self.bias is not None: output output self.bias return output def test_operations(): 测试各种矩阵向量运算 print( * 60) print(Testing Matrix-Vector Operations) print( * 60) ops MatrixVectorOps() # 测试1矩阵乘向量 print(\n--- Test 1: Matrix Vector ---) matrix torch.randn(3, 4) vector torch.randn(4) result ops.mat_vec(matrix, vector) print(fMatrix {matrix.shape} Vector {vector.shape} - {result.shape}) # 测试2批量矩阵乘向量 print(\n--- Test 2: Batch Matrix Vector ---) batch_matrix torch.randn(10, 3, 4) vector torch.randn(4) result ops.batch_mat_vec(batch_matrix, vector) print(fBatch Matrix {batch_matrix.shape} Vector {vector.shape} - {result.shape}) # 测试3批量矩阵乘批量向量 print(\n--- Test 3: Batch Matrix Batch Vector ---) batch_matrix torch.randn(10, 3, 4) batch_vector torch.randn(10, 4) result ops.batch_mat_batch_vec(batch_matrix, batch_vector) print(fBatch Matrix {batch_matrix.shape} Batch Vector {batch_vector.shape} - {result.shape}) # 测试4注意力机制 print(\n--- Test 4: Attention Mechanism ---) Q torch.randn(2, 8, 4, 64) K torch.randn(2, 8, 4, 64) V torch.randn(2, 8, 4, 64) output ops.attention(Q, K, V) print(fAttention output: {output.shape}) # 测试5自定义线性层 print(\n--- Test 5: Custom Linear Layer ---) layer LinearLayer(10, 5) x torch.randn(4, 10) output layer(x) print(fLinear output: {output.shape}) # 测试6各种乘法函数对比 print(\n--- Test 6: Function Comparison ---) A torch.randn(3, 4) B torch.randn(4, 5) r1 torch.mm(A, B) r2 torch.matmul(A, B) r3 A B print(fmm vs matmul vs : {torch.allclose(r1, r2) and torch.allclose(r2, r3)}) # 测试7einsum 对比 print(\n--- Test 7: Einsum Comparison ---) matrix torch.randn(3, 4) vector torch.randn(4) r1 torch.matmul(matrix, vector) r2 torch.einsum(ij,j-i, matrix, vector) print(fmatmul vs einsum: {torch.allclose(r1, r2)}) # 测试8广播矩阵乘法 print(\n--- Test 8: Broadcasting ---) A torch.randn(2, 3, 4) # 批量矩阵 B torch.randn(4, 5) # 普通矩阵 result torch.matmul(A, B) print(fBroadcast [2,3,4] [4,5] - {result.shape}) print(\n * 60) print(All tests passed!) print( * 60) if __name__ __main__: test_operations()常见陷阱与注意事项1. mm vs matmul vs bmmA torch.randn(3, 4) B torch.randn(4, 5) # mm: 只支持 2D x 2D torch.mm(A, B) # OK # matmul: 支持各种维度 torch.matmul(A, B) # OK # bmm: 只支持 3D x 3D try: torch.bmm(A, B) # Error: 3D expected except RuntimeError as e: print(fError: {e})2. mul vs matmulA torch.randn(3, 4) B torch.randn(3, 4) # mul: 逐元素乘法 result torch.mul(A, B) # 或 A * B print(fmul: {result.shape}) # [3, 4] # matmul: 矩阵乘法 B2 torch.randn(4, 5) result torch.matmul(A, B2) print(fmatmul: {result.shape}) # [3, 5]3. 转置操作A torch.randn(3, 4) # 注意矩阵乘法需要维度匹配 # A A 会报错4 ! 3 try: torch.matmul(A, A) except RuntimeError as e: print(fError: {e}) # 正确A A^T 或 A^T A print(fA A^T: {torch.matmul(A, A.T).shape}) # [3, 3] print(fA^T A: {torch.matmul(A.T, A).shape}) # [4, 4]4. 向量的 unsqueezev torch.randn(4) # 1D 向量 print(f1D: {v.shape}) # [4] # 列向量 [4, 1] v_col v.unsqueeze(1) print(fColumn: {v_col.shape}) # [4, 1] # 行向量 [1, 4] v_row v.unsqueeze(0) print(fRow: {v_row.shape}) # [1, 4]5. 梯度传播A torch.randn(3, 4, requires_gradTrue) v torch.randn(4, requires_gradTrue) # matmul 支持自动微分 result torch.matmul(A, v) loss result.sum() loss.backward() print(fA.grad: {A.grad.shape}) # [3, 4] print(fv.grad: {v.grad.shape}) # [4]总结在 PyTorch 中进行矩阵与向量的乘法关键在于选择正确的函数和理解广播规则torch.matmul/最通用的矩阵乘法支持 1D 到 ND 的各种维度组合推荐优先使用。torch.mm只支持 2D x 2D 的矩阵乘法不支持广播。torch.bmm只支持 3D x 3D 的批量矩阵乘法。torch.mul/*逐元素乘法不是矩阵乘法。torch.einsum使用 Einstein 求和约定适合复杂的张量运算。F.linear神经网络专用的线性变换y xW^T b。最佳实践优先使用torch.matmul或运算符需要精确控制时使用mm或bmm复杂运算使用einsum注意维度匹配和转置操作神经网络中使用F.linear或nn.Linear通过本文的详细分析和完整代码示例你应该能够正确地在 PyTorch 中进行各种矩阵与向量的乘法运算避免常见的维度错误和函数误用。