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
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
|
local VolumetricCrossEntropyCriterion, parent = torch.class('cudnn.VolumetricCrossEntropyCriterion', 'nn.Criterion')
--[[
This criterion does the VolumetricCrossEntropyCriterion across
the feature dimension for a N-channel 3D image/video of TxHxW in size.
It only supports mini-batches (5D input, 4D target)
It does a LogSoftMax on the input (over the channel dimension),
so no LogSoftMax is needed in the network at the end
input = batchSize x nClasses x T x H x W
target = batchSize x T x H x W
]]--
function VolumetricCrossEntropyCriterion:__init(weights)
parent.__init(self)
self.scec = cudnn.SpatialCrossEntropyCriterion(weights)
end
local foldInput = function(input)
-- Fold time and height into one dimension
-- bdthw -> bd(t*h)w
input = input:view(input:size(1), input:size(2),
input:size(3)*input:size(4), input:size(5))
return input
end
local foldTarget = function(target)
-- Fold time and height into one dimension
-- bthw -> b(t*h)w
target = target:view(target:size(1), target:size(2)*target:size(3),
target:size(4))
return target
end
function VolumetricCrossEntropyCriterion:updateOutput(input, target)
assert(input:dim() == 5, 'mini-batch supported only')
assert(target:dim() == 4, 'mini-batch supported only')
assert(input:size(1) == target:size(1), 'input and target should be of same size')
assert(input:size(3) == target:size(2), 'input and target should be of same size')
assert(input:size(4) == target:size(3), 'input and target should be of same size')
assert(input:size(5) == target:size(4), 'input and target should be of same size')
-- Fold inputs and use spatial cross entropy criterion
self.scec:updateOutput(foldInput(input), foldTarget(target))
self.output = self.scec.output
return self.output
end
function VolumetricCrossEntropyCriterion:updateGradInput(input, target)
assert(input:dim() == 5, 'mini-batch supported only')
assert(target:dim() == 4, 'mini-batch supported only')
assert(input:size(1) == target:size(1), 'input and target should be of same size')
assert(input:size(3) == target:size(2), 'input and target should be of same size')
assert(input:size(4) == target:size(3), 'input and target should be of same size')
assert(input:size(5) == target:size(4), 'input and target should be of same size')
local originalInputSize = input:size()
self.scec:updateGradInput(foldInput(input), foldTarget(target))
self.gradInput = self.scec.gradInput:view(originalInputSize)
return self.gradInput
end
|
