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 不是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，然后后續涉及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倍。