Ошибка при использовании MultiWorkerMirroredStrategy для обучения модели исследования обнаружения объектов ssd_mobilenet_v1_fpn_640x640_coco17_tpu-8

Я пытаюсь обучить исследовательскую модель ssd_mobilenet_v1_fpn_640x640_coco17_tpu-8, используя MultiWorkerMirroredStrategy (путем установки --num_workers=2 при вызове model_main_tf2.py). Я пытаюсь обучить двух рабочих (0 и 1), каждый с одним графическим процессором. Однако, когда я пытаюсь это сделать, я получаю следующую ошибку, всегда на рабочем 1:

Traceback (most recent call last):
  File "C:\Users\JS\.conda\envs\tensor2\lib\site-packages\tensorflow\python\distribute\input_lib.py", line 553, in __next__
    return self.get_next()
  File "C:\Users\JS\.conda\envs\tensor2\lib\site-packages\tensorflow\python\distribute\input_lib.py", line 610, in get_next
    return self._get_next_no_partial_batch_handling(name)
  File "C:\Users\JS\.conda\envs\tensor2\lib\site-packages\tensorflow\python\distribute\input_lib.py", line 642, in _get_next_no_partial_batch_handling
    replicas.extend(self._iterators[i].get_next_as_list(new_name))
  File "C:\Users\JS\.conda\envs\tensor2\lib\site-packages\tensorflow\python\distribute\input_lib.py", line 1594, in get_next_as_list
    return self._format_data_list_with_options(self._iterator.get_next())
  File "C:\Users\JS\.conda\envs\tensor2\lib\site-packages\tensorflow\python\data\ops\multi_device_iterator_ops.py", line 580, in get_next
    result.append(self._device_iterators[i].get_next())
  File "C:\Users\JS\.conda\envs\tensor2\lib\site-packages\tensorflow\python\data\ops\iterator_ops.py", line 889, in get_next
    return self._next_internal()
  File "C:\Users\JS\.conda\envs\tensor2\lib\site-packages\tensorflow\python\data\ops\iterator_ops.py", line 819, in _next_internal
    ret = gen_dataset_ops.iterator_get_next(
  File "C:\Users\JS\.conda\envs\tensor2\lib\site-packages\tensorflow\python\ops\gen_dataset_ops.py", line 2922, in iterator_get_next
    _ops.raise_from_not_ok_status(e, name)
  File "C:\Users\JS\.conda\envs\tensor2\lib\site-packages\tensorflow\python\framework\ops.py", line 7186, in raise_from_not_ok_status
    raise core._status_to_exception(e) from None  # pylint: disable=protected-access
tensorflow.python.framework.errors_impl.OutOfRangeError: End of sequence [Op:IteratorGetNext]

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "C:\Users\JS\Desktop\Tensorflow\models\research\object_detection\model_main_tf2.py", line 114, in <module>
    tf.compat.v1.app.run()
  File "C:\Users\JS\.conda\envs\tensor2\lib\site-packages\tensorflow\python\platform\app.py", line 36, in run
    _run(main=main, argv=argv, flags_parser=_parse_flags_tolerate_undef)
  File "C:\Users\JS\.conda\envs\tensor2\lib\site-packages\absl\app.py", line 312, in run
    _run_main(main, args)
  File "C:\Users\JS\.conda\envs\tensor2\lib\site-packages\absl\app.py", line 258, in _run_main
    sys.exit(main(argv))
  File "C:\Users\JS\Desktop\Tensorflow\models\research\object_detection\model_main_tf2.py", line 105, in main
    model_lib_v2.train_loop(
  File "C:\Users\JS\.conda\envs\tensor2\lib\site-packages\object_detection\model_lib_v2.py", line 605, in train_loop
    load_fine_tune_checkpoint(
  File "C:\Users\JS\.conda\envs\tensor2\lib\site-packages\object_detection\model_lib_v2.py", line 401, in load_fine_tune_checkpoint
    _ensure_model_is_built(model, input_dataset, unpad_groundtruth_tensors)
  File "C:\Users\JS\.conda\envs\tensor2\lib\site-packages\object_detection\model_lib_v2.py", line 161, in _ensure_model_is_built
    features, labels = iter(input_dataset).next()
  File "C:\Users\JS\.conda\envs\tensor2\lib\site-packages\tensorflow\python\distribute\input_lib.py", line 549, in next
    return self.__next__()
  File "C:\Users\JS\.conda\envs\tensor2\lib\site-packages\tensorflow\python\distribute\input_lib.py", line 555, in __next__
    raise StopIteration
StopIteration

Рабочий 0 в конечном итоге выходит из строя после обнаружения того, что рабочий 1 вышел из строя.

Эта ошибка возникает независимо от физических машин, на которых работают два рабочих процесса. Другими словами, я вижу это, если я запускаю обоих рабочих на одной машине (используя локальный хост) ИЛИ на разных машинах в одной сети.

Судя по трассировке в сообщениях об ошибках, ошибка возникает всякий раз, когда цикл обучения пытается выполнить итерацию по обучающим данным, сгенерированным Strategy.experimental_distribute_datasets_from_function. Обратите внимание, что если я изменю стратегию на MirroredStrategy, она будет нормально работать на одной машине (других изменений не будет). Я не уверен, что я делаю что-то не так или есть ошибка в API обнаружения объектов.

Моя настройка на обеих машинах идентична (в основном я следовал инструкциям по настройке на веб-сайте обнаружения объектов):

1. Windows 10
2. Тензорный поток 2.8.0
3. Инструментарий Cuda 11.2
4. 8.1

Кто-нибудь когда-нибудь видел эту ошибку раньше? Если да, то есть ли способ обойти это?