Ensemble network with categorical distribution in tensorflow

Question

I have n networks, each with the same input / output. I want to randomly select one of the outputs according to a categorical distribution. Tfp.Categorical outputs only integers and I tried to do something like

act_dist = tfp.distributions.Categorical(logits=act_logits) # act_logits are all the same, so the distribution is uniform
rand_out = act_dist.sample()
x = nn_out1 * tf.cast(rand_out == 0., dtype=tf.float32) + ... # for all my n networks

But rand_out == 0. is always false, as well as the other conditions.

Any idea for achieving what I need?

Answer 1

You might also look at MixtureSameFamily, which does a gather under the covers for you.

nn_out1 = tf.expand_dims(nn_out1, axis=2)
...
outs = tf.concat([nn_out1, nn_nout2, ...], axis=2)
probs = tf.tile(tf.reduce_mean(tf.ones_like(nn_out1), axis=1, keepdims=True) / n, [1, n]) # trick to have ones of shape [None,1]
dist = tfp.distributions.MixtureSameFamily(
        mixture_distribution=tfp.distributions.Categorical(probs=probs),
        components_distribution=tfp.distributions.Deterministic(loc=outs))
x = dist.sample()