I am trying to implement a Bayesian neural network in Flux. I have custom layers which return two outputs. One is the output of the layer and another is a conditional probability. I use these probabilities to define a loss function using the negative log-likelihood. My question is that how can I define a custom layer with two outputs instead of one?
1 Like
just return both: return out1, out2
If the layer is stacked into a Chain, make sure the following layer can work with an input tuple of length 2
So is this code true? All the layers in the network are of the same type.
`
> struct BayesianLayer
> model
> end
>
> function (l::BayesianLayer)(x)
> len = length(x)
> if len == 1
> out1 = l.model(x)
> out2 = somefunction(out1)
> elseif len == 2
> out1 = l.model(x[1])
> out2 = somefunction(out1)*x[2]
> end
> return out1, out2
> end
`