实用指南:PyTorch中matmul函数使用详解和示例代码

torch.matmulPyTorch 中用于执行矩阵乘法的函数,它根据输入张量的维度自动选择适当的矩阵乘法方式,包括:

  • 向量内积(1D @ 1D)
  • 矩阵乘向量(2D @ 1D)
  • 向量乘矩阵(1D @ 2D)
  • 矩阵乘矩阵(2D @ 2D)
  • 批量矩阵乘法(>2D)

函数原型

torch.matmul(input
, other, *
, out=None
) → Tensor
  • input:第一个张量
  • other:第二个张量
  • out(可选):指定输出张量

详细说明

torch.matmul(a, b) 根据 ab 的维度规则如下:

a 维度b 维度操作类型
1D1D向量点积
2D1D矩阵和向量相乘
1D2D向量和矩阵相乘
2D2D标准矩阵乘法
≥3D≥3D批量矩阵乘法(batch)

示例代码

1. 向量点积(1D @ 1D)

import torch
a = torch.tensor([1.0
, 2.0
, 3.0]
)
b = torch.tensor([4.0
, 5.0
, 6.0]
)
result = torch.matmul(a, b)
print(result) # 输出:32.0

2. 矩阵乘向量(2D @ 1D)

a = torch.tensor([[1.0
, 2.0]
, [3.0
, 4.0]]
)
b = torch.tensor([5.0
, 6.0]
)
result = torch.matmul(a, b)
print(result) # 输出:[17.0, 39.0]

3. 向量乘矩阵(1D @ 2D)

a = torch.tensor([5.0
, 6.0]
)
b = torch.tensor([[1.0
, 2.0]
, [3.0
, 4.0]]
)
result = torch.matmul(a, b)
print(result) # 输出:[23.0, 34.0]

4. 矩阵乘矩阵(2D @ 2D)

a = torch.tensor([[1.0
, 2.0]
, [3.0
, 4.0]]
)
b = torch.tensor([[5.0
, 6.0]
, [7.0
, 8.0]]
)
result = torch.matmul(a, b)
print(result)
# 输出:
# [[19.0, 22.0],
# [43.0, 50.0]]

5. 批量矩阵乘法(3D @ 3D)

a = torch.randn(10
, 3
, 4
)
b = torch.randn(10
, 4
, 5
)
result = torch.matmul(a, b)
print(result.shape) # 输出:torch.Size([10, 3, 5])

综合示例:自定义线性层(类似 nn.Linear

下面是一个使用 torch.matmul 构建自定义线性层的完整示例,适合理解如何手动定义一个具有权重、偏置、支持自动求导的神经网络层,适合自定义网络结构或深入理解 PyTorch 的底层机制。

功能描述

  • 实现线性变换:y = x @ W^T + b
  • 使用 torch.matmul 执行矩阵乘法
  • 权重和偏置作为可训练参数
  • 支持 GPU 和自动求导

代码实现

import torch
import torch.nn as nn
class MyLinear(nn.Module):
def __init__(self, in_features, out_features):
super(MyLinear, self).__init__(
)
self.weight = nn.Parameter(torch.randn(out_features, in_features)
) # shape: [out, in]
self.bias = nn.Parameter(torch.zeros(out_features)
) # shape: [out]
def forward(self, x):
# x: shape [batch_size, in_features]
# weight: shape [out_features, in_features]
# transpose weight -> shape [in_features, out_features], then matmul
out = torch.matmul(x, self.weight.t(
)
) + self.bias
return out

使用示例

batch_size = 4
in_dim = 6
out_dim = 3
x = torch.randn(batch_size, in_dim)
layer = MyLinear(in_dim, out_dim)
output = layer(x)
print(output.shape) # torch.Size([4, 3])

与官方 nn.Linear 等效性验证(可选)

# 官方线性层
torch.manual_seed(0
)
official = nn.Linear(in_dim, out_dim)
# 自定义线性层,使用相同参数初始化
custom = MyLinear(in_dim, out_dim)
custom.weight.data.copy_(official.weight.data)
custom.bias.data.copy_(official.bias.data)
# 比较输出
x = torch.randn(2
, in_dim)
out1 = official(x)
out2 = custom(x)
print(torch.allclose(out1, out2)
) # True

说明

内容
torch.matmul用于实现 x @ W.T 矩阵乘法
nn.Parameter注册为可训练参数,自动加入 .parameters()
Module.forward()用于定义前向传播逻辑

注意事项

  • 输入张量必须满足矩阵乘法的维度匹配规则。
  • 对于 >2D 的张量,PyTorch 会自动按 batch size 广播执行多组矩阵乘法。
  • torch.matmul 不支持标量乘法(标量乘张量可用 * 运算符)。

posted on 2025-07-17 17:55  ljbguanli  阅读(228)  评论(0)    收藏  举报