PyTorch torch.tensordot Función

Pytorch torch 参考手册Pytorch torch Manual de referencia

torch.tensordotEs una función en PyTorch utilizada para calcular el producto punto de dos tensores en dimensiones especificadas. Es una operación de contracción de tensores (Tensor Contraction).

Definición de la función

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

Descripción de parámetros:

  • input: El primer tensor de entrada
  • other: El segundo tensor de entrada
  • dims: El número de dimensiones a contraer o una lista de pares de dimensiones

Ejemplo de uso

Ejemplo

import torch

# Crear dos tensores unidimensionales
a = torch.tensor([1, 2, 3])
b = torch.tensor([4, 5, 6])

# Calcular el producto punto
y = torch.tensordot(a, b, dims=1)
print(y)

El resultado de salida es:

tensor(32)

Ejemplo

import torch

# Crear dos tensores bidimensionales
a = torch.tensor([[1, 2], [3, 4]])
b = torch.tensor([[5, 6], [7, 8]])

# Calcular la multiplicación de matrices (contraer dos dimensiones)
y = torch.tensordot(a, b, dims=2)
print(y)

El resultado de salida es:

tensor(70)

Ejemplo

import torch

# Crear un tensor tridimensional
a = torch.randn(2, 3, 4)
b = torch.randn(3, 4, 5)

# Especificar los pares de dimensiones a contraer
y = torch.tensordot(a, b, dims=[[1, 2], [0, 1]])
print(y.shape)

El resultado de salida es:

torch.Size([2, 5])

Pytorch torch 参考手册Pytorch torch Manual de referencia

Otras extensiones