如何在PyTorch中实现与TensorFlow tf.keras.dot等效的点积运算

要实现输入形状为N×T×D和N×D×T的张量，输出和tf.keras.layers.Dot完全一致的N×T×T结果，且无冗余计算，最优方案是使用PyTorch内置的批量矩阵乘法算子torch.bmm。

实现原理 #

torch.bmm是专门针对3维批量张量设计的矩阵乘法接口，运算时会将张量的第一维作为batch维度，仅对同batch内的后两维做矩阵乘法，不会产生跨batch的冗余计算，内存占用和计算效率与TensorFlow的Dot层完全对齐。

代码示例 #

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.bmm(x1, x2)

print(x1.shape, x2.shape)
print(dotted.shape)
print(dotted)

输出结果 #

torch.Size([2, 4, 3]) torch.Size([2, 3, 4])
torch.Size([2, 4, 4])
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)

方案优势 #

• 内存占用仅为tensordot+切片方案的1/N，仅分配目标N×T×T大小的内存，无冗余内存开销
• 没有跨batch的无效计算，耗时和TensorFlow原生Dot层基本一致
• 代码简洁，符合PyTorch官方最佳实践

通用封装（适配任意axes参数） #

如果需要和tf.keras.layers.Dot一样支持自定义点积轴，可以封装为通用函数：

def tf_equiv_dot(x1: torch.Tensor, x2: torch.Tensor, axes: tuple[int, int]) -> torch.Tensor:
    """
    实现和tf.keras.layers.Dot完全等价的点积运算
    :param x1: 输入张量1
    :param x2: 输入张量2
    :param axes: 长度为2的元组，分别指定x1和x2上要做点积的轴
    :return: 和TensorFlow输出一致的点积结果
    """
    # 调整x1的维度顺序，将待点积的轴放到最后一位
    x1_perm = list(range(x1.ndim))
    x1_perm.append(x1_perm.pop(axes[0]))
    x1_reshaped = x1.permute(x1_perm)
    
    # 调整x2的维度顺序，将待点积的轴放到倒数第二位
    x2_perm = list(range(x2.ndim))
    x2_perm.insert(-1, x2_perm.pop(axes[1]))
    x2_reshaped = x2.permute(x2_perm)
    
    return torch.bmm(x1_reshaped, x2_reshaped)

调用方式和TensorFlow完全一致：

dotted = tf_equiv_dot(x1, x2, axes=(2, 1))

内容的提问来源于stack exchange，提问作者N1h1l1sT