Если вы просто хотите получить быстрый номер, вы можете просто добавить
grads = tf.gradients(C, [A, B])
в код @Tobias Scheck для построения узлов вычисления градиента. Затем вычтите новое число (с градиентными операциями) из исходного (без градиента операций), чтобы получить предполагаемые флопы.
Предостережение об использовании этого метода в более крупных проектах. Этот метод использует статический анализ всего графика. Это имеет несколько проблем, включая:
- Флопы из операций в цикле while будут добавлены только один раз.
- Будут добавлены операции, которые обычно не выполняются (некоторые функции TF могут оставить операции удаления мусора на графике).
- Этот анализ сильно зависит от вывода формы. Это может быть доступно не для всех операций.
- Этот анализ зависит от регистрации функций, которые могут оценить флопы данной операции. Могут быть операции без таких функций, и такие функции не могут точно моделировать флопы, сделанные фактическим ядром, которое TF выберет для выполнения операции.
Подробнее см .: https://github.com/tensorflow/tensorflow/blob/r1.8/tensorflow/core/profiler/g3doc/profile_model_architecture.md
Лучше использовать это в сочетании с фактической записью выполнения (RunMetadata) или использовать подход, основанный исключительно на времени выполнения, например, Могу ли я измерить время выполнения отдельных операций с TensorFlow? и выполнить некоторую фильтрацию / агрегацию результатов.