PyTorch torch.tensordot 函数
torch.tensordot 是 PyTorch 中用于计算两个张量在指定维度上的点积的函数。它是 Tensor Contraction(张量收缩)操作。
函数定义
torch.tensordot(input, other, dims=2)
参数说明:
input: 第一个输入张量other: 第二个输入张量dims: 要收缩的维度数量或维度对列表
使用示例
实例
import torch
# 创建两个一维张量
a = torch.tensor([1, 2, 3])
b = torch.tensor([4, 5, 6])
# 计算点积
y = torch.tensordot(a, b, dims=1)
print(y)
# 创建两个一维张量
a = torch.tensor([1, 2, 3])
b = torch.tensor([4, 5, 6])
# 计算点积
y = torch.tensordot(a, b, dims=1)
print(y)
输出结果为:
tensor(32)
实例
import torch
# 创建两个二维张量
a = torch.tensor([[1, 2], [3, 4]])
b = torch.tensor([[5, 6], [7, 8]])
# 计算矩阵乘法(收缩两个维度)
y = torch.tensordot(a, b, dims=2)
print(y)
# 创建两个二维张量
a = torch.tensor([[1, 2], [3, 4]])
b = torch.tensor([[5, 6], [7, 8]])
# 计算矩阵乘法(收缩两个维度)
y = torch.tensordot(a, b, dims=2)
print(y)
输出结果为:
tensor(70)
实例
import torch
# 创建三维张量
a = torch.randn(2, 3, 4)
b = torch.randn(3, 4, 5)
# 指定收缩的维度对
y = torch.tensordot(a, b, dims=[[1, 2], [0, 1]])
print(y.shape)
# 创建三维张量
a = torch.randn(2, 3, 4)
b = torch.randn(3, 4, 5)
# 指定收缩的维度对
y = torch.tensordot(a, b, dims=[[1, 2], [0, 1]])
print(y.shape)
输出结果为:
torch.Size([2, 5])

Pytorch torch 参考手册