A simple lena-denoising auto-encoder SpatialUnpooling test.

denoising_auto_encoder.lua

require 'nn'
require 'image'
require 'optim'
require 'cunn'
require 'SpatialUnpooling'

l = image.lena()

a = nn.SpatialConvolution(3,64,7,7)
b = nn.SpatialFullConvolution(64,3,7,7)


c = nn.SpatialConvolution(64,32,3,3)
d = nn.SpatialFullConvolution(32,64,3,3)


sp = nn.SpatialMaxPooling(2,2)
up = nn.SpatialUnpooling(2,2)
net = nn.Sequential()
net:add(nn.Dropout(0.9)) -- zero out randomly 90% of input
net:add(a)
net:add(nn.PReLU())
net:add(sp)
net:add(c)
net:add(nn.PReLU())
net:add(d)
net:add(up)
net:add(b)

mse = nn.MSECriterion()

net:cuda()
mse:cuda()
l=l:cuda()
up.indices = sp.indices

-- weight tying
d.weight = c.weight
d.gradWeight = c.gradWeight

b.weight = a.weight
b.gradWeight = a.gradWeight

function create_filter_image(m, kh, kw, filters)
  local w = m.weight --:view(3, filters, kh, kw)
  local img = m.weight.new(kh*3, kw*filters)
  for i=1,filters do
    local f = w[i]:clone()
    -- normalize for better inspection
    local lb,ub = f:min(), f:max()
    f:add(-lb):div(ub-lb+1e-10)
    for j=1,3 do -- input channels
      img[{{(j-1)*kh+1, j*kh}, {(i-1)*kw+1, i*kw}}] = f[j]     
    end
  end
  return img
end


w, dw = net:getParameters()
print(#w)

local optimization_target = function(w_)

  dw:zero()

  o = net:forward(l)
  err = mse:forward(o, l)
  go = mse:backward(o, l)
  net:backward(l, go)

  return err, dw
end

wb = net:get(2).weight:clone()
f = create_filter_image(net:get(2), 7, 7, 64)
image.save('weights1.png', f)

rmsprop_state = { learningRate = 0.0005, alpha = 0.9 }
for i=1,10000 do
  local _, loss = optim.rmsprop(optimization_target, w, rmsprop_state)
  print(string.format('%d: %f', i, loss[1]))
 
  if i%100 == 0 then
    image.save('weights2.png', create_filter_image(net:get(2), 7, 7, 64))
  end
end

--print(torch.abs(net:get(2).weight - wb))

image.save('test2.png', o)

SpatialUnpooling.lua

local SpatialUnpooling, parent = torch.class('nn.SpatialUnpooling', 'nn.Module')

function SpatialUnpooling:__init(kW, kH, dW, dH, padW, padH)
   parent.__init(self)
   self.dW = dW or kW
   self.dH = dH or kH
   self.padW = padW or 0
   self.padH = padH or 0
   self.indices = torch.LongTensor()
   self._indexTensor = torch.LongTensor()
end

function SpatialUnpooling:updateOutput(input)
  local n, d, h, w, oh, ow 
  if input:nDimension() == 4 then -- batch
    n, d, h, w = input:size(1), input:size(2), input:size(3), input:size(4)
    oh, ow = h * self.dH + 2 * self.padH, w * self.dW + 2 * self.padW
    self.output:resize(n, d, oh, ow)
  else
    n, d, h, w = 1, input:size(1), input:size(2), input:size(3)
    oh, ow = h * self.dH + 2 * self.padH, w * self.dW + 2 * self.padW
    self.output:resize(d, oh, ow)
  end

  local in_cols, out_cols, rows = h * w, oh * ow, n * d
  self.output:zero()
  self.output:view(rows, out_cols):scatter(
    2, 
    self.indices:view(rows, in_cols):typeAs(self._indexTensor), 
    input:view(rows, in_cols)
  )
  return self.output
end

function SpatialUnpooling:updateGradInput(input, gradOutput)
  self.gradInput:resizeAs(input)
  local n, d, h, w, oh, ow 
  if input:nDimension() == 4 then -- batch
    n, d, h, w, oh, ow = input:size(1), input:size(2), input:size(3), input:size(4), gradOutput:size(3), gradOutput:size(4)
  else
    n, d, h, w, oh, ow = 1, input:size(1), input:size(2), input:size(3), gradOutput:size(2), gradOutput:size(3)
  end

  local in_cols, out_cols, rows = h * w, oh * ow, n * d
  self.gradInput:view(rows, in_cols):gather(
    gradOutput:view(rows, out_cols),
    2, 
    self.indices:view(rows, in_cols):typeAs(self._indexTensor)
  )
  return self.gradInput
end

function SpatialUnpooling:type(type, tensorCache)
  parent.type(self, type, tensorCache)
  if type == 'torch.CudaTensor' then
    self._indexTensor:type(type)
  else
    self._indexTensor = torch.LongTensor()
  end
end
   
function SpatialUnpooling:__tostring__()
   return string.format('%s(%d,%d)', torch.type(self), self.kW, self.kH)
end

--[[
--simple test:
x = torch.rand(2, 1,10,10)
mp = nn.SpatialMaxPooling(2,2)
y = mp:forward(x)
up = nn.SpatialUnpooling(2,2)
up.indices = mp.indices
x_ = up:forward(y)
y_ = up:backward(y, x_)

-- cuda:
x = x:cuda()
up:cuda()
mp:cuda()
up.indices = mp.indices
y = mp:forward(x)
x_ = up:forward(y)
y_ = up:backward(y, x_)
]]