PyTorch 在行为方面等同于 Tensorflow 的 tf.keras.dot()?

PyTorch equivalent to Tensorflow's tf.keras.dot() in terms of behaviour?

在张量流中,如果你有 2 个形状分别为 NxTxD 和 NxDxT 的张量(N=batch_size,T=SequenceLength,D=NumberOfFeatures),你可以点它们并得到 NxTxT 的输出,如图所示下面:

import tensorflow as tf
import numpy as np

x1 = np.arange(2 * 4 * 3).reshape(2, 4, 3)
x2 = np.flip(np.arange(2 * 4 * 3).reshape(2, 3, 4), 1).copy()
print(x1.shape, x2.shape)
dotted = tf.keras.layers.Dot(axes=(2, 1))([x1, x2])
print(dotted.shape)
dotted
(2, 4, 3) (2, 3, 4)
(2, 4, 4)
<tf.Tensor: shape=(2, 4, 4), dtype=int32, numpy=
array([[[   4,    7,   10,   13],
        [  40,   52,   64,   76],
        [  76,   97,  118,  139],
        [ 112,  142,  172,  202]],

       [[ 616,  655,  694,  733],
        [ 760,  808,  856,  904],
        [ 904,  961, 1018, 1075],
        [1048, 1114, 1180, 1246]]])>

如果你尝试在 PyTorch 中做同样的事情,结果是不同的:

import torch
import numpy as np

x1 = torch.from_numpy(np.arange(2 * 4 * 3).reshape(2, 4, 3))
x2 = torch.from_numpy(np.flip(np.arange(2 * 4 * 3).reshape(2, 3, 4), 1).copy())
dotted = torch.tensordot(x1, x2, dims=([2], [1]))
print(x1.shape, x2.shape)
print(dotted.shape)
dotted
torch.Size([2, 4, 3]) torch.Size([2, 3, 4])
torch.Size([2, 4, 2, 4])
tensor([[[[   4,    7,   10,   13],
          [  40,   43,   46,   49]],

         [[  40,   52,   64,   76],
          [ 184,  196,  208,  220]],

         [[  76,   97,  118,  139],
          [ 328,  349,  370,  391]],

         [[ 112,  142,  172,  202],
          [ 472,  502,  532,  562]]],


        [[[ 148,  187,  226,  265],
          [ 616,  655,  694,  733]],

         [[ 184,  232,  280,  328],
          [ 760,  808,  856,  904]],

         [[ 220,  277,  334,  391],
          [ 904,  961, 1018, 1075]],

         [[ 256,  322,  388,  454],
          [1048, 1114, 1180, 1246]]]], dtype=torch.int32)

现在,Tensorflow 的结果存在于 pytorch 产生的结果中(它是它的一个子集)。 事实上,tensorflow 的结果基本上是某种更高维度的“对角线”。 PyTorch 的输出是 NxTxNxT,因此要获得与 Tensorflow 完全相同的结果,您可以这样做:

torch.stack([dotted[i, :, i, :] for i in range(len(dotted))])
tensor([[[   4,    7,   10,   13],
         [  40,   52,   64,   76],
         [  76,   97,  118,  139],
         [ 112,  142,  172,  202]],

        [[ 616,  655,  694,  733],
         [ 760,  808,  856,  904],
         [ 904,  961, 1018, 1075],
         [1048, 1114, 1180, 1246]]], dtype=torch.int32)

但这并不能否定你们俩的事实:

  1. 为 NxTxNxT 张量而不是 NxTxT 分配内存
  2. 计算量complexity/time急剧增加

有没有办法在不计算 4 维张量的情况下获得与 tensorflow 从 pytorch 给出的相同的 3 维结果?

我希望您正在寻找批量矩阵乘法 (bmm),它将两批矩阵相乘 - 两个张量必须是 3D。 https://pytorch.org/docs/stable/generated/torch.bmm.html