How to randomly set inputs to zero in keras during training autoencoder (callback)?

Question

I am training 2 autoencoders with 2 separate input paths jointly and I would like to randomly set one of the input paths to zero.

I use tensorflow with keras backend (functional API).

I am computing a joint loss (sum of two losses) for backpropagation.

A -> A' & B ->B'

loss => l2(A,A')+l2(B,B')

networks taking A and B are connected in latent space. I would like to randomly set A or B to zero and compute the loss only on the corresponding path, meaning if input path A is set to zero loss be computed only by using outputs of only path B and vice versa; e.g.:

0 -> A' & B ->B'

loss: l2(B,B')

How do I randomly set input path to zero? How do I write a callback which does this?

Answer 1

Maybe try the following:

import random
def decision(probability):
  return random.random() < probability

Define a method that makes a random decision based on a certain probability x and make your loss calculation depend on this decision.

if current_epoch == random.choice(epochs):

  keep_mask = tf.ones_like(A.input, dtype=float32)
  throw_mask = tf.zeros_like(A.input, dtype=float32)

  if decision(probability=0.5):
      total_loss = tf.reduce_sum(reconstruction_loss_a * keep_mask
                               + reconstruction_loss_b * throw_mask)
  else:
      total_loss = tf.reduce_sum(reconstruction_loss_a * throw_mask 
                               + reconstruction_loss_b * keep_mask)    
else:
  total_loss = tf.reduce_sum(reconstruction_loss_a + reconstruction_loss_b)
      
  

I assume that you do not want to set one of the paths to zero every time you update your model parameters, as then there is a risk that one or even both models will not be sufficiently trained. Also note that I use the input of A to create zero_like and one_like tensors as I assume that both inputs have the same shape; if this is not the case, it can easily be adjusted.

Depending on what your goal is, you may also consider replacing your input of A or B with a random tensor e.g. tf.random.normal based on a random decision. This creates noise in your model, which may be desirable, as your model would be forced to look into the latent space to try reconstruct your original input. This means precisely that you still calculate your reconstruction loss with A.input and A.output, but in reality your model never received the A.input, but rather the random tensor.

Note that this answer serves as a simple conceptual example. A working example with Tensorflow can be found here.

Answer 2

You can set an input to 0 simply:

A = A*random.choice([0,1])

This code can be used inside a loss function