blob: 183b34353cf20e39d0931a404feff90ad33457af (
plain)
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 WeightedMSECriterion, parent = torch.class('nn.WeightedMSECriterion','nn.MSECriterion')
function WeightedMSECriterion:__init(w)
parent.__init(self)
self.weight = w:clone()
self.buffer = torch.Tensor()
end
function WeightedMSECriterion:updateOutput(input,target)
self.buffer:resizeAs(input):copy(target)
if input:dim() - 1 == self.weight:dim() then
for i=1,input:size(1) do
self.buffer[i]:cmul(self.weight)
end
else
self.buffer:cmul(self.weight)
end
return input.nn.MSECriterion_updateOutput(self, input, self.buffer)
end
function WeightedMSECriterion:updateGradInput(input, target)
self.buffer:resizeAs(input):copy(target)
if input:dim() - 1 == self.weight:dim() then
for i=1,input:size(1) do
self.buffer[i]:cmul(self.weight)
end
else
self.buffer:cmul(self.weight)
end
return input.nn.MSECriterion_updateGradInput(self, input, self.buffer)
end
|