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

如果你真的想用,arrayfun你可以尝试这样的技巧:

arrayfun(@(i)trace(A(:,:,i)), 1:size(A,3))

但要注意的arrayfun是不是矢量化 !!,它只是一个循环的包装,往往比由于增加的开销循环慢.

虽然同样不是矢量化,但是你应该改变它的第二种方式cellfun.即如果A{i} = ones(3,3)那么

cellfun(@(x)trace(x), A)

Answer 2

bsxfun基于矢量化的解决方案,ab(使用)如何trace定义 - sum of diagonal elements-

%// 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

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

因此,完整的代码看起来像这样 -

[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倍.