矩阵相乘torch.einsum()

einsumEinstein summation 的缩写,来源于爱因斯坦求和约定(Einstein summation convention)。这是物理学家阿尔伯特·爱因斯坦引入的一种简便记号,用于描述张量运算,特别是涉及多维数组的运算。

示例1:矩阵乘法

矩阵乘法 C=AB

python 复制代码
A = torch.randn(2, 3)
B = torch.randn(3, 4)
C = torch.einsum('ik,kj->ij', A, B)
print(C.size())  # 输出: torch.Size([2, 4])

这里,'ik,kj->ij' 的含义是:

  • A 的形状为 (2, 3),对应 ikik 分别表示第一个和第二个维度。
  • B 的形状为 (3, 4),对应 kjkj 分别表示第一个和第二个维度。
  • ->ij 表示输出张量的模式,结果为 (2, 4)
示例2:向量点积

向量点积 c=a⋅b

python 复制代码
a = torch.randn(3)
b = torch.randn(3)
c = torch.einsum('i,i->', a, b)
print(c.size())  # 输出: torch.Size([])

这里,'i,i->' 的含义是:

  • ab 都是向量,对应模式 i
  • -> 后面为空,表示结果是一个标量。
示例3:批量矩阵乘法

批量矩阵乘法

python 复制代码
A = torch.randn(10, 2, 3)
B = torch.randn(10, 3, 4)
C = torch.einsum('bij,bjk->bik', A, B)
print(C.size())  # 输出: torch.Size([10, 2, 4])

这里,'bij,bjk->bik' 的含义是:

  • A 的形状为 (10, 2, 3),对应 bijb 表示批次维度,ij 分别表示矩阵的行和列。
  • B 的形状为 (10, 3, 4),对应 bjkb 表示批次维度,jk 分别表示矩阵的行和列。
  • ->bik 表示输出张量的模式,结果为 (10, 2, 4)

示例4:逐元素相乘(哈达玛积)A.B或A × B

python 复制代码
A = torch.randn(3, 4)
B = torch.randn(3, 4)

C = torch.einsum('ij,ij->ij', A, B)
print(C.size())  # 输出: torch.Size([3, 4])

'ij,ij->ij' 表示:

  • AB 都是形状为 [3, 4] 的矩阵,用 ij 表示。
  • 结果 C 也是形状为 [3, 4] 的矩阵。
  • 没有重复索引,所以不进行求和。
相关推荐
AC赳赳老秦1 天前
代码生成超越 GPT-4:DeepSeek-V4 编程任务实战与 2026 开发者效率提升指南
数据库·数据仓库·人工智能·科技·rabbitmq·memcache·deepseek
液态不合群1 天前
推荐算法中的位置消偏,如何解决?
人工智能·机器学习·推荐算法
饭饭大王6661 天前
当 AI 系统开始“自省”——在 `ops-transformer` 中嵌入元认知能力
人工智能·深度学习·transformer
ujainu1 天前
CANN仓库中的AIGC可移植性工程:昇腾AI软件栈如何实现“一次开发,多端部署”的跨生态兼容
人工智能·aigc
初恋叫萱萱1 天前
CANN 生态实战指南:从零构建一个高性能边缘 AI 应用的完整流程
人工智能
Lethehong1 天前
CANN ops-nn仓库深度解读:AIGC时代的神经网络算子优化实践
人工智能·神经网络·aigc
lbb 小魔仙1 天前
面向 NPU 的高性能矩阵乘法:CANN ops-nn 算子库架构与优化技术
线性代数·矩阵·架构
开开心心就好1 天前
AI人声伴奏分离工具,离线提取伴奏K歌用
java·linux·开发语言·网络·人工智能·电脑·blender
TechWJ1 天前
CANN ops-nn神经网络算子库技术剖析:NPU加速的基石
人工智能·深度学习·神经网络·cann·ops-nn
凌杰1 天前
AI 学习笔记:LLM 的部署与测试
人工智能