如何在PyTorch中实现批量点积运算？

这个需求完全可以实现，核心是利用广播机制或者爱因斯坦求和约定实现对应维度的匹配运算，最终可以精准得到尺寸为[B, N]的输出结果。

运算逻辑说明 #

你的两个张量维度含义如下：

• 输入张量[B, N, 3]：可理解为B组数据，每组包含N个长度为3的向量
• 测试张量[N, 3]：包含N个长度为3的向量
  我们需要让输入张量中每组的第i个3维向量，和测试张量的第i个3维向量执行点积，最终输出的每个[b, i]位置的元素就是对应点积的计算结果，维度刚好匹配[B, N]的要求。

不同工具的实现示例 #

• NumPy实现：
  方法1：使用爱因斯坦求和

  import numpy as np
  output = np.einsum('bnd,nd->bn', input_tensor, test_tensor)
  

  方法2：广播乘法+维度求和

  output = (input_tensor * test_tensor[np.newaxis, :, :]).sum(axis=-1)
  

• PyTorch实现：

  import torch
  # 方法1：爱因斯坦求和
  output = torch.einsum('bnd,nd->bn', input_tensor, test_tensor)
  # 方法2：广播+求和
  output = torch.sum(input_tensor * test_tensor.unsqueeze(0), dim=-1)
  

• TensorFlow实现：

  import tensorflow as tf
  output = tf.reduce_sum(input_tensor * test_tensor[tf.newaxis, ...], axis=-1)
  

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