Каков рекомендуемый способ смешивания кода TensorFlow и TensorFlow Federated?

Question

TensorFlow (TF) и TensorFlow Federated (TFF) - это разные уровни функциональности, которые предназначены для хорошего взаимодействия (как следует из названия).

Тем не менее, это разные вещи, предназначенные для решения разных проблем.

Интересно, как лучше описать вычисления так, чтобы они могли использоваться как ванильными TF, так и в рабочих нагрузках TFF, а также типами подводных камней, которых можно было бы избежать.

Answer 1

Отличный вопрос Действительно, есть по крайней мере 3 способа приблизиться к составлению кода TensorFlow для использования с TFF, каждый со своими достоинствами.

1. Рекомендуется использовать композиционный механизм TensorFlow (defuns), если он работает для вашей конкретной ситуации. У TensorFlow уже есть механизмы для составления кода, и мы не хотим изобретать велосипед. Причина, по которой мы создали наш собственный композиционный механизм в TFF (@tff.tf_computation), заключалась в том, чтобы справиться с конкретными ограничениями (такими как отсутствие поддержки наборов данных на уровне интерфейса в TF и ​​необходимость взаимодействия компонентов TF с остальная часть TFF), и в идеале мы бы ограничивали использование этого механизма только теми ситуациями, которые действительно его требуют.

Когда это возможно, декорируйте компоненты TensorFlow с помощью @tf.function и оберните весь блок TensorFlow как вычисление @ tff.tf_ только на верхнем уровне, прежде чем встраивать его в @tff.federated_computation. Одним из многих преимуществ этого является то, что он позволяет вам тестировать компоненты вне TFF, используя стандартные инструменты TensorFlow.

Итак, следующие поощряются и предпочтительнее:

# here using TensorFlow's compositional mechanism (defuns)
# rather than TFF's to decorate "foo"
@tf.function(...)
def foo(...):
  ...

@tff.tf_computation(...)
def bar(...):
  # here relying on TensorFlow to embed "foo" as a component of "bar"
  ...foo(...)...

2. Использование композиционного механизма Python (обычные недекорированные функции Python) также является хорошим вариантом, хотя он менее предпочтителен, чем (1), поскольку он просто приводит к тому, что одна часть кода встраивается в другую во время определения, когда TFF отслеживает все TFF-декорации Python-функций для создания сериализованного представления вычислений для выполнения без предоставления вам изоляции или каких-либо других специальных преимуществ.

Вы все еще можете использовать этот шаблон, чтобы разрешить тестирование компонентов вне TFF или в ситуациях, когда ни (1), ни (3) не работает.

Итак, следующее является альтернативой, которую вы должны рассмотреть в первую очередь, если (1) не работает:

# here composing things in Python, no special TF or TFF mechanism employed
def foo(...):
  # keep in mind that in this case, "foo" can access and tamper with
  # the internal state of "bar" - you get no isolation benefits
  ... 

@tff.tf_computation(...)
def bar(...):
  # here effectively just executing "foo" within "bar" at the
  # time "bar" is traced
  ...foo(...)...

3. Использование композиционного механизма TFF (@tff.tf_computation) не рекомендуется, за исключением - как отмечалось выше - в ситуациях, когда это требуется, например, когда компонент TensorFlow должен принимать набор данных в качестве параметра или если он будет вызываться только из @tff.federated_computation. Имейте в виду, что поддержка TFF наборов данных в качестве параметров все еще является экспериментальной, и хотя в некоторых случаях это может быть единственным решением, вы все равно можете столкнуться с проблемами. Вы можете ожидать, что реализация будет развиваться.

Не рекомендуется (хотя в настоящее время иногда необходимо):

# here using TFF's compositional mechanism
@tff.tf_computation(...)
def foo(...):
  # here you do get isolation benefits - "foo" is traced and
  # serialized by TFF, but you can expect that e.g., some
  # tf.data.Dataset features won't work
  ...

@tff.tf_computation(...)
def bar(...):
  # here relying on TFF to embed "foo" within "bar"
  ...foo(...)...