burn.lua

require 'torch'
require 'nn'
require 'sys'

local cmd = torch.CmdLine()

-- Basic options
cmd:option('-image_size', 224, 'Maximum height / width of generated image')
cmd:option('-gpu', 0, 'Zero-indexed ID of the GPU to use; for CPU mode set -gpu = -1')

-- Other options
cmd:option('-backend', 'cudnn', 'nn|cl|cudnn')
cmd:option('-model', 'A')
cmd:option('-batch_size', 128)
cmd:option('-sync', 0)

local function create_model(modelType, imageSize)
   -- from the convnet_benchmarks scripts
   -- Create tables describing VGG configurations A, B, D, E
   local cfg = {}
   if modelType == 'A' then
      cfg = {64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'}
   elseif modelType == 'B' then
      cfg = {64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'}
   elseif modelType == 'D' then
      cfg = {64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'}
   elseif modelType == 'E' then
      cfg = {64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M'}
--   elseif modelType == 'hugh1' then
--      cfg = {8, 8, 'M', 16, 16, 'M', 32, 32, 'M', 64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M'}
   elseif modelType == 'unit' then
      cfg = {'M', 'M', 'M', 'M', 32, 'M'}
   else
      error('Unknown model type: ' .. modelType .. ' | Please specify a modelType A or B or D or E')
   end

   local model = nn.Sequential()
   local size = imageSize
   local numLayers = nil
   do
      local iChannels = 3;
      for k,v in ipairs(cfg) do
         if v == 'M' then
            model:add(nn.SpatialMaxPooling(2,2,2,2):ceil())
--            size = math.ceil(size / 4) * 2
            size = math.ceil(size / 2)
         else
            local oChannels = v;
            model:add(nn.SpatialConvolutionMM(iChannels,oChannels,3,3,1,1,1,1))
            numLayers = oChannels
            model:add(nn.ReLU(true))
            iChannels = oChannels;
         end
      end
   end

--   features:get(1).gradInput = nil

--   local classifier = nn.Sequential()
   model:add(nn.View(numLayers * size * size))
   model:add(nn.Linear(numLayers * size * size, 4096))
   model:add(nn.ReLU())
   model:add(nn.Dropout(0.5))
--   model:add(nn.Threshold(0, 1e-6))
   -- classifier:add(nn.Dropout(0.5))
   model:add(nn.Linear(4096, 4096))
   model:add(nn.ReLU())
   model:add(nn.Dropout(0.5))
--   model:add(nn.Threshold(0, 1e-6))
   -- classifier:add(nn.Dropout(0.5))
   model:add(nn.Linear(4096, 1000))
   -- classifier:add(nn.LogSoftMax())

--   local model = nn.Sequential()
--   model:add(features):add(classifier)
   return model
end

local function sync(params)
  if params.backend ~= 'cl' then
    cutorch.synchronize()
  else
    cltorch.synchronize()
  end
end

local function main(params)
  if params.gpu >= 0 then
    if params.backend == 'nn' then
      require 'cutorch'
      require 'cunn'
      cutorch.setDevice(params.gpu + 1)
    elseif params.backend == 'cudnn' then
      require 'cutorch'
      require 'cunn'
      require 'cudnn'
      cutorch.setDevice(params.gpu + 1)
      cudnn.fastest = True
    elseif params.backend == 'cl' then
      require 'cltorch'
      require 'clnn'
      cltorch.setDevice(params.gpu + 1)
    end
  else
    error('params.gpu should be more than 0')
--    params.backend = 'nn-cpu'
  end

  local cnn = create_model(params.model, params.image_size)

  local img = torch.FloatTensor(params.batch_size, 3, 224, 224)
  if params.gpu >= 0 then
    if params.backend == 'nn' then
      cnn:cuda()
      img = img:cuda()
    elseif params.backend == 'cudnn' then
--      cnn:cuda()
      cnn:cuda()
      cudnn.convert(cnn, cudnn)
      img = img:cuda()
    elseif params.backend == 'cl' then
      cnn:cl()
      img = img:cl()
    else
      error('unknown backend ' .. params.backend)
    end
  else
    error('params.gpu should be more than 0')
  end

  print('cnn', cnn)

  print('torch.type(img)', torch.type(img))
  print('img:size()', img:size())

--  local output = cnn:forward(img)

  local input = img
  local currentOutput = input
--  for warmup=1,3 do
  while true do
    sys.tic()
    cnn:forward(img)    
    sync(params)
    local thisTime = sys.toc()
    print('time', thisTime)
  end
--  print('forward')
--  local totalTime = 0
--  for i=1,#cnn.modules do
--    local module = cnn.modules[i]
--    print('module', module)
--    local inputDims = currentOutput:dim()
--    local inputSize = currentOutput:size(inputDims)
--    if params.sync then
--      sync(params)
--    end
--    sys.tic()
--    currentOutput = module:updateOutput(currentOutput)
--    if params.sync then
--      sync(params)
--    end
--    local thisTime = sys.toc()
--    totalTime = totalTime + thisTime
--    local outputDims = currentOutput:dim()
--    local outputSize = currentOutput:size(outputDims)
--    local moduleName = tostring(module):gsub(' ', '')
--    print('{"layer": ' .. i .. ', "time_us": ' .. math.floor(thisTime*1e6 + 0.5) ..
--       ', "module": "' .. moduleName .. '", "inputsize": ' .. inputSize .. ', "outputsize": ' .. outputSize .. '}')
--  end
--  local output = currentOutput
--  print('total time: ' .. math.floor(totalTime*1e6 + 0.5))

--  print('backward')
--  local gradInput = cnn:backward(img, output)
end
  
local params = cmd:parse(arg)
main(params)