1
0
Fork 0
mirror of synced 2024-05-17 03:12:18 +12:00

Add ScaleTable for Squeeze and Excitation Networks

This commit is contained in:
nagadomi 2018-10-18 19:35:16 +00:00
parent 0d51c40f6b
commit eea286059f
2 changed files with 35 additions and 0 deletions

34
lib/ScaleTable.lua Normal file
View file

@ -0,0 +1,34 @@
local ScaleTable, parent = torch.class("w2nn.ScaleTable", "nn.Module")
function ScaleTable:__init()
parent.__init(self)
self.gradInput = {}
self.grad_tmp = torch.Tensor()
self.scale = torch.Tensor()
end
function ScaleTable:updateOutput(input)
assert(#input == 2)
assert(input[1]:size(2) == input[2]:size(2))
self.scale:resizeAs(input[1]):expandAs(input[2], input[1])
self.output:resizeAs(self.scale):copy(self.scale)
self.output:cmul(input[1])
return self.output
end
function ScaleTable:updateGradInput(input, gradOutput)
self.gradInput[1] = self.gradInput[1] or input[1].new()
self.gradInput[1]:resizeAs(input[1]):copy(gradOutput)
self.gradInput[1]:cmul(self.scale)
self.grad_tmp:resizeAs(input[1]):copy(gradOutput)
self.grad_tmp:cmul(input[1])
self.gradInput[2] = self.gradInput[2] or input[2].new()
self.gradInput[2]:resizeAs(input[2]):sum(self.grad_tmp:reshape(self.grad_tmp:size(1), self.grad_tmp:size(2), self.grad_tmp:size(3) * self.grad_tmp:size(4)), 3):resizeAs(input[2])
for i=#input+1, #self.gradInput do
self.gradInput[i] = nil
end
return self.gradInput
end

View file

@ -84,6 +84,7 @@ else
require 'RandomBinaryConvolution'
require 'RandomBinaryCriterion'
require 'EdgeFilter'
require 'ScaleTable'
return w2nn
end