Блог

Я использую Julia 1.0. Пожалуйста, рассмотрите следующий код:

 using LinearAlgebra
using Distributions

## create random data
const data = rand(Uniform(-1,2), 100000, 2)

function test_function_1(data)
    theta = [1 2]
    coefs = theta * data[:,1:2]'
    res   = coefs' .* data[:,1:2]
    return sum(res, dims = 1)'
end

function test_function_2(data)
    theta   = [1 2]
    sum_all = zeros(2)
    for i = 1:size(data)[1]
        sum_all .= sum_all   (theta * data[i,1:2])[1] *  data[i,1:2]
    end
    return sum_all
end
  

После первого запуска я рассчитал время

 julia> @time test_function_1(data)
  0.006292 seconds (16 allocations: 5.341 MiB)
2×1 Adjoint{Float64,Array{Float64,2}}:
 150958.47189289227
 225224.0374366073

julia> @time test_function_2(data)
  0.038112 seconds (500.00 k allocations: 45.777 MiB, 15.61% gc time)
2-element Array{Float64,1}:
 150958.4718928927
 225224.03743660534
  

test_function_1 значительно превосходит как по распределению, так и по скорости, но test_function_1 не девекторизован. Я ожидал test_function_2 лучшей производительности. Обратите внимание, что обе функции выполняют одно и то же.

У меня есть подозрение, что это из-за того, что test_function_2 я использую sum_all .= sum_all ... , но я не уверен, почему это проблема. Могу я получить подсказку?

Итак, сначала позвольте мне прокомментировать, как бы я написал вашу функцию, если бы хотел использовать цикл:

 function test_function_3(data)
    theta   = (1, 2)
    sum_all = zeros(2)
    for row in eachrow(data)
        sum_all . = dot(theta, row) .*  row
    end
    return sum_all
end
  

Далее, здесь приведено сравнительное сравнение трех вариантов:

 julia> @benchmark test_function_1($data)
BenchmarkTools.Trial: 
  memory estimate:  5.34 MiB
  allocs estimate:  16
  --------------
  minimum time:     1.953 ms (0.00% GC)
  median time:      1.986 ms (0.00% GC)
  mean time:        2.122 ms (2.29% GC)
  maximum time:     4.347 ms (8.00% GC)
  --------------
  samples:          2356
  evals/sample:     1

julia> @benchmark test_function_2($data)
BenchmarkTools.Trial: 
  memory estimate:  45.78 MiB
  allocs estimate:  500002
  --------------
  minimum time:     16.316 ms (7.44% GC)
  median time:      16.597 ms (7.63% GC)
  mean time:        16.845 ms (8.01% GC)
  maximum time:     34.050 ms (4.45% GC)
  --------------
  samples:          297
  evals/sample:     1

julia> @benchmark test_function_3($data)
BenchmarkTools.Trial: 
  memory estimate:  96 bytes
  allocs estimate:  1
  --------------
  minimum time:     777.204 μs (0.00% GC)
  median time:      791.458 μs (0.00% GC)
  mean time:        799.505 μs (0.00% GC)
  maximum time:     1.262 ms (0.00% GC)
  --------------
  samples:          6253
  evals/sample:     1
  

Далее вы можете работать немного быстрее, если явно реализуете dot в цикле:

 julia> function test_function_4(data)
           theta   = (1, 2)
           sum_all = zeros(2)
           for row in eachrow(data)
               @inbounds sum_all . = (theta[1]*row[1] theta[2]*row[2]) .*  row
           end
           return sum_all
       end
test_function_4 (generic function with 1 method)

julia> @benchmark test_function_4($data)
BenchmarkTools.Trial: 
  memory estimate:  96 bytes
  allocs estimate:  1
  --------------
  minimum time:     502.367 μs (0.00% GC)
  median time:      502.547 μs (0.00% GC)
  mean time:        505.446 μs (0.00% GC)
  maximum time:     806.631 μs (0.00% GC)
  --------------
  samples:          9888
  evals/sample:     1
  

Чтобы понять различия, давайте взглянем на эту строку вашего кода:

 sum_all .= sum_all   (theta * data[i,1:2])[1] *  data[i,1:2]
  

Давайте посчитаем, сколько памяти вы выделяете в этом выражении:

 sum_all .= 
    sum_all
      # allocation of a new vector as a result of addition
    (theta
     *  # allocation of a new vector as a result of multiplication
     data[i,1:2] # allocation of a new vector via getindex
    )[1]
    * # allocation of a new vector as a result of multiplication
    data[i,1:2] # allocation of a new vector via getindex
  

Итак, вы можете видеть, что на каждой итерации цикла вы выделяете пять раз.
Выделения стоят дорого. И вы можете видеть это в тестах, что у вас есть 5000002 выделения в процессе:

• 1 выделение sum_all
• 1 выделение theta
• 500000 выделений в цикле (5 * 100000)

Кроме того, вы выполняете индексацию, подобную data[i,1:2] , которая выполняет проверку границ, что также требует небольших затрат (но незначительных по сравнению с распределениями).

Теперь в функции test_function_3 я использую eachrow(data) . На этот раз я также получаю строки data матрицы, но они возвращаются как представления (не новые матрицы), поэтому внутри цикла выделение не происходит. Затем я снова использую dot функцию, чтобы избежать выделения, которое ранее было вызвано умножением матрицы (я изменил theta на Tuple из a, Matrix поскольку тогда dot это немного быстрее, но это вторично). Наконец я пишу, um_all . = dot(theta, row) .* row и в этом случае все операции транслируются, поэтому Julia может выполнять широковещательное объединение (опять же — никаких распределений не происходит).

В test_function_4 я просто заменяю dot на развернутый цикл, поскольку мы знаем, что у нас есть два элемента для вычисления скалярного произведения. На самом деле, если вы полностью развернете все и будете использовать @simd , это станет еще быстрее:

 julia> function test_function_5(data)
          theta   = (1, 2)
          s1 = 0.0
          s2 = 0.0
          @inbounds @simd for i in axes(data, 1)
               r1 = data[i, 1]
               r2 = data[i, 2]
               mul = theta[1]*r1   theta[2]*r2
               s1  = mul * r1
               s2  = mul * r2
          end
          return [s1, s2]
       end
test_function_5 (generic function with 1 method)

julia> @benchmark test_function_5($data)
BenchmarkTools.Trial: 
  memory estimate:  96 bytes
  allocs estimate:  1
  --------------
  minimum time:     22.721 μs (0.00% GC)
  median time:      23.146 μs (0.00% GC)
  mean time:        24.306 μs (0.00% GC)
  maximum time:     100.109 μs (0.00% GC)
  --------------
  samples:          10000
  evals/sample:     1
  

Итак, вы можете видеть, что таким образом вы примерно в 100 раз быстрее, чем с test_function_1 . Все еще уже test_function_3 выполняется относительно быстро и является полностью универсальным, поэтому, вероятно, обычно я бы написал что-то вроде test_function_3 , если только мне действительно не нужно было работать очень быстро и я знал, что размеры моих данных фиксированы и малы.