Блог

Я пытаюсь обновить оптимизаторы, подобные тем, которые https://github.com/openai/blocksparse/blob/master/blocksparse/optimize.py для поддержки параметров, хранящихся в виде переменных ресурса. Глядя на исходный код TensorFlow для встроенных операций оптимизатора, кажется необходимым иметь относительно сложную логику блокировки и копирования при записи для обработки различных состояний, в которых может находиться переменная ресурса. Для этого TensorFlow использует вспомогательные функции в tensorflow/core/kernels/training_op_helpers.h, вызывая MaybeLockVariableInputMutexesInOrder и GetInputTensorFromVariable .

К сожалению, training_op_helpers.h этого нет в общедоступных заголовках TensorFlow, поэтому я не могу просто #include это сделать. Хуже того, это зависит от внутренних функций, которые libtensorflow_framework.so даже не экспортируются (в частности, tensorflow::functor::DenseUpdate для копирования при записи), поэтому копирования заголовков также недостаточно.

Есть ли предполагаемый способ сделать что-то подобное? На данный момент у меня возникает соблазн скомпилировать собственную ветку TensorFlow, которая экспортирует эти функции. (Я думаю, их переопределение было бы не так уж плохо.)