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

如果您正在使用矩阵，将x 视为矩阵向量而不是 3D 数组可能是合适的。那你就可以了

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 倍的加速，运行时间为 12us。

Answer 2

mapslices(i->i*y, x, (1,2)) 可能更“干净”，但会更慢。

读作：将函数“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 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

所以使用sum 的tst3 具有更好的性能（大约是tst 的7 倍和大约是tst2 的15 倍）。

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