main.diff

diff --git a/main.lua b/main.lua
index 5e827c0..6305a20 100644
--- a/main.lua
+++ b/main.lua
@@ -10,6 +10,7 @@ require 'torch'
 require 'paths'
 require 'optim'
 require 'nn'
+local json = require 'cjson'
 local DataLoader = require 'dataloader'
 local models = require 'models/init'
 local Trainer = require 'train'
@@ -41,6 +42,13 @@ if opt.testOnly then
    return
 end
 
+local logfile = io.open(paths.concat(opt.save_folder, 'log.txt'), 'w')
+
+local function log(t)
+   logfile:write('json_stats: '..json.encode(tablex.merge(t,opt,true))..'\n')
+   logfile:flush()
+end
+
 local startEpoch = checkpoint and checkpoint.epoch + 1 or opt.epochNumber
 local bestTop1 = math.huge
 local bestTop5 = math.huge
@@ -59,6 +67,15 @@ for epoch = startEpoch, opt.nEpochs do
       print(' * Best model ', testTop1, testTop5)
    end
 
+   log{
+      epoch = epoch,
+      trainTop1 = trainTop1,
+      trainTop5 = trainTop5,
+      testTop1  = testTop1,
+      testTop5 = testTop5,
+      trainLoss = trainLoss,
+   }
+
    checkpoints.save(epoch, model, trainer.optimState, bestModel, opt)
 end