Function for creating a neural network with n hidden layers in Flux

Hi,

I want to create a function, where I can set the number of hidden layers with an argument. One option that kinda works is

using Flux
function make_nn(length_in = 80, length_out = 1, nodes = 128, hid_lay = 3, act_fun = relu, act_fun_last = sigmoid)
	q = Dense(length_in, nodes, act_fun)
	for i in 1:hid_lay-1
		q = Chain(q, Dense(nodes, nodes, act_fun))
	end
	q = Chain(q, Dense(nodes, length_out, act_fun_last))
end

but the output is a nested Chain which is not particular beautiful to me:

Chain(
  Chain(
    Chain(
      Dense(80, 128, relu),             # 10_368 parameters
      Dense(128, 128, relu),            # 16_512 parameters
    ),
    Dense(128, 128, relu),              # 16_512 parameters
  ),
  Dense(128, 1, σ),                     # 129 parameters
)     

Is there a better way to do this in Flux?

Best regards.

You could create the layers first and then stack them, something like this maybe

using Flux

function make_nn(length_in = 80, length_out = 1, nodes = 128, hid_lay = 3, act_fun = relu, act_fun_last = sigmoid)
	first_layer = Dense(length_in, nodes, act_fun)
    intermediate_layers = [Dense(nodes,nodes,act_fun) for _ in 1:hid_lay-1]
    last_layer = Dense(nodes, length_out, act_fun_last)

	return Chain(
        first_layer,
        intermediate_layers...,
        last_layer
    )
end
2 Likes