init.lua 2.65 KB
Newer Older
Alexis Brenon's avatar
Alexis Brenon committed
1
--- A module to build neural network from a @{NetworkDescription|description}.
Alexis Brenon's avatar
Alexis Brenon committed
2
--
3
-- An example of description can be found in this file: @{Inference.lua}
4
-- @module network
5
-- @alias module
6
7
-- @author Alexis BRENON <alexis.brenon@imag.fr>

8
local module = require('arcades.utils.package_loader')(...)
Alexis Brenon's avatar
Alexis Brenon committed
9
10
11

local torch = require('torch')
local nn = require('nn')
12

Alexis Brenon's avatar
Alexis Brenon committed
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
--- Data Types
-- @section data-types

--- Description of a neural networks.
-- @tfield table input_size Size of the input @{torch.Tensor/tensor.md/|Tensor} `{d, w, h}`
-- @tfield {ConvolutionDescription,...} conv_layers List of @{nn.Convolution/convolution.md#nn.SpatialModules/|convolution} layers
-- @tfield {number,...} fc_layers List of width of @{nn.Linear/simple.md#nn.Linear/} fully connected layers
-- @tfield string nl_layer Name of the @{nn.Transfert/transfer.md/|transfert function} to use
-- @tfield table output_size Size of the output @{torch.Tensor/tensor.md/|Tensor} `{w, h}`
-- @table NetworkDescription

--- Description of a convolutionnal layer.
-- @tfield int n_filters Number of filters to use
-- @tfield DimenTable field_size Size of the receptive field
-- @tfield DimenTable stride Size of the strides
-- @tfield DimenTable zero_padding Size of the paddings
-- @table ConvolutionDescription

--- Description of a 2D dimension.
-- @tfield int width
-- @tfield int height
-- @table DimenTable

--- Public Methods
-- @section public-methods

--- Build a network.
-- @tparam NetworkDescription args Description of the network to build
-- @treturn nn.Sequential  Built @{nn.Module/containers.md#nn.Sequential/|network}
42
43
44
45
46
47
48
49
50
51
52
53
54
55
function module.create_network(args)

    local net = nn.Sequential()
    net:add(nn.View(unpack(args.input_size)))

    -- Convolutionnal layers
    for i, layer in ipairs(args.conv_layers) do
      net:add(nn.SpatialConvolution(
        (args.conv_layers[i-1] and args.conv_layers[i-1].n_filters) or args.input_size[1],
        layer.n_filters,
        layer.field_size.width, layer.field_size.height,
        layer.stride.width, layer.stride.height,
        layer.zero_padding.width, layer.zero_padding.height
      ))
Alexis Brenon's avatar
Alexis Brenon committed
56
      net:add(nn[args.nl_layer](true))
57
58
59
60
61
62
63
64
65
    end

    -- Convert multidimensionnal output to 1D input
    local nelements = net:forward(torch.zeros(1,unpack(args.input_size))):nElement()
    net:add(nn.View(nelements))

    -- Hidden Linear layers
    for _, layer_size in ipairs(args.fc_layers) do
      net:add(nn.Linear(nelements, layer_size))
Alexis Brenon's avatar
Alexis Brenon committed
66
      net:add(nn[args.nl_layer](true))
67
68
69
70
71
72
73
74
75
76
      nelements = layer_size
    end

    -- Output layer
    net:add(nn.Linear(nelements, torch.Tensor(args.output_size):prod()))
    net:add(nn.View(unpack(args.output_size)))
    return net
end

return module