Как сделать избирательное обратное распространение в мини-партии в Tensorflow?

Question

Недавно я работаю над проектом "Предсказание будущих траекторий объектов по их прошлым траекториям с использованием LSTM в Tensorflow". (Здесь траектория означает последовательность двухмерных позиций.)

Вход в LSTM - это, конечно, "прошлые траектории", а выход - "будущие траектории".

Размер мини-партии фиксируется при обучении. Однако количество прошедших траекторий в мини-партии может быть разным. Например, пусть размер мини-партии равен 10. Если у меня есть только 4 прошлых траектории для текущей итерации обучения, 6 из 10 в мини-партии дополняются нулевым значением.

При расчете потерь для обратного распространения я допускаю, что потери от 6 равны нулю, так что только 4 способствуют обратному распространению.

Проблема, с которой я сталкиваюсь, заключается в том, что Tensorflow все еще вычисляет градиенты для 6, даже если их потеря равна нулю. В результате скорость обучения становится меньше, когда я увеличиваю размер мини-партии, даже если я использую те же данные тренировки.

Я также использовал функцию tf.where при расчете потерь. Однако время тренировки не уменьшается.

Как я могу сократить время обучения?

Здесь я приложил свой псевдокод для обучения.

# For each frame in a sequence
for f in range(pred_length):

    # For each element in a batch
    for b in range(batch_size):


        with tf.variable_scope("rnnlm") as scope:
            if (f > 0 or b > 0):
                scope.reuse_variables()

            # for each pedestrian in an element
            for p in range(MNP):

                # ground-truth position
                cur_gt_pose = ...

                # loss mask
                loss_mask_ped = ... # '1' or '0'

                # go through RNN decoder
                output_states_dec_list[b][p], zero_states_dec_list[b][p] = cell_dec(cur_embed_frm_dec,
                                                                                    zero_states_dec_list[b][p])

                # fully connected layer for output
                cur_pred_pose_dec = tf.nn.xw_plus_b(output_states_dec_list[b][p], output_wd, output_bd)

                # go through embedding function for the next input
                prev_embed_frms_dec_list[b][p] = tf.reshape(tf.nn.relu(tf.nn.xw_plus_b(cur_pred_pose_dec, embedding_wd, embedding_bd)), shape=(1, rnn_size))

                # calculate MSE loss
                mse_loss = tf.reduce_sum(tf.pow(tf.subtract(cur_pred_pose_dec, cur_gt_pose_dec), 2.0))

                # only valid ped's traj contributes to the loss
                self.loss += tf.multiply(mse_loss, loss_mask_ped)

Answer 1

Я думаю, что вы ищете функцию tf.stop_gradient. Используя это, вы можете сделать что-то вроде tf.where(loss_mask, tensor, tf.stop_gradient(tensor)) для достижения желаемого результата, при условии, что размеры правильные.

Тем не менее, похоже, что это не ваша проблема. Кажется, что для каждого элемента в вашем наборе данных вы определяете новые узлы графа. Это не то, как TensorFlow должен функционировать, у вас должен быть только один график, построенный заранее, который выполняет некоторую фиксированную функцию, независимо от размера пакета. Определенно не следует определять новые узлы для каждого элемента в пакете, поскольку это не может эффективно использовать преимущества параллелизма.