Spaces:
Build error
Build error
| local GDLCriterion, parent = torch.class('nn.GDLCriterion', 'nn.Criterion') | |
| function GDLCriterion:__init(alpha, sizeAverage) | |
| parent.__init(self) | |
| self.alpha = alpha or 1 | |
| self.sizeAverage = sizeAverage or true | |
| assert(self.alpha == 1 or self.alpha == 2, "alpha should be 1 or 2") | |
| local hNet = nn.Sequential() | |
| hNet:add(nn.ConcatTable():add(nn.Narrow(3, 1, -2)):add(nn.Narrow(3, 2, -1))) | |
| hNet:add(nn.CSubTable()):add(nn.Abs()) | |
| local wNet = nn.Sequential() | |
| wNet:add(nn.ConcatTable():add(nn.Narrow(4, 1, -2)):add(nn.Narrow(4, 2, -1))) | |
| wNet:add(nn.CSubTable()):add(nn.Abs()) | |
| self.inputNet = nn.ConcatTable():add(hNet):add(wNet) | |
| self.targetNet = self.inputNet:clone() | |
| self.criterion = {} | |
| if self.alpha == 1 then | |
| self.criterion[1] = nn.AbsCriterion(self.sizeAverage) | |
| self.criterion[2] = nn.AbsCriterion(self.sizeAverage) | |
| else | |
| self.criterion[1] = nn.MSECriterion(self.sizeAverage) | |
| self.criterion[2] = nn.MSECriterion(self.sizeAverage) | |
| end | |
| end | |
| function GDLCriterion:updateOutput(input, target) | |
| assert( input:nElement() == target:nElement(), | |
| "input and target size mismatch") | |
| self.inputNetOutput = self.inputNet:forward(input) | |
| self.targetNetOutput = self.targetNet:forward(target) | |
| self.output = self.criterion[1]:forward(self.inputNetOutput[1], self.targetNetOutput[1]) | |
| self.output = self.output + self.criterion[2]:forward(self.inputNetOutput[2], self.targetNetOutput[2]) | |
| return self.output | |
| end | |
| -- must have called updateOutput with the same input/target pair right before | |
| function GDLCriterion:updateGradInput(input, target) | |
| assert( input:nElement() == target:nElement(), | |
| "input and target size mismatch") | |
| local gradCriterion = {} | |
| gradCriterion[1] = self.criterion[1]:backward(self.inputNetOutput[1], self.targetNetOutput[1]) | |
| gradCriterion[2] = self.criterion[2]:backward(self.inputNetOutput[2], self.targetNetOutput[2]) | |
| self.gradInput = self.inputNet:backward(input, gradCriterion) | |
| return self.gradInput | |
| end | |