Spaces:
Build error
Build error
| local CrossConvolveParallel, parent = torch.class('nn.CrossConvolveParallel', 'nn.Module') | |
| function CrossConvolveParallel:__init(nInputPlane, nOutputPlane, kW, kH) | |
| parent.__init(self) | |
| kH = kH or kW | |
| -- get args | |
| self.nInputPlane = nInputPlane or 1 | |
| self.nOutputPlane = nOutputPlane or 1 | |
| self.kW = kW | |
| self.kH = kH | |
| -- padding values | |
| self.padH = math.floor(kH / 2) | |
| self.padW = math.floor(kW / 2) | |
| self.convolver = {} | |
| self.gradInput = {} | |
| end | |
| function CrossConvolveParallel:updateOutput(input) | |
| assert(input[2]:size(5) == self.kH) | |
| assert(input[2]:size(4) == self.kW) | |
| -- Create convolver | |
| if #self.convolver < input[1]:size(1) then | |
| for i=1,input[1]:size(1) do | |
| local convolver = nn.Sequential() | |
| convolver:add(nn.SpatialZeroPadding(self.padW, self.padW, self.padH, self.padH)) | |
| convolver:add(nn.SpatialConvolution(self.nInputPlane, self.nOutputPlane, self.kW, self.kH)) | |
| -- set bias | |
| convolver.modules[2].bias:zero() | |
| if self:type() == 'torch.CudaTensor' then | |
| convolver:cuda() | |
| end | |
| -- Add to list | |
| table.insert(self.convolver, convolver) | |
| end | |
| end | |
| -- compute output | |
| local nSample = input[2]:size(1) | |
| for i=1,nSample do | |
| self.convolver[i].modules[2].weight = input[2][i] | |
| local tmp = self.convolver[i]:updateOutput(input[1][i]) | |
| if (self.output == nil) or (self.output:size():size()==0) or (self.output:size(1)~=nSample) then | |
| local h = tmp:size(2) | |
| local w = tmp:size(3) | |
| if self:type() == 'torch.CudaTensor' then | |
| self.output = torch.CudaTensor(nSample,self.nOutputPlane,h,w) | |
| else | |
| self.output = torch.DoubleTensor(nSample,self.nOutputPlane,h,w) | |
| end | |
| end | |
| self.output[i] = tmp | |
| end | |
| -- done | |
| return self.output | |
| end | |
| function CrossConvolveParallel:updateGradInput(input, gradOutput) | |
| -- resize grad | |
| for i = 1, #input do | |
| if self.gradInput[i] == nil then | |
| self.gradInput[i] = input[i].new() | |
| end | |
| self.gradInput[i]:resizeAs(input[i]):zero() | |
| end | |
| -- backprop | |
| for i = 1, input[1]:size(1) do | |
| self.convolver[i]:zeroGradParameters() | |
| self.gradInput[1][i]:add(self.convolver[i]:updateGradInput(input[1][i], gradOutput[i])) | |
| self.convolver[i]:accGradParameters(input[1][i], gradOutput[i]) | |
| self.gradInput[2][i]:add(self.convolver[i].modules[2].gradWeight) | |
| end | |
| -- done | |
| return self.gradInput | |
| end | |
| function CrossConvolveParallel:clearState() | |
| self.convolver:clearState() | |
| return parent.clearState(self) | |
| end | |
| function CrossConvolveParallel:__tostring__() | |
| return torch.type(self) .. string.format(' (%dx%dx%dx%d)', | |
| self.nInputPlane, self.nOutputPlane, self.kW, self.kH) | |
| end | |