PyTorch中matmul函数使用详解和示例代码

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

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

函数原型

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

详细说明

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

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

示例代码

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

python 复制代码
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)

python 复制代码
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)

python 复制代码
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)

python 复制代码
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)

python 复制代码
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 和自动求导

代码实现

python 复制代码
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

使用示例

python 复制代码
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 等效性验证(可选)

python 复制代码
# 官方线性层
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 不支持标量乘法(标量乘张量可用 * 运算符)。

相关推荐
白-胖-子1 小时前
深入剖析大模型在文本生成式 AI 产品架构中的核心地位
人工智能·架构
想要成为计算机高手2 小时前
11. isaacsim4.2教程-Transform 树与Odometry
人工智能·机器人·自动驾驶·ros·rviz·isaac sim·仿真环境
mortimer2 小时前
安装NVIDIA Parakeet时,我遇到的两个Pip“小插曲”
python·github
@昵称不存在3 小时前
Flask input 和datalist结合
后端·python·flask
静心问道3 小时前
InstructBLIP:通过指令微调迈向通用视觉-语言模型
人工智能·多模态·ai技术应用
宇称不守恒4.03 小时前
2025暑期—06神经网络-常见网络2
网络·人工智能·神经网络
赵英英俊3 小时前
Python day25
python
东林牧之3 小时前
Django+celery异步:拿来即用,可移植性高
后端·python·django
何双新3 小时前
基于Tornado的WebSocket实时聊天系统:从零到一构建与解析
python·websocket·tornado
小楓12014 小时前
醫護行業在未來會被AI淘汰嗎?
人工智能·醫療·護理·職業