特征张量的批处理矩阵乘法

时间:2020-08-05 20:50:34

标签: c++ eigen eigen3

我想通过切片大张量来批量矩阵乘法。

说我的A的形状为[N,1,4],B的形状为[N,4,4]。我想首先沿批处理维度对它们进行切片,获得[b,1,4]和[b,4,4],它们不一定是连续的,而是通过分批进行矩阵乘法来获得形状[b,4]的结果。有没有办法使用Eigen做到这一点?

1 个答案:

答案 0 :(得分:0)

我不确定这是否是对本征张量执行批处理矩阵乘法的有效方法,但一种解决方案可能是将张量页面映射为矩阵并执行常规矩阵乘法:

#include <Eigen/Dense>
#include <unsupported/Eigen/CXX11/Tensor>

typedef Eigen::Tensor<double, 3> Tensor3d;

inline void batchedTensorMultiplication(const Tensor3d& A, const Tensor3d& B, const std::vector<int>& batchIndices, Tensor3d& C)
{
    Eigen::DenseIndex memStepA = A.dimension(0) * A.dimension(1);
    Eigen::DenseIndex memStepB = B.dimension(0) * B.dimension(1);
    Eigen::DenseIndex memStepC = C.dimension(0) * C.dimension(1);
    int outputBatchIndex = 0;

    for (int batchIndex : batchIndices)
    {
        Eigen::Map<const Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic>> pageA(A.data() + batchIndex * memStepA, A.dimension(0), A.dimension(1));
        Eigen::Map<const Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic>> pageB(B.data() + batchIndex * memStepB, B.dimension(0), B.dimension(1));
        Eigen::Map<Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic>> pageC(C.data() + outputBatchIndex * memStepC, C.dimension(0), C.dimension(1));

        outputBatchIndex++;

        pageC.noalias() = pageA * pageB;
    }
}

int main() 
{
    constexpr int N = 50;
    std::vector<int> batchIndices = { 0,1,2,3,4,9,10,11,12,13 };

    Tensor3d A(1, 4, N), B(4, 4, N), C(1, 4, (int)batchIndices.size());

    batchedTensorMultiplication(A, B, batchIndices, C);

    return 0;
}