Как оптимизировать использование памяти для цикла в тензорном потоке?

Question

Возьмем, например, среднеквадратичную ошибку (MSE). Обычно эта функция определяется следующим образом:

def exp_loss(batch_p, batch_t):
    loss_val = tf.reduce_mean(tf.squared_difference(batch_p, batch_t))
    return loss_val

Но когда я использую цикл for для вычисления ошибки поэлементно, вот так:

def exp_loss_for(batch_p, batch_t):
    loss_val = 0
    ns = int(batch_p.get_shape()[0])  # batch_size
    sl = int(batch_p.get_shape()[1])  # sequence_length
    nd = int(batch_p.get_shape()[2])  # num_dim
    for i in range(ns):
        for j in range(sl):
            for k in range(nd):
                loss_val += tf.square(tf.subtract(batch_p[i, j, k], batch_t[i, j, k]))
    loss_val = loss_val / (ns * sl * nd)
    return loss_val

tenorflow будет потреблять слишком много памяти на этапе построения графа. Если мне нужно настроить собственную функцию потерь, используя цикл for, например exp_loss_for, но более сложный, есть ли способ расставить точки, чтобы уменьшить использование памяти?

Answer 1

Избегайте петель. Еще более сложные функции потерь обычно можно выразить с помощью операций с тензорным потоком без использования циклов for.