PyTorch torch.tensordot 関数

Pytorch torch 参考手册Pytorch torch リファレンスマニュアル

torch.tensordotPyTorch において、2つのテンソルの指定した次元におけるドット積を計算するための関数です。これは Tensor Contraction（テンソル縮約）操作です。

関数定義

torch.tensordot(input, other, dims=2)

パラメータの説明：

  • input: 最初の入力テンソル
  • other: 2番目の入力テンソル
  • dims: 縮約する次元の数、または次元ペアのリスト

使用例

インスタンス

import torch

# 2つの1次元テンソルを作成
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

# 2つの2次元テンソルを作成
a = torch.tensor([[1, 2], [3, 4]])
b = torch.tensor([[5, 6], [7, 8]])

# 行列乗算を計算（2つの次元を縮約）
y = torch.tensordot(a, b, dims=2)
print(y)

出力結果は次のとおりです：

tensor(70)

インスタンス

import torch

# 3次元テンソルを作成
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 参考手册Pytorch torch リファレンスマニュアル

その他の拡張