Julia中的批量矩阵乘法答案

【问题标题】：Batch matrix multiplication in JuliaJulia中的批量矩阵乘法
【发布时间】：2019-08-27 16:29:27
【问题描述】：

我正在尝试将 Julia 中的 N 维 (N>=3) 数组作为矩阵批次相乘，即沿最后两个维度执行矩阵乘法，同时保持其他维度不变。

例如，如果x 的维度为(d1,d2,4,3)，y 的维度为(d1,d2,3,2)，则乘法的结果应该为(d1,d2,4,2)，即应该执行一批矩阵乘法。

这正是 Python 的 numpy.matmul 中发生的情况：

如果任一参数为 N-D，N > 2，则将其视为驻留在最后两个索引中的矩阵堆栈并相应地广播。

np.matmul(randn(10,10,4,3), randn(10,10,3,2)).shape
(10, 10, 4, 2)

有没有办法在 Julia 中重现 numpy.matmul 的行为？

我希望.* 可以工作，但是：

julia> randn(10,10,4,3) .* randn(10,10,3,2)
ERROR: DimensionMismatch("arrays could not be broadcast to a common size")
Stacktrace:
 [1] _bcs1 at ./broadcast.jl:485 [inlined]
 [2] _bcs at ./broadcast.jl:479 [inlined] (repeats 3 times)
 [3] broadcast_shape at ./broadcast.jl:473 [inlined]
 [4] combine_axes at ./broadcast.jl:468 [inlined]
 [5] instantiate at ./broadcast.jl:256 [inlined]
 [6] materialize(::Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{4},Nothing,typeof(*),Tuple{Array{Float64,4},Array{Float64,4}}}) at ./broadcast.jl:798
 [7] top-level scope at REPL[80]:1

我知道列表推导可能适用于 3-D，但在更高维度上会变得非常混乱。重塑（或查看）除最后两个维度之外的所有维度、使用列表理解并将其重新整形的最佳解决方案是什么？还是有更好的办法？

附：我能找到的最接近的东西是this，但并不完全相同。对 Julia 来说是新手，因此可能会遗漏一些对 Julia 用户来说显而易见的东西。

【问题讨论】：

标签： numpy julia matrix-multiplication

【解决方案1】：

我不知道有任何此类功能，但很可能在某些软件包中。我认为在 Julia 中，将数据组织为矩阵数组并在它们上广播矩阵乘法更为自然：

D = [rand(50, 60) for i in 1:4, j in 1:3]
E = [rand(60, 70) for i in 1:4, j in 1:3]
D .* E  # now you can use dot broadcasting!

也就是说，自己制作很容易。不过，我会做一个改变。 Julia 是主要列，而 numpy 是“最后一个主要维度”，因此您应该让矩阵位于 first 两个维度，而不是最后两个维度。

首先，我将定义一个与数组C 相乘的就地方法，然后定义一个调用就地版本的非就地方法（我将跳过维度检查等）：

# In-place version, note the use of the @views macro, 
# which is essential to get in-place behaviour

using LinearAlgebra: mul!  # fast in-place matrix multiply

function batchmul!(C, A, B)
    for j in axes(A, 4), i in axes(A, 3)
        @views mul!(C[:, :, i, j], A[:, :, i, j], B[:, :, i, j])
    end
    return C
end

# The non-in-place version
function batchmul(A, B)
    T = promote_type(eltype(A), eltype(B))
    C = Array{T}(undef, size(A, 1), size(B)[2:end]...)
    return batchmul!(C, A, B)
end

您也可以将其设为多线程。 ~~在我的计算机上，4 个线程提供了 2.5 倍的加速（实际上，对于最后两个维度的较大值，我得到了 3.5 倍的加速）~~你获得多少加速取决于涉及的数组：

function batchmul!(C, A, B)
    Threads.@threads for j in axes(A, 4)
        for i in axes(A, 3)
            @views mul!(C[:, :, i, j], A[:, :, i, j], B[:, :, i, j])
        end
    end
    return C
end

编辑：我刚才注意到您需要一般的 N-D，而不仅仅是 4-D。不过，不应该太难概括。无论如何，更有理由选择矩阵数组，其中广播将自动适用于所有维度。

Edit2： 不能离开它，所以这里有一个 N-D 案例（还有更多工作要做，比如处理基于非 1 的索引（更新：axes 应该解决这个问题） ):

function batchmul!(C, A, B)
    Threads.@threads for I in CartesianIndices(axes(A)[3:end])
        @views mul!(C[:, :, Tuple(I)...], A[:, :, Tuple(I)...], B[:, :, Tuple(I)...])
    end
    return C
end

【讨论】：

我更喜欢 N-D 数组而不是数组数组，因为它们更容易检查（例如 size 给了我一个不错的 N 元组），但在这种情况下，第一个解决方案似乎是最自然的。不过，肯定会保留通用的 N-D 解决方案以供参考，感谢您的出色回答！
对于“数组矩阵”解决方案，请记住，您可以使用 mul! 就地完成它：for i in eachindex(C) mul!(C[i], A[i], B[i]) end 不幸的是，点不会就地工作。

【解决方案2】：

对于 N=3，您正在寻找 NNlib.batched_mul。请注意（如上所述）Julia 的数组是列优先的，因此将最后一个不是第一个索引视为运行在批处理上通常是有意义的：

julia> using NNlib

julia> randn(4,3,100) ⊠ randn(3,2,100)
4×2×100 Array{Float64, 3}:
[:, :, 1] =
  0.9292     -0.223521
...

这只是一个类似batchmul!(C, A, B) 的循环，但它也会为 GPU CuArrays 调用适当的库函数。

将其扩展到超过 3 个维度并不难，但必须有人这样做，并决定规则。对于第 3 维，它的行为类似于广播：

julia> randn(4,3,100) ⊠ randn(3,2) |> size
(4, 2, 100)

julia> randn(4,3,100) ⊠ randn(3,2,1) |> size
(4, 2, 100)

julia> try randn(4,3,100) ⊠ randn(3,2,2) catch e println(e) end
DimensionMismatch("batch size mismatch: A != B")

【讨论】：