Julia 将每个矩阵乘以 dim

Question

我有一个 3 维数组

x = rand(6,6,2^10)

我想将沿三维的每个矩阵乘以一个向量。有没有比以下更干净的方法：

y = rand(6,1)
z = zeros(6,1,2^10)
for i in 1:2^10
    z[:,:,i] = x[:,:,i] * y
end

Answer 1

mapslices(i->i*y, x, (1,2)) 可能是 "cleaner" 但它会更慢。

读作：将函数"times by y"应用于前两个维度的每个切片。

function tst(x,y)
   z = zeros(6,1,2^10)
   for i in 1:2^10
       z[:,:,i] = x[:,:,i] * y
   end
   return z
end

tst2(x,y) = mapslices(i->i*y, x, (1,2))

time tst(x,y); 0.002152 秒（4.10 k 分配：624.266 KB）

@time tst2(x,y); 0.005720 秒（13.36 k 分配：466.969 KB）

Answer 2

如果您使用的是矩阵，则将 x 视为矩阵向量而不是 3 维数组可能比较合适。那么你可以做

x = [rand(6,6) for _ in 1:2^10]
y = [rand(6)]
z = x .* y

z 现在是向量的向量。

如果 z 是预分配的，那就是

z .= x .* y

而且，如果你想要它非常快，使用 StaticArrays

的向量

using StaticArrays

x = [@SMatrix rand(6, 6) for _ in 1:2^10]
y = [@SVector rand(6)]
z = x .* y

在我的计算机上显示了 10 倍的加速，运行在 12 微秒内。

Answer 3

sum(x.*y',2) 是一个干净的简短解决方案。

它还具有良好的速度和内存特性。诀窍是将矩阵-向量乘法视为按向量元素缩放的矩阵列的线性组合。我们没有对矩阵 x[:,:,i] 进行每个线性组合，而是对 x[:,i,:] 使用相同的尺度 y[i]。在代码中：

const x = rand(6,6,2^10);
const y = rand(6,1);
function tst(x,y)
    z = zeros(6,1,2^10)
    for i in 1:2^10
        z[:,:,i] = x[:,:,i]*y
    end
    return z
end
tst2(x,y) = mapslices(i->i*y,x,(1,2))
tst3(x,y) = sum(x.*y',2)

基准测试得出：

julia> using BenchmarkTools
julia> z = tst(x,y); z2 = tst2(x,y); z3 = tst3(x,y);
julia> @benchmark tst(x,y)
  BenchmarkTools.Trial: 
    memory estimate:  688.11 KiB
    allocs estimate:  8196
    --------------
    median time:      759.545 μs (0.00% GC)
    samples:          6068
julia> @benchmark tst2(x,y)
  BenchmarkTools.Trial: 
    memory estimate:  426.81 KiB
    allocs estimate:  10798
    --------------
    median time:      1.634 ms (0.00% GC)
    samples:          2869
julia> @benchmark tst3(x,y)
  BenchmarkTools.Trial: 
    memory estimate:  336.41 KiB
    allocs estimate:  12
    --------------
    median time:      114.060 μs (0.00% GC)
    samples:          10000

因此 tst3 使用 sum 具有更好的性能（比 tst 高出约 7 倍，比 tst2 高出约 15 倍）。

按照@DNF 的建议使用 StaticArrays 也是一种选择，最好将其与此处的解决方案进行比较。

Julia 将每个矩阵乘以 dim

Julia multiply each matrix along dim

matrix-multiplication

julia