#!/usr/bin/env luajit

require 'torch'
require 'nn'

require 'profiler'

local model = nn.Sequential()
model:add(nn.Linear(1000, 1000))
model:add(nn.ReLU())
model:add(nn.Linear(1000, 100))

profiler.decor(model)

local input = torch.Tensor(1000, 1000)
local target = torch.Tensor(input:size(1), 100)
local criterion = nn.MSECriterion()

local nbSamples = 0
local modelTime = 0
local dataTime = 0

for k = 1, 5 do
   local t1 = sys.clock()
   input:uniform(-1, 1)
   target:uniform()

   local t2 = sys.clock()

   local output = model:forward(input)
   local loss = criterion:forward(output, target)
   local dloss = criterion:backward(output, target)
   model:backward(input, dloss)

   local t3 = sys.clock()

   dataTime = dataTime + (t2 - t1)
   modelTime = modelTime + (t3 - t2)

   nbSamples = nbSamples + input:size(1)
end

profiler.print(model, nbSamples)

print('----------------------------------------------------------------------')
print(string.format('Total model time %.02fs', modelTime))
print(string.format('Total data time %.02fs', dataTime))