快速的方法來初始化torch7張量

Question

我要初始化與torch7的索引相關功能的三維張量，即

func = function(i,j,k) --i, j is the index of an element in the tensor 
    return i*j*k  --do operations within func which're dependent of i, j 
end 

然後我初始化一個三維張量是這樣的：

for i=1,A:size(1) do 
    for j=1,A:size(2) do 
     for k=1,A:size(3) do 
      A[{i,j,k}] = func(i,j,k) 
     end 
    end 
end 

但是這段代碼運行速度很慢，我發現它佔總運行時間的92％。在torch7中有沒有更有效的方法來初始化3D張量？

Answer 1

參見用於Tensor:apply

這些功能的文檔應用的函數，在 張量的每個元件，其被調用的方法（自）。這些方法比使用Lua中的for循環的 快得多。

文檔中的示例基於其索引i（在內存中）初始化2D數組。以下是三維的擴展示例，低於N維張量的示例。使用應用方法多，多更快我的機器上：

require 'torch' 

A = torch.Tensor(100, 100, 1000) 
B = torch.Tensor(100, 100, 1000) 

function func(i,j,k) 
    return i*j*k  
end 

t = os.clock() 
for i=1,A:size(1) do 
    for j=1,A:size(2) do 
     for k=1,A:size(3) do 
      A[{i, j, k}] = i * j * k 
     end 
    end 
end 
print("Original time:", os.difftime(os.clock(), t)) 

t = os.clock() 
function forindices(A, func) 
    local i = 1 
    local j = 1 
    local k = 0 
    local d3 = A:size(3) 
    local d2 = A:size(2) 
    return function() 
    k = k + 1 
    if k > d3 then 
     k = 1 
     j = j + 1 
     if j > d2 then 
     j = 1 
     i = i + 1 
     end 
    end 
    return func(i, j, k) 
    end 
end 

B:apply(forindices(A, func)) 
print("Apply method:", os.difftime(os.clock(), t)) 

---

編輯

這對於任何張量的目標工作：

function tabulate(A, f) 
    local idx = {} 
    local ndims = A:dim() 
    local dim = A:size() 
    idx[ndims] = 0 
    for i=1, (ndims - 1) do 
    idx[i] = 1 
    end 
    return A:apply(function() 
    for i=ndims, 0, -1 do 
     idx[i] = idx[i] + 1 
     if idx[i] <= dim[i] then 
     break 
     end 
     idx[i] = 1 
    end 
    return f(unpack(idx)) 
    end) 
end 

-- usage for 3D case. 
tabulate(A, function(i, j, k) return i * j * k end)