пакетное умножение матриц для собственных тензоров - PullRequest
Новая
       88

пакетное умножение матриц для собственных тензоров

0 голосов
/ 05 августа 2020

Я хотел бы пакетное умножение матриц, взяв срезы больших тензоров.

Скажем, у меня есть A формы [N, 1, 4], B формы [N, 4, 4]. Я хотел бы сначала разрезать их по размерности партии, получая [b, 1, 4] и [b, 4, 4], которые не обязательно являются смежными, но получить результаты формы [b, 4], выполняя матричное умножение партиями. . Есть ли способ сделать это с помощью Eigen?

1 Ответ

0 голосов
/ 14 августа 2020

Я не уверен, что это эффективный способ выполнения пакетного умножения матриц для собственных тензоров, но одним из решений может быть отображение тензорных страниц как матриц и выполнение общего умножения матриц: 

#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;
}
...