在 torch7 中初始化张量的快捷方式

2015-5-31 14:16:44
收藏：0
阅读：109
评论：1

我需要在torch7中使用依赖于索引的函数来初始化一个3D张量，例如：

func = function(i,j,k)  -- i，j是张量元素的索引
    return i*j*k        -- 对于i，j执行函数内部操作
end

然后我像下面这样初始化一个3D张量A:

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%。有没有更有效的方法来初始化一个3D张量在torch7中？

点赞
用户2726734
用户2726734

查看Tensor:apply的文档。

这些函数会将它们被调用时（self）的 tensor 的每个元素应用到一个函数上。在Lua中，这些方法比使用for loop要快得多。

文档中的示例基于索引i（在内存中）初始化2D数组。以下是一个扩展示例，用于3个维度，下面是用于N-D张量的示例。在我的计算机上使用apply方法要快得多：

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))

EDIT

这适用于任何Tensor对象：

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

-- 用于3D情况的用法。
tabulate(A, function(i, j, k) return i * j * k end)
2015-05-31 17:51:03

作者:

用户1056274

用户1056274
推荐文章