如何在Eigen中沿维度1拼接形状各异的[B,n,C]类型张量

需求说明 #

我们有一批形状为[B, n, C]的三维张量，其中批次大小B和通道数C是固定值，但每个张量的中间维度n可以不同（比如n ∈ [1,5]）。需要将所有这些张量沿中间维度（轴1）拼接成一个形状为[B, k, C]的单一张量，其中k是所有输入张量的n值之和。

NumPy参考实现 #

先看NumPy中的实现方式，方便对比理解：

import numpy as np

batch_size = 2
num_channels = 4

values = []

for i in range(0, 5):
  tensor = np.ones([batch_size, i+1, num_channels]) * i
  values.append(tensor)

print(np.concatenate(values, axis=1))

运行结果：

[[[0. 0. 0. 0.]
  [1. 1. 1. 1.]
  [1. 1. 1. 1.]
  [2. 2. 2. 2.]
  [2. 2. 2. 2.]
  [2. 2. 2. 2.]
  [3. 3. 3. 3.]
  [3. 3. 3. 3.]
  [3. 3. 3. 3.]
  [3. 3. 3. 3.]
  [4. 4. 4. 4.]
  [4. 4. 4. 4.]
  [4. 4. 4. 4.]
  [4. 4. 4. 4.]
  [4. 4. 4. 4.]]

 [[0. 0. 0. 0.]
  [1. 1. 1. 1.]
  [1. 1. 1. 1.]
  [2. 2. 2. 2.]
  [2. 2. 2. 2.]
  [2. 2. 2. 2.]
  [3. 3. 3. 3.]
  [3. 3. 3. 3.]
  [3. 3. 3. 3.]
  [3. 3. 3. 3.]
  [4. 4. 4. 4.]
  [4. 4. 4. 4.]
  [4. 4. 4. 4.]
  [4. 4. 4. 4.]
  [4. 4. 4. 4.]]]

Eigen实现方案 #

Eigen没有像NumPy那样直接的concatenate函数处理可变长度的三维张量拼接，需要手动完成以下步骤：

1. 计算总长度k：遍历所有输入张量，累加每个张量的中间维度大小，得到目标张量的中间维度长度k。
2. 初始化目标张量：创建一个形状为[B, k, C]的三维张量，提前分配好内存。
3. 逐张量拷贝数据：依次将每个输入张量的数据拷贝到目标张量的对应位置，通过控制起始索引来定位拼接区域。

完整代码示例 #

#include <Eigen/Dense>
#include <vector>
#include <iostream>

int main() {
    const int batch_size = 2;
    const int num_channels = 4;
    
    // 构造输入张量列表，模拟n从1到5的情况
    std::vector<Eigen::Tensor<double, 3>> values;
    for (int i = 0; i < 5; ++i) {
        int n = i + 1;
        // 创建形状为[batch_size, n, num_channels]的张量，填充为i
        Eigen::Tensor<double, 3> tensor(batch_size, n, num_channels);
        tensor.setConstant(i);
        values.push_back(tensor);
    }
    
    // 步骤1：计算总长度k
    int total_k = 0;
    for (const auto& tensor : values) {
        total_k += tensor.dimension(1);
    }
    
    // 步骤2：初始化目标张量
    Eigen::Tensor<double, 3> result(batch_size, total_k, num_channels);
    
    // 步骤3：逐张量拷贝数据
    int current_pos = 0;
    for (const auto& tensor : values) {
        int n = tensor.dimension(1);
        // 提取目标张量的对应区域，并拷贝数据
        result.chip(current_pos, 1).slice(Eigen::array<int, 3>{0, 0, 0}, 
                                         Eigen::array<int, 3>{batch_size, n, num_channels}) = tensor;
        current_pos += n;
    }
    
    // 打印结果（可选，验证正确性）
    for (int b = 0; b < batch_size; ++b) {
        std::cout << "Batch " << b << ":\n";
        for (int k = 0; k < total_k; ++k) {
            std::cout << "[";
            for (int c = 0; c < num_channels; ++c) {
                std::cout << result(b, k, c) << " ";
            }
            std::cout << "]\n";
        }
        std::cout << "\n";
    }
    
    return 0;
}

代码说明 #

• 使用Eigen::Tensor<double, 3>表示三维张量，对应NumPy中的三维数组。
• chip(current_pos, 1)用于获取目标张量在中间维度（轴1）从current_pos开始的切片，再通过slice限定拷贝的区域大小，和输入张量的尺寸匹配。
• setConstant(i)用于初始化张量值，和NumPy中的ones(...) * i效果一致。

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