arrays - Matlab:带矩阵的 Arrayfun

Question

我正在尝试使用 arrayfun 重写以下代码

A = ones(3,3,3)
for i = 1:3
    B(i) = trace(A(:,:,i));
end

我希望我尝试过

f = @(x) trace(x)
B = arrayfun(f, A);

但这只是(如您所料)跟踪每个单独的 A(i,j,k) 而不是 A(:,:,i)，因为我想。然后我尝试将 A{i}=ones(3,3) 声明为一个单元格并传递给 arrayfun 但这也不起作用。

如何在 Matlab 中向量化矩阵函数？

Answer 1

bsxfun 基于矢量化解决方案，ab(使用)如何定义 trace - 对角线元素的总和 -

%// Get size of A
[m,n,r] = size(A) 

%// Get indices of the diagonal elements for each 3D "slice" as columns of idx
idx = bsxfun(@plus,[1:m+1:m*n]',[0:r-1]*m*n) %//'

%// Thus, for your 3 x 3 x 3 case, idx would be -
%//idx =
%//      1    10    19
%//      5    14    23
%//      9    18    27
%// and these are the linear indices to the diagonal elements to each `3D` slide.

%//Index into A with idx and sum along columns to get each element of desired output
B = sum(A(idx),1)

如果你想避免工作空间因不必要的额外变量而变得困惑，请避免使用 idx -

B = sum(A(bsxfun(@plus,[1:m+1:m*n]',[0:r-1]*m*n)),1)

---

用于使用 GPU

如果您必须使用 GPU，您可以使用 gpuArray(A) 将它们声明为 gpuArrays，然后涉及 A 的后续工作将在 GPU 上完成，您将获得作为 gpuArray 的输出，您可以使用 gather(..) 将其作为 CPU 变量返回。

因此，完整的代码如下所示 -

[m,n,r] = size(A); %// Get size
gpu_A = gpuArray(A); %// copy data from CPU to GPU

%// Perform calculations on GPU
gpu_B = sum(gpu_A(bsxfun(@plus,[1:m+1:m*n]',[0:r-1]*m*n)),1); %//'

B = gather(gpu_B); %// get back output onto CPU

快速测试:使用 GTX 750 Ti(我可以访问)，这似乎使我的循环代码速度提高了 3 倍。