BCECriterion.lua


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31

local BCECriterion, parent = torch.class('nn.BCECriterion', 'nn.Criterion')

function BCECriterion:__init()
   parent.__init(self)
   self.sizeAverage = true
end

function BCECriterion:updateOutput(input, target)
    -- log(input) * target + log(1 - input) * (1 - target)
    self.output = torch.log(input):cmul(target)
    
    self.output:add(torch.add(-input,1):log():cmul(torch.add(-target,1)))

    if self.sizeAverage then
    	self.output:div(target:size(1))
    end

    return self.output:sum()
end

function BCECriterion:updateGradInput(input, target)
    -- target / input - (1 - target) / (1 - input)
    self.gradInput = torch.cdiv(target,input)
    self.gradInput:add(-1,torch.cdiv(torch.add(-target,1),torch.add(-input,1)))

    if self.sizeAverage then
    	self.gradInput:div(target:size(1))
    end
    
    return self.gradInput
end