
--  Copyright (c) 2016, Facebook, Inc.
--  All rights reserved.
--
--  This source code is licensed under the BSD-style license found in the
--  LICENSE file in the root directory of this source tree. An additional grant
--  of patent rights can be found in the PATENTS file in the same directory.
--
--  The ResNet model definition
--

local depth =  34 -- 18           -- Imagenet: 18, 34, 50, 101, 152 | Cifar: 20, 32, 44, 56, 110, 1202
local shortcutType = 'B'   -- B/C
local dataset = 'imagenet' -- imagenet/cifar10
local iChannels

local nn = require 'nn'
require 'cunn'
require 'cudnn' -- jh : cudnn.version check 

local Convolution = nn.SpatialConvolution
local Avg = nn.SpatialAveragePooling
local ReLU = nn.ReLU
local Max = nn.SpatialMaxPooling
local SBatchNorm = nn.SpatialBatchNormalization

-- The shortcut layer is either identity or 1x1 convolution
local function shortcut(nInputPlane, nOutputPlane, stride)
   local useConv = shortcutType == 'C' or
   (shortcutType == 'B' and nInputPlane ~= nOutputPlane)
   if useConv then
      -- 1x1 convolution
      return nn.Sequential()
      :add(Convolution(nInputPlane, nOutputPlane, 1, 1, stride, stride))
      :add(SBatchNorm(nOutputPlane))
   elseif nInputPlane ~= nOutputPlane then
      -- Strided, zero-padded identity shortcut
      return nn.Sequential()
      :add(nn.SpatialAveragePooling(1, 1, stride, stride))
      :add(nn.Concat(2)
      :add(nn.Identity())
      :add(nn.MulConstant(0)))
   else
      return nn.Identity()
   end
end

-- The basic residual layer block for 18 and 34 layer network, and the
-- CIFAR networks
local function basicblock(n, stride)
   local nInputPlane = iChannels
   iChannels = n
	-- Changed to Preactivation!
   local s = nn.Sequential()
   s:add(SBatchNorm(n))
   s:add(ReLU(true))
   s:add(Convolution(nInputPlane,n,3,3,stride,stride,1,1))
   s:add(SBatchNorm(n))
   s:add(ReLU(true))
   s:add(Convolution(n,n,3,3,1,1,1,1))

   return nn.Sequential()
   :add(nn.ConcatTable()
   :add(s)
   :add(shortcut(nInputPlane, n, stride)))
   :add(nn.CAddTable(true))
   -- :add(ReLU(true))
end

-- The bottleneck residual layer for 50, 101, and 152 layer networks
local function bottleneck(n, stride)
   local nInputPlane = iChannels
   iChannels = n * 4

   local s = nn.Sequential()
   s:add(Convolution(nInputPlane,n,1,1,1,1,0,0))
   s:add(SBatchNorm(n))
   s:add(ReLU(true))
   s:add(Convolution(n,n,3,3,stride,stride,1,1))
   s:add(SBatchNorm(n))
   s:add(ReLU(true))
   s:add(Convolution(n,n*4,1,1,1,1,0,0))
   s:add(SBatchNorm(n * 4))

   return nn.Sequential()
   :add(nn.ConcatTable()
   :add(s)
   :add(shortcut(nInputPlane, n * 4, stride)))
   :add(nn.CAddTable(true)) -- removed ReLU for Preactivation
end

-- Creates count residual blocks with specified number of features
local function layer(block, features, count, stride)
   local s = nn.Sequential()
   for i=1,count do
      s:add(block(features, i == 1 and stride or 1))
   end
   return s
end

local model = nn.Sequential()
if dataset == 'imagenet' then
   -- Configurations for ResNet:
   --  num. residual blocks, num features, residual block function
   local cfg = {
      [18]  = {{2, 2, 2, 2}, 512, basicblock},
      -- [34]  = {{3, 4, 6, 3}, 512, basicblock},
      [34]  = {{5, 2, 1, 7}, 512, basicblock}, -- Deconv
      [50]  = {{3, 4, 6, 3}, 2048, bottleneck},
      [101] = {{3, 4, 23, 3}, 2048, bottleneck},
      [152] = {{3, 8, 36, 3}, 2048, bottleneck},
   }

   assert(cfg[depth], 'Invalid depth: ' .. tostring(depth))
   local def, nFeatures, block = table.unpack(cfg[depth])
   iChannels = 64
   print('ResNet-' .. depth .. ' ImageNet')

   -- The ResNet ImageNet model
   model:add(Convolution(3,64,9,9,1,1,4,4))
   -- model:add(Convolution(3,64,7,7,1,1,3,3))
   -- jh: first 3x3 .. model:add(Convolution(3,64,3,3,1,1,1,1))
   --
   --
   -- jh: off model:add(SBatchNorm(64))
   model:add(ReLU(true))
   -- model:add(Max(3,3,2,2,1,1))
   model:add(layer(block, 64, def[1]))
   model:add(layer(block, 64, def[2])) --, 2))
   model:add(layer(block, 64, def[3])) --, 2))
   model:add(layer(block, 64, def[4])) --, 2))
   -- model:add(Avg(7, 7, 1, 1))
   -- model:add(nn.View(nFeatures):setNumInputDims(3))
   -- model:add(nn.Linear(nFeatures, 1000))
   model:add( nn.SpatialFullConvolution(64,64, 4, 4, 2, 2, 1, 1, 0,0))
   model:add(ReLU(true))
   model:add( nn.SpatialFullConvolution(64,64, 4, 4, 2, 2, 1,1, 0,0))
   model:add(ReLU(true))

   -- Final conv 64 -> 3 color channel
   -- model:add(Convolution(64,3,3,3,1,1,1,1))
   -- model:add(Convolution(64,3,5,5,1,1,2,2))
   -- model:add(Convolution(64,3,7,7,1,1,3,3))
   model:add(Convolution(64,3,9,9,1,1,4,4))
elseif dataset == 'cifar10' then
   -- Model type specifies number of layers for CIFAR-10 model
   assert((depth - 2) % 6 == 0, 'depth should be one of 20, 32, 44, 56, 110, 1202')
   local n = (depth - 2) / 6
   iChannels = 16
   print('ResNet-' .. depth .. ' CIFAR-10')

   -- The ResNet CIFAR-10 model
   model:add(Convolution(3,16,3,3,1,1,1,1))
   model:add(SBatchNorm(16))
   model:add(ReLU(true))
   model:add(layer(basicblock, 16, n))
   model:add(layer(basicblock, 32, n, 2))
   model:add(layer(basicblock, 64, n, 2))
   model:add(Avg(8, 8, 1, 1))
   model:add(nn.View(64):setNumInputDims(3))
   model:add(nn.Linear(64, 10))
else
   error('invalid dataset: ' .. dataset)
end

local function ConvInit(name)
   for k,v in pairs(model:findModules(name)) do
      local n = v.kW*v.kH*v.nOutputPlane
      v.weight:normal(0,math.sqrt(2/n))
      if cudnn.version >= 4000 then
         v.bias = nil
         v.gradBias = nil
      else
         v.bias:zero()
      end
   end
end
local function BNInit(name)
   for k,v in pairs(model:findModules(name)) do
      v.weight:fill(1)
      v.bias:zero()
   end
end

ConvInit('nn.SpatialConvolution')
ConvInit('nn.SpatialConvolution')
BNInit('fbnn.SpatialBatchNormalization')
BNInit('nn.SpatialBatchNormalization')
BNInit('nn.SpatialBatchNormalization')
for k,v in pairs(model:findModules('nn.Linear')) do
   v.bias:zero()
end
model:cuda()
 
  -- if opt.cudnn == 'deterministic' then
     model:apply(function(m)
        if m.setMode then m:setMode(1,1,1) end
     end)
  -- end

model:get(1).gradInput = nil

return model
