在 TFF 中执行 Python 函数
本示例旨在着重展示如何将 Python 函数转换为 TFF 计算以及如何通过 TFF 评估该计算。
从用户角度:
tff.backends.native.set_local_python_execution_context() # 3
@tff.tf_computation(tf.int32) # 2
def add_one(x): # 1
return x + 1
result = add_one(2) # 4
编写 Python 函数。
使用
@tff.tf_computation
装饰 Python 函数。注:目前的重点仅在于对 Python 函数进行装饰,暂不深究装饰器本身的细节;下文将予以详细说明。
设置 TFF 上下文。
调用 Python 函数。
从 TFF 角度:
解析 Python 时,@tff.tf_computation
装饰器将跟踪 Python 函数并构造 TFF 计算。
调用装饰的 Python 函数时,即调用 TFF 计算,TFF 将在设置的上下文中编译和执行计算。
TF、TFF 与 Python 间的对比
tff.backends.native.set_local_python_execution_context()
@tff.tf_computation(tf.int32)
def add_one(x):
return x + 1
@tff.federated_computation(tff.type_at_clients(tf.int32))
def add_one_to_all_clients(values):
return tff.federated_map(add_one, values)
values = [1, 2, 3]
values = add_one_to_all_clients(values)
values = add_one_to_all_clients(values)
>>> [3, 4, 5]
TODO(b/153500547):描述 TF、TFF 与 Python 间对比情况的示例。