Tensorflow - Невозможно использовать BasicLSTMCell с распределением MirroredStrategy в Estimator

Question

У меня есть tf.nn.rnn_cell.BasicLSTMCell как часть моей архитектуры NN. Я использую цикл for, потому что он повторяется через ввод фиксированного количества временных шагов. Примерно так:

    lstm_cell = tf.nn.rnn_cell.BasicLSTMCell(num_units=lstm_dimensionality, name="forward_lstm")
    _, (lstm_memory, lstm_hidden) = lstm_cell(input_m, state=[lstm_memory, lstm_hidden])

    for i in range(3):
        # HERE is where the error is thrown
        _, (lstm_memory, lstm_hidden) = lstm_cell(input_m, state=[lstm_memory, lstm_hidden])

Локально работает на одном устройстве. Он также отлично работает в Google ML Engine на одном графическом процессоре. Однако, когда я пытаюсь распределить между 4 графическими процессорами, используя tf.distribute.MirroredStrategy, выдается исключение

ValueError: At least one of name (None) and default_name (None) must be provided.

. Вызываемый элемент lstm_cell даже не принимает параметр name, поэтому он сбивает с толку.

Здесь не так много места для деталей, поэтому в этом репозитории Github я создал игрушечный пример для воспроизведения ошибки в ML Engine. Именно на этой строке , где выдается ошибка.

Тензор потока: 1.13.1 ML Двигатель: --runtime-version 1.13

Answer 1

В вашем коде здесь вы используете область действия в функции compute_initial_lstm_state.

Вы повторно используете 2 возвращаемых значения здесь

Вы используете область для генерации значений и назначаете их без области.

Это должно быть вашей корневой ошибкой. С одним графическим процессором область может быть определена автоматически. Но с мультипроцессором это невозможно и не получится.