Search code examples
luaneural-networktorch

How to add additional layers to a graph module in torch


How can I add new nodes to a graph module (gModule) in the nngraph package in torch? I tried to use the add function and this added the node to the modules slot in the gModules object. However the output is still taken from the previous last node.

Simplified code:

require "nn"
require "nngraph"

-- Function that builds a gModule
function buildModule(input_size,hidden_size)
    local x = nn.Identity()()
    local out = x - nn.Linear(input_size,hidden_size) - nn.Tanh()
    return nn.gModule({x},{out})
end

network = buildModule(5,3)
-- Additional layer to add
l2 = nn.Linear(3,10)
network:add(l2)

-- Expected a tensor of size 10 but got one with size 3
print(network:forward(torch.randn(5)))

Solution

  • gModule is not really supposed to be mutated. The fact it supports :add is actually a side-effect of being a child class of nn.Container, and rather not a design decision. In general once you create gModule you should not modify its internal structure, as you would have to modify some internal attributes to make verything work. Instead - if you want to add something "on top" just define new container that takes the previous one as an input.

    -- Function that builds a gModule
    function buildModule(input_size,hidden_size)
        local x = nn.Identity()()
        local out = x - nn.Linear(input_size,hidden_size) - nn.Tanh()
        return nn.gModule({x},{out})
    end
    
    network = buildModule(5,3)
    
    new_network = nn.Sequential()
    new_network:add(network)
    new_network:add(nn.Linear(3,10))