A class wraps a concrete function to handle different distributed contexts. The reason for wrapping a concrete function is because the _captured_inputs fields used for in-replica context and cross-replica context are different. When `load()` is called from within a tf.distribute.strategy scop
| 63 | |
| 64 | |
| 65 | class _WrapperFunction(function.ConcreteFunction): |
| 66 | """A class wraps a concrete function to handle different distributed contexts. |
| 67 | |
| 68 | The reason for wrapping a concrete function is because the _captured_inputs |
| 69 | fields used for in-replica context and cross-replica context are different. |
| 70 | When `load()` is called from within a tf.distribute.strategy scope, the |
| 71 | captured inputs are distributed variables. When using these distributed |
| 72 | variables during calling the function, we need different approaches when it is |
| 73 | in-replica and when it is not in-replica. When it is in replica, naturally we |
| 74 | should use the corresponding component of the distributed variable; when it is |
| 75 | not in-replica, calling the function should mean that it is constructing a |
| 76 | graph that is not actually going to be used. A typical use case is when |
| 77 | constructing a functional model. In this case, return a placeholder with a |
| 78 | control dependency to ensure that is is never accessed. |
| 79 | """ |
| 80 | |
| 81 | def __init__(self, concrete_function): |
| 82 | # Shallow copy the concrete_function |
| 83 | self.__dict__.update(vars(concrete_function)) |
| 84 | |
| 85 | def _call_flat(self, args, captured_inputs, cancellation_manager=None): |
| 86 | |
| 87 | def get_in_replica_handle(x): |
| 88 | return x.handle if ds_values.is_distributed_variable(x) else x |
| 89 | |
| 90 | def get_cross_replica_handle(x): |
| 91 | return _unused_handle() if ds_values.is_distributed_variable(x) else x |
| 92 | |
| 93 | if ds_context.get_replica_context() is not None: # in-replica context |
| 94 | captured_inputs = list(map(get_in_replica_handle, captured_inputs)) |
| 95 | else: # cross-replica context |
| 96 | captured_inputs = list( |
| 97 | map(get_cross_replica_handle, captured_inputs)) |
| 98 | return super(_WrapperFunction, self)._call_flat(args, captured_inputs, |
| 99 | cancellation_manager) |
| 100 | |
| 101 | |
| 102 | class Loader(object): |