Reformer — Эффективный Трансформер

Colab-ноутбук демонстрации обработки «Преступления и наказания» нерабочий:

в конце 7го блока кода необходимо писать my_inputs(trax.fastmath.ops.local_device_count() вместо my_inputs(trax.fastmath.device_count(); в восьмом блоке должно быть ReformerLM.pos_axial_shape = (512, 1024) ReformerLM.pos_d_axial_embs= (64, 192) вместо ReformerLM.axial_pos_shape = (512, 1024) ReformerLM.d_axial_pos_embs= (64, 192); выполнение 11го и 12го блоков прерываются ошибкой "RuntimeError: Resource exhausted: Ran out of memory in memory space hbm. Used 8.73G of 8.00G hbm. Exceeded hbm capacity by 745.30M". С наскока вменяемый результат получить не удалось.

"Возможно, читателям этой статьи удастся найти новое применение Reformer'а." Если Reformer превратиться в рабочий инструмент, то его можно будет применять, например, в криминалистике для восстановления внешнего вида человека по ДНК.

Попытки использовать Reformer в качестве основы для модели sequence-to-sequence для перевода или предсказывания последовательностей не приводят к какому-либо существенному улучшению по сравнению с традиционными трансформерами. Видимо поэтому сами авторы Reformer(а) не приводят пример реализации sequence-to-sequence.