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
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
|
local PixelSort, parent = torch.class("nn.PixelSort", "nn.Module")
-- Reverse pixel shuffle, based on the torch nn.PixelShuffle module (i'd attribute code, but not sure who wrote that)
-- Converts a [batch x channel x m x p] tensor to [batch x channel*r^2 x m/r x p/r]
-- tensor, where r is the downscaling factor.
-- Useful as an alternative to pooling & strided convolutions, as it doesn't discard information
-- if used with bottleneck convolution, you can discard half of the information, as opposed to 3/4 in pooling
-- also avoids the 'checkerboard' sampling issues found with strided convolutions.
-- @param downscaleFactor - the downscaling factor to use
function PixelSort:__init(downscaleFactor)
parent.__init(self)
self.downscaleFactor = downscaleFactor
self.downscaleFactorSquared = self.downscaleFactor * self.downscaleFactor
end
-- Computes the forward pass of the layer i.e. Converts a
-- [batch x channel x m x p] tensor to [batch x channel*r^2 x m/r x p/r] tensor.
-- @param input - the input tensor to be sorted of size [b x c x m x p]
-- @return output - the sorted tensor of size [b x c*r^2 x m/r x p/r]
function PixelSort:updateOutput(input)
self._intermediateShape = self._intermediateShape or torch.LongStorage(6)
self._outShape = self.outShape or torch.LongStorage()
self._shuffleOut = self._shuffleOut or input.new()
local batched = false
local batchSize = 1
local inputStartIdx = 1
local outShapeIdx = 1
if input:nDimension() == 4 then
batched = true
batchSize = input:size(1)
inputStartIdx = 2
outShapeIdx = 2
self._outShape:resize(4)
self._outShape[1] = batchSize
else
self._outShape:resize(3)
end
local channels = input:size(inputStartIdx)
local inHeight = input:size(inputStartIdx + 1)
local inWidth = input:size(inputStartIdx + 2)
self._intermediateShape[1] = batchSize
self._intermediateShape[2] = channels
self._intermediateShape[3] = inHeight / self.downscaleFactor
self._intermediateShape[4] = self.downscaleFactor
self._intermediateShape[5] = inWidth / self.downscaleFactor
self._intermediateShape[6] = self.downscaleFactor
self._outShape[outShapeIdx] = channels * self.downscaleFactorSquared
self._outShape[outShapeIdx + 1] = inHeight / self.downscaleFactor
self._outShape[outShapeIdx + 2] = inWidth / self.downscaleFactor
local inputView = torch.view(input, self._intermediateShape)
self._shuffleOut:resize(inputView:size(1), inputView:size(2), inputView:size(4),
inputView:size(6), inputView:size(3), inputView:size(5))
self._shuffleOut:copy(inputView:permute(1, 2, 4, 6, 3, 5))
self.output = torch.view(self._shuffleOut, self._outShape)
return self.output
end
-- Computes the backward pass of the layer, given the gradient w.r.t. the output
-- this function computes the gradient w.r.t. the input.
-- @param input - the input tensor of shape [b x c x m x p]
-- @param gradOutput - the tensor with the gradients w.r.t. output of shape [b x c*r^2 x m/r x p/r]
-- @return gradInput - a tensor of the same shape as input, representing the gradient w.r.t. input.
function PixelSort:updateGradInput(input, gradOutput)
self._intermediateShape = self._intermediateShape or torch.LongStorage(6)
self._shuffleIn = self._shuffleIn or input.new()
local batchSize = 1
local inputStartIdx = 1
if input:nDimension() == 4 then
batchSize = input:size(1)
inputStartIdx = 2
end
local channels = input:size(inputStartIdx)
local height = input:size(inputStartIdx + 1)
local width = input:size(inputStartIdx + 2)
self._intermediateShape[1] = batchSize
self._intermediateShape[2] = channels
self._intermediateShape[3] = self.downscaleFactor
self._intermediateShape[4] = self.downscaleFactor
self._intermediateShape[5] = height /self.downscaleFactor
self._intermediateShape[6] = width /self.downscaleFactor
local gradOutputView = torch.view(gradOutput, self._intermediateShape)
self._shuffleIn:resize(gradOutputView:size(1), gradOutputView:size(2), gradOutputView:size(5),
gradOutputView:size(4), gradOutputView:size(6), gradOutputView:size(3))
self._shuffleIn:copy(gradOutputView:permute(1, 2, 5, 3, 6, 4))
self.gradInput = torch.view(self._shuffleIn, input:size())
return self.gradInput
end
function PixelSort:clearState()
nn.utils.clear(self, {
"_intermediateShape",
"_outShape",
"_shuffleIn",
"_shuffleOut",
})
return parent.clearState(self)
end
|
