Векторизация умножения в Tensorflow

Question

У меня следующая проблема в Tensorflow:

#y_true.shape = (8, 256, 512) 
#y_pred.shape = (8, 256, 512, 39) where 39 is the different classes

focal_loss_weight = tf.multiply(y_true, tf.pow(
   tf.subtract(1., y_pred), 0.3))

Я пытаюсь найти общую потерю, поэтому мне нужно создать для l oop итерацию по 39 классам из y_pred и добавить их вместе в полная потеря измерения (8,256,512). Как бы векторизовать путь?

Answer 1

Довольно просто, если я правильно понял:

y_true_broadcast = y_true[..., tf.newaxis]  # add axis at the end to make dimensions match
# in this next line, y_true will be replicated over the 39 classes.
# I hope this is what you were asking for...
loss_per_class = y_true_broadcast * tf.pow(1. - y_pred, 0.3)
total_loss = tf.reduce_sum(loss_per_class, axis=-1)  # sum over classes

Это полностью параллельно и должно быть довольно быстрым.