优化不同大小的矩阵乘法

Question

假设我有以下数据生成过程

using Random
using StatsBase

m_1    = [1.0 2.0]
m_2    = [1.0 2.0; 3.0 4.0]
DD     = []
y      = zeros(2,200)

for i in 1:100
    rand!(m_1)
    rand!(m_2)
    push!(DD, m_1)
    push!(DD, m_2)
end

idxs   = sample(1:200,10)
for i in idxs
    DD[i] = DD[1]
end

假设给定数据，我有以下函数


function test(y, DD, n)
    v_1   = [1 2]
    v_2   = [3 4]
    for j in 1:n
        for i in 1:size(DD,1)
            if size(DD[i],1) == 1
                y[1:size(DD[i],1),i] .= (v_1 * DD[i]')[1]
            else
                y[1:size(DD[i],1),i] = (v_2 * DD[i]')'
            end
        end
    end
end

我正在努力优化 test 的速度。特别是，内存分配随着我的增加而增加 n。但是，我并没有真正分配任何新东西。

数据生成过程捕获了一个事实，即我事先不确定 DD[i] 的大小。也就是说，我第一次调用test时，DD[1]可能是一个2x2的矩阵。我第二次调用 test，DD[1] 可能是一个 1x2 矩阵。我认为这可能是内存分配问题的一部分：Julia 事先不知道大小。

我完全卡住了。我试过 @inbounds 但没有用。有什么办法可以改善吗？

Answer 1

检查性能的一件重要事情是 Julia 可以理解类型。您可以通过运行 @code_warntype test(y, DD, 1) 检查这一点，输出将清楚地表明 DD 是 Any[] 类型（因为您是这样声明的）。使用 Any 会导致相当大的性能损失，因此声明 DD = Matrix{Float64}[] 会将我的测试时间缩短到三分之一。

我不确定这个示例与您要编写的实际代码有多接近，但在这种特殊情况下，size(DD[i],1) == 1 分支可以替换为对 LinearAlgebra.dot 的调用：

y[1:size(DD[i],1),i] .= dot(v_1, DD[i])

这又为我节省了 50% 的时间。最后，您可以通过使用 mul! 执行其他乘法来多挤出一点点：

mul!(view(y, 1:size(DD[i],1),i:i), DD[i], v_2')

完整示例：

using Random
using LinearAlgebra

DD = [rand(i,2) for _ in 1:100 for i in 1:2]
y  = zeros(2,200)
shuffle!(DD)

function test(y, DD, n)
    v_1   = [1 2]
    v_2   = [3 4]'
    for j in 1:n
        for i in 1:size(DD,1)
            if size(DD[i],1) == 1
                y[1:size(DD[i],1),i] .= dot(v_1, DD[i])
            else
                mul!(view(y, 1:size(DD[i],1),i:i), DD[i], v_2)
            end
        end
    end
end

优化不同大小的矩阵乘法

Optimizing matrix multiplication with varying sizes

performance

matrix-multiplication

julia