自定义损失函数显着减慢了多 GPU 模型的训练

Question

I'm training variational autoencoders on protein structures using Keras' multi_gpu_model. When switching from normal AEs to VAEs, my model takes >3x longer per epoch to train.

I identified the problem to be the loss function, changing it back to a built-in mse results in the same speed as seen before.

I'm using more or less the same vae_loss implementation as seen in many tutorials:

def vae_loss(y_true, y_pred):
    reconstruction_loss = recon_loss(y_true, y_pred)
    kl_loss = beta * K.mean(
             1 + K.flatten(z_log_var) - K.square(K.flatten(z_mean)) - K.exp(K.flatten(z_log_var)), axis=-1)
    kl_loss /= kl_loss_scaling # divide kl_loss by size of output dimension
    total_loss = K.mean(reconstruction_loss + kl_loss)
    return total_loss

When monitoring the GPU usage I realized, that they are well utilized and then drop to zero utilization after each epoch. The batchsize is adjusted for the number of GPUs and the exact same setup but with mse as loss works fine. It seems like the GPUs are waiting for the loss to be computed and hence have a considerable amount of downtime. (The effect is more strongly noticeable with smaller batchsizes, so increasing this parameter is somewhat of a solution, but I think this is far from optimal).

Is this unavoidable since this loss is more expensive to compute or is there something I can adjust to achieve better performance?

Answer 1

The root-cause might be the supports of custom loss function in Keras. If you use predefined loss in Keras, it works just fine. One thing you can try is to rewrite loss function to a Lambda layer and change model to be multi outputs, one is you original output and one is model loss.