nn.LogSoftMax
具体来说,当输入张量的大小很小时,我不想使用 omp 。我有一个小脚本来测试运行时间。
require 'nn'
my_lsm = function(t)
o = torch.zeros((#t)[1])
sum = 0.0
for i = 1,(#t)[1] do
o[i] = torch.exp(t[i])
sum = sum + o[i]
end
o = o / sum
return torch.log(o)
end
ii=torch.randn(arg[1])
m=nn.LogSoftMax()
timer = torch.Timer()
timer:stop()
timer:reset()
timer:resume()
my_lsm(ii)
print(timer:time().real)
timer:stop()
timer:reset()
timer:resume()
m:forward(ii)
print(timer:time().real)
如果arg[1]
是 10,那么我的基本 log softmax 函数运行得更快:
0.00021696090698242
0.033425092697144
但是一次arg[1]
是 10,000,000,omp 确实有很大帮助:
29.561321973801
0.11547803878784
所以我怀疑 omp 开销非常高。如果我的代码必须使用少量输入多次调用 log softmax(例如张量大小仅为 3),则将花费太多时间。有没有办法在某些情况下(但并非总是)手动禁用 omp 使用?