Tf 2.0 MirroredStrategy на модели Albert TF Hub (несколько графических процессоров)

Question

Я пытаюсь запустить версию концентратора Albert Tensorflow на нескольких графических процессорах на одном компьютере. Модель отлично работает на одном GPU.

Это структура моего кода:

strategy = tf.distribute.MirroredStrategy()
print('Number of devices: {}'.format(strategy.num_replicas_in_sync)) # it prints 2 .. correct
if __name__ == "__main__":
    with strategy.scope():
        run()

Где в функции run() я читаю данные, строю модель и подгоняю ее.

I получаю эту ошибку:

Traceback (most recent call last):
  File "Albert.py", line 130, in <module>
    run()
  File "Albert.py", line 88, in run
    model = build_model(bert_max_seq_length)
  File "Albert.py", line 55, in build_model
    model.compile(loss="categorical_crossentropy", optimizer=optimizer, metrics=["accuracy"])
  File "/home/****/py_transformers/lib/python3.5/site-packages/tensorflow_core/python/training/tracking/base.py", line 457, in _method_wrapper
    result = method(self, *args, **kwargs)
  File "/home/bighanem/py_transformers/lib/python3.5/site-packages/tensorflow_core/python/keras/engine/training.py", line 471, in compile
    '  model.compile(...)'% (v, strategy))
ValueError: Variable (<tf.Variable 'bert/embeddings/word_embeddings:0' shape=(30000, 128) dtype=float32>) was not created in the distribution strategy scope of (<tensorflow.python.distribute.mirrored_strategy.MirroredStrategy object at 0x7f62e399df60>). It is most likely due to not all layers or the model or optimizer being created outside the distribution strategy scope. Try to make sure your code looks similar to the following.
with strategy.scope():
  model=_create_model()
  model.compile(...)

Возможно ли, что эта ошибка происходит, потому что Albert модель была подготовлена ​​ранее командой tenorflow (собрана и скомпилирована)?

Отредактировано:

Если быть точным, версия Tensorflow имеет вид 2.1.

Кроме того, я загружаю предварительно обученную модель Альберта:

features = {"input_ids": in_id, "input_mask": in_mask, "segment_ids": in_segment, }

albert = hub.KerasLayer(
    "https://tfhub.dev/google/albert_xxlarge/3",
    trainable=False, signature="tokens", output_key="pooled_output",
)
x = albert(features)

После этого урока : SavedModels из TF Hub в TensorFlow 2

Answer 1

Ответ из двух частей:

1) TF Hub поддерживает две версии ALBERT (каждая в нескольких размерах):

• https://tfhub.dev/google/albert_base/3 et c. от исследовательской группы Google, которая первоначально разработала ALBERT, поставляется в формате hub.Module для TF1 . Это, вероятно, не будет работать со стратегией распространения TF2.

• https://tfhub.dev/tensorflow/albert_en_base/1 et c. от TensorFlow Model Garden поставляется в пересмотренном формате TF2 SavedModel . Пожалуйста, попробуйте этот для использования в TF2 со стратегией распространения.

2) Тем не менее, немедленная проблема заключается в том, что объясняется в сообщении об ошибке ( сокращено):

Переменная 'bert / embeddings / word_embeddings' не была создана в области стратегии распространения ... Попробуйте убедиться, что ваш код выглядит следующим образом.

with strategy.scope():
  model = _create_model()
  model.compile(...)

Для SavedModel (из TF Hub или другого) это загрузка, которая должна происходить в рамках стратегии распространения, потому что это то, что заново создает объекты tf.Variable в текущей программе. В частности, любой из следующих способов загрузки TF2 SavedModel из TF Hub должен выполняться в рамках стратегии распространения для распространения на работу:

• tf.saved_model.load();
• hub.load() , который просто вызывает tf.saved_model.load() (после загрузки, если необходимо);
• hub.KerasLayer при использовании со строковым дескриптором модели, для которого он затем вызывает hub.load().