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 NOT VECTORIZATION !!，它只是循环的包装器，由于增加了开销，通常比循环慢。

虽然同样不是矢量化，但要按照第二种方式进行操作，您应该更改为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，可以将它们声明为带有gpuArray(A)的gpuArrays，然后在GPU上完成涉及A的后续工作，您将获得输出作为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倍。

Matlab：带矩阵的Arrayfun

2 个答案: