Moved a part of the code to fftb.lua.
[dyncnn.git] / dyncnn.lua
index 6625579..0bb780c 100755 (executable)
@@ -29,43 +29,7 @@ require 'nn'
 require 'optim'
 require 'image'
 
-require 'img'
-
-----------------------------------------------------------------------
-
-function printf(f, ...)
-   print(string.format(f, unpack({...})))
-end
-
-colors = sys.COLORS
-
-function printfc(c, f, ...)
-   print(c .. string.format(f, unpack({...})) .. colors.black)
-end
-
-function logCommand(c)
-   print(colors.blue .. '[' .. c .. '] -> [' .. sys.execute(c) .. ']' .. colors.black)
-end
-
-----------------------------------------------------------------------
--- Environment variables
-
-local defaultNbThreads = 1
-local defaultUseGPU = false
-
-if os.getenv('TORCH_NB_THREADS') then
-   defaultNbThreads = os.getenv('TORCH_NB_THREADS')
-   print('Environment variable TORCH_NB_THREADS is set and equal to ' .. defaultNbThreads)
-else
-   print('Environment variable TORCH_NB_THREADS is not set, default is ' .. defaultNbThreads)
-end
-
-if os.getenv('TORCH_USE_GPU') then
-   defaultUseGPU = os.getenv('TORCH_USE_GPU') == 'yes'
-   print('Environment variable TORCH_USE_GPU is set and evaluated as ' .. tostring(defaultUseGPU))
-else
-   print('Environment variable TORCH_USE_GPU is not set, default is ' .. tostring(defaultUseGPU))
-end
+require 'fftb'
 
 ----------------------------------------------------------------------
 -- Command line arguments
@@ -77,6 +41,7 @@ cmd:text('General setup')
 cmd:option('-seed', 1, 'initial random seed')
 cmd:option('-nbThreads', defaultNbThreads, 'how many threads (environment variable TORCH_NB_THREADS)')
 cmd:option('-useGPU', defaultUseGPU, 'should we use cuda (environment variable TORCH_USE_GPU)')
+cmd:option('-fastGPU', true, 'should we go as fast as possible, possibly non-deterministically')
 
 cmd:text('')
 cmd:text('Log')
@@ -109,68 +74,21 @@ cmd:text('Problem to solve')
 
 cmd:option('-dataDir', './data/10p-mg', 'data directory')
 
-------------------------------
--- Log and stuff
-
 cmd:addTime('DYNCNN','%F %T')
 
 params = cmd:parse(arg)
 
-if params.rundir == '' then
-   params.rundir = cmd:string('exp', params, { })
-end
-
-paths.mkdir(params.rundir)
-
-if not params.noLog then
-   -- Append to the log if there is one
-   cmd:log(io.open(params.rundir .. '/log', 'a'), params)
-end
-
-----------------------------------------------------------------------
--- The experiment per se
-
-if params.predictGrasp then
-   params.targetDepth = 2
-else
-   params.targetDepth = 1
-end
-
 ----------------------------------------------------------------------
--- Initializations
-
-torch.setnumthreads(params.nbThreads)
-torch.setdefaulttensortype('torch.FloatTensor')
-torch.manualSeed(params.seed)
-
-----------------------------------------------------------------------
--- Dealing with the CPU/GPU
-
--- mynn will take entries in that order: mynn, cudnn, cunn, nn
-
-mynn = {}
 
-setmetatable(mynn,
-             {
-                __index = function(table, key)
-                   return (cudnn and cudnn[key]) or (cunn and cunn[key]) or nn[key]
-                end
-             }
-)
+fftbInit(cmd, params)
 
--- These are the tensors that can be kept on the CPU
-mynn.SlowTensor = torch.Tensor
-
--- These are the tensors that should be moved to the GPU
-mynn.FastTensor = torch.Tensor
-
-if params.useGPU then
-   require 'cutorch'
-   require 'cunn'
-   require 'cudnn'
-   cudnn.benchmark = true
-   cudnn.fastest = true
-   mynn.FastTensor = torch.CudaTensor
+for _, c in pairs({
+      'date',
+      'uname -a',
+      'git log -1 --format=%H'
+                 })
+do
+   logCommand(c)
 end
 
 ----------------------------------------------------------------------
@@ -185,8 +103,8 @@ function loadData(first, nb, name)
    data.width = 64
    data.height = 64
 
-   data.input = mynn.SlowTensor(data.nbSamples, 2, data.height, data.width)
-   data.target = mynn.SlowTensor(data.nbSamples, 1, data.height, data.width)
+   data.input = ffnn.SlowTensor(data.nbSamples, 2, data.height, data.width)
+   data.target = ffnn.SlowTensor(data.nbSamples, 1, data.height, data.width)
 
    for i = 1, data.nbSamples do
       local n = i-1 + first-1
@@ -226,8 +144,8 @@ function collectAllOutputs(model, collection, which)
 end
 
 function saveInternalsImage(model, data, n)
-   -- Explicitely copy to keep input as a mynn.FastTensor
-   local input = mynn.FastTensor(1, 2, data.height, data.width)
+   -- Explicitely copy to keep input as a ffnn.FastTensor
+   local input = ffnn.FastTensor(1, 2, data.height, data.width)
    input:copy(data.input:narrow(1, n, 1))
 
    local output = model:forward(input)
@@ -278,8 +196,8 @@ function saveResultImage(model, data, nbMax)
       criterion:cuda()
    end
 
-   local input = mynn.FastTensor(1, 2, data.height, data.width)
-   local target = mynn.FastTensor(1, 1, data.height, data.width)
+   local input = ffnn.FastTensor(1, 2, data.height, data.width)
+   local target = ffnn.FastTensor(1, 1, data.height, data.width)
 
    local nbMax = nbMax or 50
 
@@ -293,14 +211,14 @@ function saveResultImage(model, data, nbMax)
 
    for n = 1, nb do
 
-      -- Explicitely copy to keep input as a mynn.FastTensor
+      -- Explicitely copy to keep input as a ffnn.FastTensor
       input:copy(data.input:narrow(1, n, 1))
       target:copy(data.target:narrow(1, n, 1))
 
       local output = model:forward(input)
       local loss = criterion:forward(output, target)
 
-      output = mynn.SlowTensor(output:size()):copy(output)
+      output = ffnn.SlowTensor(output:size()):copy(output)
 
       -- We use our magical img.lua to create the result images
 
@@ -338,61 +256,60 @@ function createTower(filterSize, nbChannels, nbBlocks)
 
    else
 
-      tower = mynn.Sequential()
+      tower = ffnn.Sequential()
 
       for b = 1, nbBlocks do
-         local block = mynn.Sequential()
+         local block = ffnn.Sequential()
 
-         block:add(mynn.SpatialConvolution(nbChannels,
+         block:add(ffnn.SpatialConvolution(nbChannels,
                                            nbChannels,
                                            filterSize, filterSize,
                                            1, 1,
                                            (filterSize - 1) / 2, (filterSize - 1) / 2))
-         block:add(mynn.SpatialBatchNormalization(nbChannels))
-         block:add(mynn.ReLU(true))
+         block:add(ffnn.SpatialBatchNormalization(nbChannels))
+         block:add(ffnn.ReLU(true))
 
-         block:add(mynn.SpatialConvolution(nbChannels,
+         block:add(ffnn.SpatialConvolution(nbChannels,
                                            nbChannels,
                                            filterSize, filterSize,
                                            1, 1,
                                            (filterSize - 1) / 2, (filterSize - 1) / 2))
 
-         local parallel = mynn.ConcatTable()
-         parallel:add(block):add(mynn.Identity())
+         local parallel = ffnn.ConcatTable()
+         parallel:add(block):add(ffnn.Identity())
 
-         tower:add(parallel):add(mynn.CAddTable(true))
+         tower:add(parallel):add(ffnn.CAddTable(true))
 
-         tower:add(mynn.SpatialBatchNormalization(nbChannels))
-         tower:add(mynn.ReLU(true))
+         tower:add(ffnn.SpatialBatchNormalization(nbChannels))
+         tower:add(ffnn.ReLU(true))
       end
 
    end
 
    return tower
-
 end
 
 function createModel(imageWidth, imageHeight,
                      filterSize, nbChannels, nbBlocks)
 
-   local model = mynn.Sequential()
+   local model = ffnn.Sequential()
 
    -- Encode the two input channels (grasping image and starting
    -- configuration) into the internal number of channels
-   model:add(mynn.SpatialConvolution(2,
+   model:add(ffnn.SpatialConvolution(2,
                                      nbChannels,
                                      filterSize, filterSize,
                                      1, 1,
                                      (filterSize - 1) / 2, (filterSize - 1) / 2))
 
-   model:add(mynn.SpatialBatchNormalization(nbChannels))
-   model:add(mynn.ReLU(true))
+   model:add(ffnn.SpatialBatchNormalization(nbChannels))
+   model:add(ffnn.ReLU(true))
 
    -- Add the resnet modules
    model:add(createTower(filterSize, nbChannels, nbBlocks))
 
    -- Decode down to a single channel, which is the final image
-   model:add(mynn.SpatialConvolution(nbChannels,
+   model:add(ffnn.SpatialConvolution(nbChannels,
                                      1,
                                      filterSize, filterSize,
                                      1, 1,
@@ -403,42 +320,11 @@ end
 
 ----------------------------------------------------------------------
 
-function fillBatch(data, first, batch, permutation)
-   local actualBatchSize = math.min(params.batchSize, data.input:size(1) - first + 1)
-
-   if actualBatchSize ~= batch.input:size(1) then
-      local size = batch.input:size()
-      size[1] = actualBatchSize
-      batch.input:resize(size)
-   end
-
-   if actualBatchSize ~= batch.target:size(1) then
-      local size = batch.target:size()
-      size[1] = actualBatchSize
-      batch.target:resize(size)
-   end
-
-   for k = 1, batch.input:size(1) do
-      local i
-      if permutation then
-         i = permutation[first + k - 1]
-      else
-         i = first + k - 1
-      end
-      batch.input[k] = data.input[i]
-      batch.target[k] = data.target[i]
-   end
-end
-
 function trainModel(model, trainSet, validationSet)
 
    local criterion = nn.MSECriterion()
    local batchSize = params.batchSize
 
-   local batch = {}
-   batch.input = mynn.FastTensor(batchSize, 2, trainSet.height, trainSet.width)
-   batch.target = mynn.FastTensor(batchSize, 1, trainSet.height, trainSet.width)
-
    local startingEpoch = 1
 
    if model.epoch then
@@ -472,6 +358,8 @@ function trainModel(model, trainSet, validationSet)
       learningRateDecay = 0
    }
 
+   local batch = {}
+
    for e = startingEpoch, params.nbEpochs do
 
       model:training()
@@ -574,57 +462,39 @@ function trainModel(model, trainSet, validationSet)
 
 end
 
-function createAndTrainModel(trainSet, validationSet)
-
-   -- Load the current training state, or create a new model from
-   -- scratch
+----------------------------------------------------------------------
+-- main
 
-   if pcall(function () model = torch.load(params.rundir .. '/model_last.t7') end) then
+local trainSet = loadData(1,
+                          params.nbTrainSamples, 'train')
 
-      printfc(colors.red,
-              'Found a model with %d epochs completed, starting from there.',
-              model.epoch)
+local validationSet = loadData(params.nbTrainSamples + 1,
+                               params.nbValidationSamples, 'validation')
 
-      if params.exampleInternals ~= '' then
-         for _, i in ipairs(string.split(params.exampleInternals, ',')) do
-            saveInternalsImage(model, validationSet, tonumber(i))
-         end
-         os.exit(0)
-      end
+local model
 
-   else
+if pcall(function () model = torch.load(params.rundir .. '/model_last.t7') end) then
 
-      model = createModel(trainSet.width, trainSet.height,
-                          params.filterSize, params.nbChannels,
-                          params.nbBlocks)
+   printfc(colors.red,
+           'Found a model with %d epochs completed, starting from there.',
+           model.epoch)
 
+   if params.exampleInternals ~= '' then
+      for _, i in ipairs(string.split(params.exampleInternals, ',')) do
+         saveInternalsImage(model, validationSet, tonumber(i))
+      end
+      os.exit(0)
    end
 
-   trainModel(model, trainSet, validationSet)
-
-   return model
-
-end
+else
 
-----------------------------------------------------------------------
--- main
+   model = createModel(trainSet.width, trainSet.height,
+                       params.filterSize, params.nbChannels,
+                       params.nbBlocks)
 
-for _, c in pairs({
-      'date',
-      'uname -a',
-      'git log -1 --format=%H'
-                 })
-do
-   logCommand(c)
 end
 
-local trainSet = loadData(1,
-                          params.nbTrainSamples, 'train')
-
-local validationSet = loadData(params.nbTrainSamples + 1,
-                               params.nbValidationSamples, 'validation')
-
-local model = createAndTrainModel(trainSet, validationSet)
+trainModel(model, trainSet, validationSet)
 
 ----------------------------------------------------------------------
 -- Test