Search code examples
luaparallel-processingdeep-learninglayertorch

Make first layer of AlexNet parallel in lua


I am trying to have parallel table for the first layer of Alexnet in Torch, Lua. I want to pass two Batches of RGB images to the network and then send the addition of them to the next layer. For example: suppose that I want to send images with 6 channels to the first layer of alexnet but in this case, I want to send two batches of 3 channels each to the first parallel layers , join them and then send the output to the next layer. the actual code is like this:

net = nn.Sequential()
net:add(nn.SpatialConvolution(3,96,11,11,4,4,2,2))       
net:add(nn.SpatialBatchNormalization(96))
net:add(nn.ReLU(true))
net:add(nn.SpatialMaxPooling(3,3,2,2))                   
net:add(nn.SpatialConvolution(96,256,5,5,1,1,2,2))       
net:add(nn.SpatialBatchNormalization(256))
net:add(nn.ReLU(true))
net:add(nn.SpatialMaxPooling(3,3,2,2))                   
net:add(nn.SpatialConvolution(256,384,3,3,1,1,1,1))     

and the code I thought it would work is :

net = nn.Sequential()
c = nn.ParallelTable()
c:add(nn.SpatialConvolution(3,48,11,11,4,4,2,2)) 
c:add(nn.SpatialConvolution(3,48,11,11,4,4,2,2))
net:add(c)
net:add(nn.JoinTable(1,8))
net:add(nn.SpatialBatchNormalization(96))
net:add(nn.ReLU(true))
net:add(nn.SpatialMaxPooling(3,3,2,2)) 

and the error I got is :

In 1 module of nn.ParallelTable: /torch/install/share/lua/5.1/cudnn/init.lua:171: assertion failed!

I was wondering where I am going wrong with this implementation and any help would be much appreciated.

Thanks


Solution

  • Your error is on the line net:add(nn.JoinTable(1,8)), it should be 3 instead of 8. This value is the number of dimensions (without counting the batch dimension as a dimension) of your input tensor. Here you feed your network with 3D images, then you should write net:add(nn.JoinTable(1,3))

    I used the following code and all goes well

    require 'nn'
    require 'cutorch'
    require 'cunn'
    require 'cudnn'
    
    net = nn.Sequential()
    c = nn.ParallelTable()
    c:add(nn.SpatialConvolution(3,48,11,11,4,4,2,2)) 
    c:add(nn.SpatialConvolution(3,48,11,11,4,4,2,2))
    net:add(c)
    net:add(nn.JoinTable(1,3))
    net:add(nn.SpatialBatchNormalization(96))
    net:add(nn.ReLU(true))
    net:add(nn.SpatialMaxPooling(3,3,2,2)) 
    
    net:cuda()
    
    input1 = torch.rand(128,3,227, 227):cuda()
    input2 = torch.rand(128,3,227,227):cuda()
    
    out = net:forward({input1, input2})
    
    print(out:size())