джулия умножение двух массивов - PullRequest
Новая
       53

джулия умножение двух массивов

2 голосов
/ 25 февраля 2020

Есть ли способ более элегантно ускорить / записать это умножение массива (которое в numpy массивах я бы написал как A * B)? 

A = rand(8,15,10)
B = rand(10,5)
C = zeros(8,15,5)
for i in 1:8
    for j in 1:15
        for k in 1:10
            for l in 1:5
                C[i,j,l] = A[i,j,:]⋅B[:,l]
            end
        end
    end
end

1 Ответ

3 голосов
/ 26 февраля 2020

Существует множество пакетов Julia, которые позволяют вам записать ваше сокращение в одну простую строку. Вот несколько примеров, основанных на Einsum.jl , OMEinsum.jl и TensorOperations.jl : 

using OMEinsum
f_omeinsum(A,B) = ein"ijk,km->ijm"(A,B)

using Einsum
f_einsum(A,B) = @einsum C[i,j,l] := A[i,j,k] * B[k,l]

using TensorOperations
f_tensor(A,B) = @tensor C[i,j,l] := A[i,j,k] * B[k,l]

Помимо этих элегантных (и быстро, см. ниже) версии, вы можете немного улучшить свой код l oop. Вот ваш код, завернутый в функцию, и улучшенная версия с комментариями: 

function f(A,B)
    C = zeros(8,15,5)
    for i in 1:8
        for j in 1:15
            for k in 1:10
                for l in 1:5
                    C[i,j,l] = A[i,j,:]⋅B[:,l]
                end
            end
        end
    end
    return C
end

function f_fast(A,B)
    # check bounds
    n1,n2,n3 = size(A)
    m1, m2 = size(B)
    @assert m1 == n3
    C = zeros(n1,n2,m2)

    # * @inbounds to skip boundchecks inside the loop
    # * different order of the loops to account for Julia's column major order
    # * written out the k-loop (dot product) explicitly to avoid temporary allocations
    @inbounds for l in 1:m2
                for k in 1:m1
                        for j in 1:n2
                            for i in 1:n1
                            C[i,j,l] += A[i,j,k]*B[k,l]
                        end
                    end
                end
            end
    return C
end

Давайте сравним все подходы. Сначала мы проверим правильность: 

using Test
@test f(A,B) ≈ f_omeinsum(A,B) # Test passed
@test f(A,B) ≈ f_einsum(A,B) # Test passed
@test f(A,B) ≈ f_tensor(A,B) # Test passed
@test f(A,B) ≈ f_fast(A,B) # Test passed

Теперь давайте проверим, используя BenchmarkTools.jl . Я поместил время на мою машину в качестве комментариев. 

using BenchmarkTools
@btime f($A,$B); # 663.500 μs (12001 allocations: 1.84 MiB)
@btime f_omeinsum($A,$B); # 33.799 μs (242 allocations: 20.20 KiB)
@btime f_einsum($A,$B); # 4.200 μs (1 allocation: 4.81 KiB)
@btime f_tensor($A,$B); # 2.367 μs (3 allocations: 4.94 KiB)
@btime f_fast($A,$B); # 7.375 μs (1 allocation: 4.81 KiB)

Как мы видим, все подходы, основанные на einsum / тензорной записи, намного быстрее, чем ваша оригинальная реализация l oop - и только один лайнер! Производительность нашего f_fast находится на том же уровне, но все еще немного отстает от f_tensor, который является самым быстрым.

Наконец, давайте go все для производительности, потому что мы можем. Используя мастер dry из LoopVectorization.jl , мы заменим @inbounds в f_fast на @avx (мы называем эту новую версию f_avx ниже) и автоматически получим еще 2-кратное ускорение относительно с показателем f_tensor выше: 

@test f(A,B) ≈ f_avx(A,B) # Test passed
@btime f_avx($A,$B); # 930.769 ns (1 allocation: 4.81 KiB)

Однако из-за его простоты я все же предпочел бы f_tensor, если в вашем приложении не учитывается каждая микросекунда.

...