Loop over a layer to do a Monte Carlo from a neural net output

Question

I've recently had to tweak a neural network. Here's how it works:

• Given an image as input, several layer turns it into a mean matrix mu and a covariance matrix sigma.
• Then, a sample z is taken from the Gaussian distribution of parameters mu, sigma.
• Several layer turns this sample into an output
• this output is compared to a given image, which gives a cost

What I want to do is to keep mu and sigma, take multiple samples z, propagate them through the rest of the NN, and compare the multiple images I get to a given image.

Note that the step z -> image output calls other package, I'd like not having to dig into these...

What I did so far :

• At first, I thought I did not need to go through all this hassle : I take a batch_size of one, it is as if I'm doing a Monte Carlo by running the NN multiple times. But I actually need the neural net to try several image before updating the weights, and thus changing mu and sigma.
• I simply sampled multiple z then propagated them through the net. But I soon discovered that I was duplicating all the layers, making the code terribly slow, and above all preventing me from taking many samples to achieve the MC I'm aiming at.

Of course, I updated the loss and data input classes to take that into account.

Do you have any ideas ? Basically, I'd like an efficient way to make z -> output multiple time, in a cost-efficient manner. I've still a lot to learn from tensorflow and keras, so I'm a little bit lost on how to do that. As usual, please apologized if an answer already exists somewhere, I did my best to look for one by myself!

Answer 1

Ok, my question was a bit stupid. So as not to duplicate layers, I created multiple slice layers, and then I simply propagated them through the net with previously declared layers. Here's my code :

# First declare layers
a = layer_A()
b = layer_B()
# And so on ...
# Generate samples
samples = generate_samples()([mu, sigma])
# for all the monte carlo samples, do : 
for i in range(mc_samples):
    cur_sample = Lambda(lambda x: K.slice(x, (0, 0, 0, 2*i), (-1, -1, -1, 2)), name="slice-%i" % i)(samples)
    cur_output = a(cur_sample)
    cur_output = b(cur_output)
    all_output.append(output)
output_of_net = keras.layers.concatenate(all_output)
return Model(inputs=inputs, outputs=output_of_net)

Simply loop over the last dimension in the loss function, average, and you're done ! A glimpse at my loss :

loss = 0
for i in range(mc_samples):
    loss += f(y_true[..., i], y_pred[..., i])
return loss/mc_samples