From 9645a1249d3bdbe8e930af62d1958120a940c31d Mon Sep 17 00:00:00 2001 From: Asim Shankar Date: Wed, 19 Sep 2018 19:06:13 -0700 Subject: [PATCH 01/15] RFC: Functions not Sessions in TensorFlow 2.0 --- rfcs/20180918-functions-not-sessions-20.md | 991 +++++++++++++++++++++ 1 file changed, 991 insertions(+) create mode 100644 rfcs/20180918-functions-not-sessions-20.md diff --git a/rfcs/20180918-functions-not-sessions-20.md b/rfcs/20180918-functions-not-sessions-20.md new file mode 100644 index 000000000..f0cf6bf23 --- /dev/null +++ b/rfcs/20180918-functions-not-sessions-20.md @@ -0,0 +1,991 @@ +# TensorFlow 2.0: Functions, not Sessions. + +| Status | (Proposed / Accepted / Implemented / Obsolete) | +:-------------- |:---------------------------------------------------- | +| **Author(s)** | ashankar@google.com, joshl@google.com | +| **Sponsor** | apassos@google.com | +| **Updated** | 2018-09-18 | + +## Objective + +This document elaborates on the proposal to make TensorFlow be more "Pythonic" in 2.0. In four bullet points, the proposal is to: + + + +* Encourage[^1] the encapsulation of graph computation as Python functions \ +(where the graph is executed when the function is invoked, instead of via `Session`) +* Align "state" in the TensorFlow runtime (e.g., resource tensors like those that back `tf.Variable` objects) with state in the Python program (e.g., Python objects corresponding to the runtime state with lifetimes attached to each other). +* Make it easy to export these encapsulations to a `GraphDef`+Checkpoint and/or `SavedModel`. +* Enable eager execution by default. +* Provide a path for incorporating existing code that uses the 1.x APIs to construct TensorFlows graphs as functions in TensorFlow 2.x programs. + +This document liberally employs the use of sample code to describe the end-user effect of proposed changes. + + + +## Design Proposal + +### Basic idea: Python functions as Graphs + +Today, the TensorFlow graph defines the union of all computation that the author of the graph may be interested in. The actual computation to execute is defined by the arguments to `tf.Session.run`. Once this subgraph is defined, the runtime can optimize and execute. For example, consider the following: + + +``` +import tensorflow as tf + +x = tf.placeholder(tf.float32) +y = tf.square(x) +z = tf.add(x, y) + +sess = tf.Session() + +z0 = sess.run([z], feed_dict={x: 2.}) # 6.0 +z1 = sess.run([z], feed_dict={x: 2., y: 2.}) # 4.0 +``` + + + \ +Though there is one `tf.Graph` object the user is interacting with (`tf.get_default_graph()`), the two `sess.run` calls are executing different programs (indeed the runtime ends up with two separate `Graph` objects in C++, one for each program), equivalent to: + + +``` +def compute_z0(x): + return tf.add(x, tf.square(x)) + +def compute_z1(x, y): + return tf.add(x, y) +``` + + +The core proposal of this document is the alignment between computation expressed in Python and the computation executed by the runtime. Instead of defining a graph and then selecting the subgraph to execute at `sess.run()` time, the exact computation of interest is encapsulated in a Python callable. For example, the program above that uses `sess.run()` to compute `z0` and `z1` can be written as: + + +``` +import tensorflow as tf + +@tf.defun +def compute_z1(x, y): + return tf.add(x, y) + +@tf.defun +def compute_z0(x): + return compute_z1(x, tf.square(x)) + +z0 = compute_z0(2.) +z1 = compute_z1(2., 2.) +``` + + +Where `tf.defun` is a decorator that "**de**fines a TensorFlow **fun**ction". A "TensorFlow function" defines a computation as a graph of TensorFlow operations, with named arguments and explicit return values. Users define the function they want TensorFlow to "accelerate" as a Python function and integrate it into their Python program like any other Python function call. + +Having the Python function correspond to what the runtime will execute reduces conceptual complexity in translating between the two domains. It also affords an opportunity to provide more helpful stacktraces on errors. More advanced features available today (e.g., carving sub-graphs, feeding intermediate values) will still be possible (discussed later), though most users should not need to think in terms of graphs, feeds, and fetches. The constructed graph also provides a natural point for accelerators/acceleration libraries (NVIDIA TensorRT, Google Cloud TPUs etc.) to hook in for rewrites. + + +### `defun`: A brief specification + +`defun` constructs a TensorFlow graph by "tracing" the TensorFlow operations executed by the Python function. Specifically: + + + +* `defun(f)` is a Python function that returns a Python callable, `C` +* When the `C` is invoked it: + 1. Determines an "input signature" \ +If one was not explicitly specified by the user (as an argument to `defun`), the signature is computed from the types of the input arguments (including `dtype` and `shape` for `Tensor` arguments) + 1. If a new input signature is encountered, then it invokes `f` to create a TensorFlow graph, `G`. If the input signature has been seen before, it looks up `G` from a cache keyed by the input signature. + 1. It executes the graph defined by `G` and feeding each argument as a value of the corresponding `Placeholder` node in the graph. + +Changes in input signature result in a new graph being traced. For example: + + +``` +@tf.defun +def f(x): + one = tf.constant(1, dtype=x.dtype) + return tf.add(x, one) + +# Traces a graph with int32 operations and executes it +f(tf.constant(1, dtype=tf.int32)) +# Traces a graph with float32 operations and executes it. +f(tf.constant(1, dtype=tf.float32)) +``` + + + +### Referencing state: Variables, tables etc. + +A `defun` decorated Python function encapsulates a graph and its execution. The Python function may reference stateful objects (i.e., state backed by `DT_RESOURCE` tensors in the runtime, e.g., `tf.Variable`) by referencing the corresponding Python object, and these will be captured as implicit inputs to the function. + +Comparing TensorFlow code today with how we propose it looks in 2.x: + + + + + + + + + + + +
TensorFlow 1.x + 2.0 +
+ + + +
W = tf.get_variable(
+  "weight", shape=[10, 10])
+b = tf.get_variable(
+  "bias", shape=[10],
+  initializer=tf.zeros_initializer())
+c = tf.get_variable(
+  "counter", shape=[],
+  dtype=tf.int32,
+  initializer=tf.zeros_initializer())
+
+x = tf.placeholder(tf.float32)
+ctr = c.assign_add(1)
+with tf.control_dependencies([ctr]):
+  y = tf.matmul(x, W) + b
+init = 
+  tf.global_variables_initializer()
+
+with tf.Session() as sess:
+  sess.run(init)
+  print(sess.run(y,
+  feed_dict={x: make_input_value()}))
+  assert int(sess.run(c)) == 1
+ + +
+ + + +
W = tf.Variable(
+  tf.glorot_uniform_initializer()(
+    (10, 10))
+b = tf.Variable(tf.zeros(10))
+c = tf.Variable(0)
+
+@tf.defun
+def f(x):
+  c.assign_add(1)
+  return tf.matmul(x, W) + b
+
+print(f(make_input_value())
+assert int(c) == 1
+ + +
+ + +Worthy of note here - in TensorFlow 1.x, the memory underlying the variables `W` and `b` in the runtime lives for the lifetime of the `Session` - unrelated to the lifetime of the Python objects. In 2.x, the lifetime of the Python objects and the runtime state are tied together. + + +### Program-order semantics / Control dependencies + +In TensorFlow graphs today, control dependencies are sometimes needed to ensure correct evaluation order. For example, consider the following: + + +``` +v = tf.Variable(1.0) +init_op = tf.global_variables_initializer() +assign_op = v.assign(2.0) +read = v.read_value() + +with tf.Session() as sess: + sess.run(init_op) + val = sess.run(read) + print(val) # Will print 1.0, the assign is ignored + val = sess.run([read, assign_op])[0] + print(val) # Non-deterministically prints 1.0 or 2.0, +``` + + +The output here is not deterministic, since `val` may evaluate to either 1.0 or 2.0 depending on whether the runtime happened to execute `assign_op` before `y` or not. `tf.control_dependencies` is a mechanism provided to add annotations at graph construction time to influence graph execution. The TensorFlow user, a Python programmer, is thus forced to think about two execution models - TensorFlow graphs and the Python interpreter. To eliminate this cognitive load, `defun` will automatically insert control dependencies to ensure that operations that produce or consume a given `DT_RESOURCE` tensor and operations that are marked stateful (`REGISTER_OP(...).SetIsStateful()`) follow graph construction order. Thus: + + +``` +v = tf.Variable(1.0) +@tf.defun +def f(): +  v.assign(2.0) +  return v.read_value() + +print(f()) # Always prints 2.0. +``` + + +A preview of this implemented in `tf.contrib.eager.defun` today (using [AutomaticControlDependencies](https://github.com/tensorflow/tensorflow/blob/2f886d17f1990da418366bd093a09fb01fe5e777/tensorflow/python/eager/function.py#L1800)). + + +### Functions that create state + +In the above code, no `tf.Variable` objects are created inside a `tf.defun` decorated function. This is makes it clear that the code will have the same semantics once wrapped. + +Note that if the function naturally creates state only on the first trace, all is well: + + +``` +v = None + +@tf.defun +def f(x): + global v + if v is None: + v = tf.Variable(1.0) + return tf.cast(x, tf.float32) + v + +f(tf.constant(1, dtype=tf.float32)) # Creates the variable, returns 2.0 +f(tf.constant(2, dtype=tf.int32)) # Reuses the variable, returns 3.0 +``` + + +To support this `defun` imposes some requirements on the decorated function: + + + +1. State (like `tf.Variable` objects) are only created the first time the function `f` is called. \ +How that is accomplished is left up to the implementation of `f`. \ +If any variables are created in the first execution of `f`, then `@tf.defun` will trace `f` a second time in order to record the behavior that will be used from then on. No variables may be created during that second trace, or any other trace after that (due to different dtypes, shapes, or non-tensor arguments). +1. The caller must make sure that any variable referenced by the function still exists whenever the function is evaluated. \ +`@tf.defun` itself will keep only weak references to these created variables. Thus, if the referenced state does not exist when the decorated function is invoked, an exception will be raised. + +In the future we may want to allow for function local `tf.Variable`s, which are created and destroyed each time the decorated function is invoked. + + +### API for `defun` + +We've introduced a single new symbol: `defun` that consumes a Python function and returns a callable Python object. The precise API of the object needs some iteration, but at a high level it will have methods to: + + + +* List out all captured state (`tf.Variable` objects, other `DT_RESOURCE` tensors used by the computation and provided as implicit inputs). +* Access the `tf.Graph` that corresponds to the graph executed by the `__call__` method of the object. +* Execute the function with custom `RunOptions` and retrieve `RunMetadata`. + +Since new graphs are traced when new input signatures are encountered, a `defun` can encapsulate multiple graphs. For example, consider the following snippet: + + +``` +@tf.defun +def f(x): + return tf.square(x) + +f(int(1)) +f(float(1.0)) +``` + + + \ +There are two graphs created here - one which corresponds to the `Square` operation applied to `DT_INT32` tensors, and one with the `Square` operation applied to `DT_FLOAT32` tensors. The object returned by `defun` encapsulates multiple graphs (lazily generated based on the type and shape of input arguments), multiplexing between them in `__call__`. + +The same holds for the case where arguments are not `Tensor`s, for example: + + +``` +@tf.defun +def f(x, use_multiply): + return tf.multiply(x, x) if use_multiply else tf.square(x) + +f(2.0, True) +f(2.0, False) +``` + + +will result in 2 graphs being created. + +Note that the "type" of `Tensor` inputs to the function also incorporates the shape. For example: + + +``` +@tf.defun +def f(x): return tf.add(x, 1.) +f([2.0]) +f([2.0, 3.0]) +f([[2.0]]) +``` + + +will result in 3 graphs being created: + + + +1. One for when the first argument is a `tf.float32` vector with 1 element \ +(input signature: `((tf.float32, [1]))`) +1. One for when the first argument is a `tf.float32` vector with 2 elements \ +(input signature: `((tf.float32, [2]]))`) +1. One for when the first argument is a `tf.float32` 1x1 matrix \ +(input signature: `((tf.float32, [1, 1]))`) + +Tracing the decorated function to create a new graph on each input shape is a conservative choice (allowing for `f` to create graphs dependent on the shape), which may be unnecessary. Users can explicitly specify an input signature to ensure that the same graph is used for multiple inputs. For example: + + +``` +@tf.defun(input_signature=((tf.float32, [None])) +def f(x): return tf.add(x, 1.) + +f([2.0]) # Returns [3.0] +f([2.0, 3.0]) # Matches the input signature as [None] matches the actual shape [2] +f([[2.0]]) # Raises an error as the arguments don't match the input signature. + +# f is backed by a single Graph since the input signature specification allowed +# for the same graph to be used when the input shape is (1,) or (2,). +``` + + + + + +### Classes + +If a member function of a class does not create variables, it may be decorated with `@tf.defun` and it will work: + + +``` +class ScalarModel(object): + def __init__(self): + self.v = tf.Variable(0) + + @tf.defun + def increment(self, amount): + self.v.assign_add(amount) + +model1 = ScalarModel() +model1.increment(tf.constant(3)) +assert int(model1.v) == 3 +model1.increment(tf.constant(4)) +assert int(model1.v) == 7 +model2 = MyModel() +model2.increment(tf.constant(5)) +assert int(model2.v) == 5 +``` + + + \ +This works since `increment()` has `self` as a non-tensor argument, and a new trace will be created for each value of `self`. However, if variables are created in a method, we want to allow a new set of variables for every instantiation of `self`. You get this behavior by using `@tf.method`: + + +``` +class AnyShapeModel(object): + def __init__(self): + self.v = None + + @tf.method + def increment(self, amount): + if self.v is None: + self.v = tf.Variable(tf.zeros_like(amount)) + self.v.assign_add(amount) + +model1 = AnyShapeModel() +model1.increment(tf.constant(3)) +assert int(model1.v) == 3 +model1.increment(tf.constant(4)) +assert int(model1.v) == 7 +model2 = MyModel() +model2.increment(tf.constant([4, 5])) +assert model2.v.numpy() == [4, 5] +``` + + +The semantics here are that each new instance is allowed to create variables in each `@tf.method` once. The simple recommendation would be "always use `@tf.method` on methods, use `@tf.defun` for functions outside of a class". In the above example, if `increment` was decorated with `@tf.defun` instead, then the `model2.increment()` call would raise an exception (as per `defun`s stated behavior of disallowing state creation on anything but the first trace). + +In addition, as long as all variable creation/initialization happens while we are tracing, we should be able to support exporting the initialization graph when exporting a `SavedModel` or `MetaGraphDef`. + + +### Transitioning from 1.x + +The definition of `tf.defun` above is careful to check that invoking a decorated Python function would have the same behavior as invoking an undecorated function. This is to guard against it being passed code from TensorFlow v1.x that expects to only be called once (and relies on things like graph collections to track which variables are created), for example: + + +``` +def f(x, do_add): + v = tf.Variable(5.0) + if do_add: + v.assign_add(x) + else: + v.assign_sub(x) + return v +``` + + +For this case, we use a different API, `tf.compat.v1.wrap_function`, that treats any created variables as static local state: + + +``` +f_add = tf.compat.v1.wrap_function(f, tf.TensorSpec(tf.float32, ()), True) + +assert float(f_add(1.0)) == 6.0 +assert float(f_add(1.0)) == 7.0 + +# Can call tf.compat.v1.wrap_function again to get a new trace, a new set +# of variables, and possibly different non-template arguments. +f_sub = tf.compat.v1.wrap_function(f, tf.TensorSpec(tf.float32, ()), False) + +assert float(f_sub(1.0)) == 4.0 +assert float(f_sub(1.0)) == 3.0 +``` + + +Note these differences from `tf.defun`: + + + +* Only ever traces `f()` once (per call to `tf.compat.v1.wrap_function`). +* The complete input tensor signature (via `tf.TensorSpec` calls) and the values of all non-tensor arguments must be specified when wrapping the function. Note: we may want a `tf.tensor_like(x)` convenience function that returns `tf.TensorSpec(x.dtype, x.shape)`. +* Will include extra TF v1.x compatibility features like collections, and access v1.x APIs like `tf.compat.v1.get_variable()` +* Will not automatically insert control dependencies to maintain program order across stateful operations/state accesses. +* May only use a function or Python constant to initialize variables, no tensors. This is a technical limitation, required by the fact that we need some way of disentangling the initializers for variables from the other operations from the function. +* Keeps strong references to variables created in f, weak references to other variables accessed by f. This is to match the v1.x graph behavior that variables have the lifetime of the graph they are created, and can generally be accessed through graph collections. Some common patterns of writing v1.x code don't leave any references to those variables around. Keeping references to those variables extends their lifetime to match that of the object returned by `tf.compat.v1.wrap_function`. +* Typically won't be used as a decorator. Calling `tf.compat.v1.wrap_function` takes some arguments, traces the function, and creates an object with state. The lifetime of the return value should be tracked explicitly by saving it in a variable. + +Treating state (like `tf.Variable`) as static local does mean that the behavior of a `tf.compat.v1.wrap_function`-decorated Python function differs from that of an undecorated one. In the above example, `f(1.0, True)` will always return 6.0 (as a scalar `Tensor`), while each call to `f_add(1.0)` will return a different value. We propose this separate `tf.compat.v1.wrap_function` endpoint specifically to make it easy to migrate TensorFlow 1.x libraries to the TensorFlow 2.0. The behavior of 2.0 `tf.defun` is restricted to cases where we can say that the behavior will match. + +We recognize that code written for TensorFlow 1.x commonly does not encapsulate state in Python objects, instead adding to hidden (graph-)global collections. We will support code that accesses collections inside a `tf.compat.v1.wrap_function`, though those collections will be local to a single trace. + +With the `tf.compat.v1.wrap_function` proposed above, most graph construction library functions written against TensorFlow 1.x can be incorporated into TensorFlow 2.x programs. + + +``` +def f(x): + W = tf.compat.v1.get_variable(name="weight", shape=[10, 10]) + b = tf.compat.v1.get_variable(name="bias", shape=[10], + initializer=tf.zeros_initializer()) + c = tf.Variable(0, dtype=tf.int32, name="counter") + with tf.control_dependencies([c.assign_add(1)]): + return tf.matmul(x, W) + b +``` + + + +``` +f = tf.compat.v1.wrap_function(f, tf.placeholder(tf.float32, None)) +print(f(make_input_value())) +assert len(f.variables) == 3 +assert f.variables[0].name == "weight" +``` + + + \ +In this case, the object returned by `tf.compat.v1.wrap_function` owns the state created within `f`, and the `__call__` method on it invokes the corresponding computation. + +Long story short, `tf.compat.v1.wrap_function` helps in incorporating graph construction code written against TensorFlow 1.x into TensorFlow 2.x programs. `wrap_function` constructs the same object as a `defun` decorated function, which provides the conceptual equivalent of graph construction and `Session.run`. + + +### Serialization: Exporting SavedModel/GraphDefs + +So far we've only considered Python programs. One of the key features of TensorFlow is the ability to integrate models created (and possibly trained) in a Python program into an application written in another programming language and/or platform (e.g., servers, mobile phones, self-driving cars). This ability will of course remain, with a smoother path to exporting models. + +In TensorFlow 1.x, "saving a model" could mean one of three things: + + + +1. Saving parameter values, but not the computation: \ +A "checkpoint" containing the values of all model parameters. \ +(`tf.train.Saver` / `tf.train.Checkpoint`) \ +Restoring this model required that the restoring program duplicate the Python code to construct the graph with the same model parameters. +1. Saving the computation graph, but not the parameter values: \ +The computation is represented by a `GraphDef` that can be exported by calls to `tf.Graph.as_graph_def()`, or `tf.train.export_meta_graph()`, and reconstructed by calls to `tf.import_graph_def()` / `tf.train.import_meta_graph()`. \ +Note that the parameter (`tf.Variable`) values are not saved, but their initializers are. +1. Saving both the computation and the parameter values: \ +The two packaged together in a SavedModel. \ +At a high level, the SavedModel format packages the `MetaGraphDef`, checkpoint, and a signature (names of input and output tensors). \ +(`tf.saved_model.simple_save` / `tf.saved_model.builder.SavedModelBuilder`) \ +This is the format preferred for exporting for serving via TensorFlow Serving or to other languages (e.g., `SavedModelBundle.load()` in Java, `LoadSavedModel` in Go) + +The objects created by `defun` encapsulate (1) the computation expressed as a `GraphDef`, (2) the state used by it. Thus, these objects are naturally suited for import/export in any of the above formats, using something like the following: + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
TensorFlow 1.x + 2.x +
Save only the parameters, not the computation +
+ + + +
W = tf.get_variable(
+  "weights", shape=[10, 10])
+
+# Presumably the train_op is
+# a little fancier 
+train_op = W.assign_add(1.)
+saver = tf.train.Saver()
+
+with tf.Session() as sess:
+  sess.run(W.initializer)
+  sess.run(train_op)
+  saver.save(sess, "/tmp/checkpoint/")
+
+with tf.Session() as sess:
+  saver.restore(sess, "/tmp/checkpoint/")
+  sess.run(train_op)
+ + +
+ + + +
W = tf.Variable(
+  tf.glorot_uniform_initializer()(
+    (10, 10)))
+
+@tf.defun
+def train():
+  W.assign_add(1.)
+
+train()
+ckpt = tf.train.Checkpoint(W=W)
+ckpt.save("/tmp/checkpoint")
+ckpt.restore("/tmp/checkpoint")
+ + +
Exporting/Importing GraphDefs +
+ + + +
W = tf.get_variable("weights", shape=[10, 10])
+x = tf.placeholder(
+  tf.float32, shape=(None, 10)))
+y = tf.matmul(x, W)
+
+graph = tf.get_default_graph()
+graph_def =  graph.as_graph_def()
+with open("/tmp/graph.pb", "w") as f:
+  f.write(
+      graph_def.SerializeToString())
+
+tf.reset_default_graph()
+
+graph_def = tf.GraphDef()
+with open("/tmp/graph.pbtxt") as f:
+  graph_def.ParseFromString(f.read())
+
+tf.import_graph_def(graph_def)
+ + +
+ + + +
W = tf.Variable(
+  tf.glorot_uniform_initializer()(
+    (10, 10)))
+
+@tf.defun
+def f(x):
+  return tf.matmul(x, W)
+
+# Retrieve the object corresponding to
+# a particular input signature:
+graph = f.graph_function(
+  (tf.float32, (None, 10)).graph
+graph_def = graph.as_graph_def()
+
+with open("/tmp/graph.pb", "w") as f:
+  f.write(graph_def.SerializeToString())
+ + \ + +
Exporting/Importing SavedModels +
+ + + +
+def save_model():
+  W = tf.get_variable("weights",
+                      shape=[10, 10])
+  x = tf.placeholder(
+    tf.float32, shape=(None, 10))
+  y = tf.matmul(x, W)
+
+  with tf.Session() as sess:
+    sess.run(
+    tf.global_variables_initializer())
+    tf.saved_model.simple_save(
+      sess,
+      "/tmp/model",
+      inputs={"x": x},
+      outputs={"y": y})
+
+def load_model():
+  sess = tf.Session()
+  with sess.as_default():
+    inputs, outputs =  tf.saved_model.simple_load(sess, "/tmp/model")
+  return inputs, outputs, sess
+ + +
To be worked on but something along the lines of: + + + +
+class Model(tf.train.Checkpointable):
+  def __init__(self):
+    self.W = tf.Variable(...)
+
+  @tf.method
+  def f(self, x):
+    return tf.matmul(x, self.W)
+
+m = Model()
+
+tf.export_saved_model(m, "/tmp/model")
+
+m = tf.import_saved_model("/tmp/model")
+
+ + +
+ + + +### Derived/Related Graphs + +One reservation expressed by TensorFlow graph/session enthusiasts today is that the ability to write generic analysis/inspection tooling on graphs, precluding the need to understand or modify the Python code that constructed the graph, is important to them. To put it differently, some find it easier to navigate the `GraphDef` program than navigating the Python program. \ + + +This ability will be maintained. `defun`-decorated Python functions have an associated graph, and new functions can be created by specifying the sub-graph of interest. For example: + + + + + + + + + + + + + + + + + + + + + +
TensorFlow 1.x + TensorFlow 2.x +
Carving out a subgraph +
+ + + +
def build_graph():
+  x = tf.placeholder(tf.float32)
+  y = tf.square(x)
+  z = tf.square(y)
+
+with tf.Session() as sess:
+  build_graph()
+  sess.run("Square_1:0",
+   feed_dict={"Square:0": 2.0})  # 4.0
+ + +
+ + + +
@tf.defun
+def f(x):
+  return tf.square(tf.square(x))
+
+# tf.Graph corresponding to "x" 
+# being a float32 tensor with unknown
+# shape
+graph = f.graph_function(
+  (tf.float32, None)).graph
+
+f2 = tf.NewGraphFunction(
+  graph,
+  inputs=["Square:0"], 
+  outputs=["Square_1:0"])
+# The above may optionally take a
+# "prune" argument to allow for
+# pruning stateful operations in
+# `graph` that are not in the path
+# from inputs to outputs.
+f2(2.0) # 4.0
+ + +
Extending a graph +
+ + + +
def build_graph():
+  x = tf.placeholder(tf.float32)
+  y = tf.square(x)
+  return y
+
+y = build_graph()
+z = tf.square(y)
+
+with tf.Session() as sess:
+  # Line below will return 16.0
+  sess.run(z, feed_dict={"Placeholder:0": 2.0))
+ + +
+ + + +
@tf.defun
+def f(x):
+  return tf.square(x)
+
+@tf.defun
+def g(x):
+  return tf.square(f(x))
+
+g(2.0) # 16.0
+ + +
+ + + +### Distributed Execution + +At the lowest level of the API, distributed execution continues to work with `tf.device` annotations, where the device name can reference remote devices as well, just like they do today. + +The `DistributionStrategy` API, typically aimed at synchronous training will continue to be the method of choice (where the API can be used inside a `defun`). Other APIs such as go/tf-replicator should also be usable. + +The author realizes that this section can do with more detail. However, to keep this document more focused, these details will be discussed separately. In particular, usage of `MonitoredSession` and session hooks today needs additional thought. + + +### `defun`-ing Python control flow + +`defun` decorates a graph construction function and transparently recreates graphs if needed. However, this does mean that if the function has data-dependent control flow then though the function will execute fine with eager execution enabled, `defun` decorating it will fail. For example: + + +``` +def f(x, y): + if tf.equal(y, 0.0): + return y + return x / y + +x = tf.constant(2.0) +y = tf.constant(2.0) + +f(x, y) # Will be 1.0 + +df = tf.defun(f) +df(x, y) # Will raise an error complaining about the data-dependent control flow +``` + + + \ +To fix this, one would have to use the graph construction APIs for control flow (`tf.cond`, `tf.while_loop`): + + +``` +def f(x, y): + return tf.cond(tf.equal(y, 0.0), lambda: y, lambda: x/y) + +x = tf.constant(2.0) +y = tf.constant(2.0) + +f(x, y) # Will be 1.0 + +df = tf.defun(f) +df(x, y) # Will be 1.0 +``` + + +This situation can be improved with the help of go/tf-autograph to allow expression of control flow in Python. Whether autograph will be enabled by default or not is still under debate, but the option will be there as a flag on defun. + + +### Summaries + +**Background**: See the updated 2.0 summaries design (go/tf-summaries-2.0) and plan (go/tf-2.0-summaries). Support for TensorFlow 1.x summaries is a non-goal. + +The summary writing operations ([tb.summary.scalar](https://www.tensorflow.org/api_docs/python/tf/contrib/summary/scalar), [tb.summary.image](https://www.tensorflow.org/api_docs/python/tf/contrib/summary/image) etc.) can be naturally placed in the graph by using them in a defun-decorated function. These operations require two "external" inputs - the summary writer resource and the condition, that will be picked up from the context (e.g., [tb.summary.create_file_writer](https://www.tensorflow.org/api_docs/python/tf/contrib/summary/create_file_writer) and [tb.summary.record_summary_every_n_global_steps](https://www.tensorflow.org/api_docs/python/tf/contrib/summary/record_summary_every_n_global_steps)). When defining the graph, these inputs are converted to placeholders, which are then resolved at function invocation time. Thus, something like this: + + +``` +writer = tf.contrib.summary.create_file_writer('/tmp/test') +with writer.as_default(), tf.contrib.summary.always_record_summaries(): + f() +with writer.as_default(), tf.contrib.summary.never_record_summaries(): + f() +``` + + +Will write one summary to `writer` whether `f` is defined as: + + +``` +def f(): + tb.summary.scalar("loss", compute_loss()) +``` + + +Or + + +``` +f = tf.contrib.eager.defun(f) +``` + + +(NOTE: As of August 2018, this is not the case, but it will be. See b/112269952). + +Note that the runtime is free to prune away the summary writing operations when the function is invoked in a context where there is no summary writer resource or the condition is false. + + +### What does that have to do with eager execution? + +So far this proposal has dealt with the encapsulation of TensorFlow graphs in Python functions with the intention of making it easier to integrate TensorFlow-accelerated computation in Python programs. + +_Additionally_, this proposal suggests enabling eager execution by default in TensorFlow 2.0. Keeping `defun` in mind, this basically means: + + + +* Inside the context of defining a TensorFlow function (i.e., within a `defun` decorated function) `tf.Tensor` objects created refer to symbolic tensors. +* Outside this context, `tf.Tensor` objects created are backed by concrete values and TensorFlow API. The underlying memory of the tensor can be backed by any device (i.e., CPU/GPU) and is not restricted to host-memory (like numpy arrays). + +See the [docstring for tf.contrib.eager.defun](https://www.tensorflow.org/api_docs/python/tf/contrib/eager/defun) - the evolving playground for the implementation of the proposal in this document. The basic takeaway is that: + + + +* For users that embrace symbolic tensors and graphs, continue doing so with your code placed inside a `defun` decorated Python function. +* We believe most users (new ones in particular) will find it more convenient to deal with `Tensor` objects backed by concrete values and then selectively "compiling" portions of their Python program into TensorFlow graphs rather than being exposed to graph metaprogramming in Python upfront. In spirit, this is similar to Swift4TensorFlow with the obvious glaring difference that[ graph program extraction](https://github.com/tensorflow/swift/blob/master/docs/DesignOverview.md#graph-program-extraction) here is manually specified (with the `defun` decoration). + +NOTE: In TensorFlow 1.x, eager execution is enabled by [tf.enable_eager_execution()](https://www.tensorflow.org/api_docs/python/tf/enable_eager_execution). Once invoked, all public API endpoints that consume or produce symbolic Tensor objects begin to produce and consume Tensor objects that are backed by a concrete value. See the "Research and Experimentation" section at [www.tensorflow.org/tutorials](http://www.tensorflow.org/tutorials) for an introduction. + + +### A few things of note + + + +* This change **only** applies to the TensorFlow **Python** frontend + * [TensorFlow.js](https://js.tensorflow.org/) is already "eager by default". + * [Switf4TensorFlow](https://github.com/tensorflow/swift) has [similar design goals](https://github.com/tensorflow/swift/blob/master/docs/DesignOverview.md#swift-for-tensorflow-design-overview), doing away with the define-then-run style of TensorFlow graphs. + * Most other language bindings ([Java](https://www.tensorflow.org/api_docs/java/reference/org/tensorflow/package-summary), [C++](https://www.tensorflow.org/api_guides/cc/guide), [Go](https://godoc.org/github.com/tensorflow/tensorflow/tensorflow/go), others) are mostly targeting deployment of defined models in applications. While an imperative style might help simplify model development and training in these languages, doing so is explicitly out of scope for TensorFlow 2.0. The notion of graphs and sessions will remain in them, as well as in the stable [C API](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/c/c_api.h). +* Users of **Estimator** will see no change + * Canned Estimators are black boxes that create and train models. Enabling eager execution will have no effect on their usage. This is true today. + * The model_fn of a regular (non-canned) Estimator will remain as a graph construction function. +* [SavedModel](https://www.tensorflow.org/guide/saved_model#save_and_restore_models) will continue to be the format encouraged for exporting trained models + * Crudely speaking, a SavedModel encapsulates a Graph, a checkpoint of variable values, and some metadata like signature information (names of input and output tensors). + * A path will be provided to easily export models in this format (e.g., via tf.keras.Model.save()). There may be instances where converting the Python code to a graph is not trivial, in which case, exporting to a SavedModel (and thus a Graph) will fail. + + +## Alternatives Considered + + +### Creating state inside a `defun` + +How state (`DT_RESOURCE` tensors) created inside a `defun` should be handled is actively being debated. Options include: + + + +1. "Lifting" state out as a static local function variable +1. Mimic the undecorated code - creating and destroying variables on each call. + + +#### "Static-local" state + +`tf.contrib.eager.defun` today treats state as function-static variables, which allows for code like: + + +``` +def f(x): + v = tf.Variable(1, dtype=x.dtype) + v.assign_add(x) + return v + +df = tf.contrib.eager.defun(f) +# tf.defun(f) proposed in this document will raise an exception on first use +x = tf.constant(1, dtype=tf.float32)) +print(df(x)) # 2.0 +print(df(x)) # 3.0 +``` + + + \ +However, the one major issue with this approach is that it behaves differently from how an undecorated function would: + + +``` +print(f(1.0), df(1.0)) # 2.0, 2.0 +print(f(1.0), df(1.0)) # 2.0, 3.0 +``` + + +To be conservative, we propose some restrictions on `defun`, such as: + + + +1. State is created only once, i.e., `defun` will fail if calling `f` a second time results in new state being created. +1. `defun` decorated functions can only produce `Tensor` return values. +1. If you want to convert TF v1.x code like `f` above, you may use `tf.compat.v1.wrap_function` which guarantees it will only trace `f` once. + + +#### Function-local state + +Another option would be to match typical Python functions, where state is created and destroyed during the call to the function. So: + + +``` +def f(x): + v = tf.Variable(1.0) + v.assign_add(x) + return v + +df = tf.defun(f) + +assert f(1.0) == df(1.0) # Both will be 2.0 +assert f(1.0) == df(1.0) # Still 2.0, since 'v' would be recreated. +``` + + + \ +This seems like an avenue definitely worth pursuing, but requires careful consideration of some additional design points such as escape analysis of return values (e.g. the lifetime of `tf.Variable` objects that are returned from a decorated function). + +For now, we propose that `defun` continue with the restricted abilities proposed in this document and a "maintain Python semantics" decorator be investigated independently. + + +## Questions and Discussion Topics + +* Naming: + * `tf.defun` or `tf.function`? + * `tf.compat.v1.wrap_function` or `tf.compat.v1.defun` or `tf.compat.v1.function` or `tf.compat.v1.wrap_graph_as_function`? + + +## Notes + +[^1]: + + We say "encourage" instead of "require" since removing the Session API from the Python frontend within a year may be an unrealistic aspiration. Particularly given the use in Estimators and the use of MonitoredSession and hooks. The `Session` API may have to stick around in `tf.compat.v1`. + From 6f265d655ca556c0a9929266247c92201589baae Mon Sep 17 00:00:00 2001 From: Asim Shankar Date: Wed, 19 Sep 2018 19:09:45 -0700 Subject: [PATCH 02/15] Formatting tweaks --- rfcs/20180918-functions-not-sessions-20.md | 60 +++++++++++----------- 1 file changed, 29 insertions(+), 31 deletions(-) diff --git a/rfcs/20180918-functions-not-sessions-20.md b/rfcs/20180918-functions-not-sessions-20.md index f0cf6bf23..f9b836eb3 100644 --- a/rfcs/20180918-functions-not-sessions-20.md +++ b/rfcs/20180918-functions-not-sessions-20.md @@ -21,6 +21,10 @@ This document elaborates on the proposal to make TensorFlow be more "Pythonic" i This document liberally employs the use of sample code to describe the end-user effect of proposed changes. +(We say "encourage" instead of "require" since removing the Session API from the +Python frontend within a year may be an unrealistic aspiration. Particularly +given the use in Estimators and the use of MonitoredSession and hooks. The +`Session` API may have to stick around in `tf.compat.v1`). ## Design Proposal @@ -30,7 +34,7 @@ This document liberally employs the use of sample code to describe the end-user Today, the TensorFlow graph defines the union of all computation that the author of the graph may be interested in. The actual computation to execute is defined by the arguments to `tf.Session.run`. Once this subgraph is defined, the runtime can optimize and execute. For example, consider the following: -``` +```python import tensorflow as tf x = tf.placeholder(tf.float32) @@ -48,7 +52,7 @@ z1 = sess.run([z], feed_dict={x: 2., y: 2.}) # 4.0 Though there is one `tf.Graph` object the user is interacting with (`tf.get_default_graph()`), the two `sess.run` calls are executing different programs (indeed the runtime ends up with two separate `Graph` objects in C++, one for each program), equivalent to: -``` +```python def compute_z0(x): return tf.add(x, tf.square(x)) @@ -60,7 +64,7 @@ def compute_z1(x, y): The core proposal of this document is the alignment between computation expressed in Python and the computation executed by the runtime. Instead of defining a graph and then selecting the subgraph to execute at `sess.run()` time, the exact computation of interest is encapsulated in a Python callable. For example, the program above that uses `sess.run()` to compute `z0` and `z1` can be written as: -``` +```python import tensorflow as tf @tf.defun @@ -97,7 +101,7 @@ If one was not explicitly specified by the user (as an argument to `defun`), the Changes in input signature result in a new graph being traced. For example: -``` +```python @tf.defun def f(x): one = tf.constant(1, dtype=x.dtype) @@ -187,7 +191,7 @@ Worthy of note here - in TensorFlow 1.x, the memory underlying the variables `W` In TensorFlow graphs today, control dependencies are sometimes needed to ensure correct evaluation order. For example, consider the following: -``` +```python v = tf.Variable(1.0) init_op = tf.global_variables_initializer() assign_op = v.assign(2.0) @@ -205,7 +209,7 @@ with tf.Session() as sess: The output here is not deterministic, since `val` may evaluate to either 1.0 or 2.0 depending on whether the runtime happened to execute `assign_op` before `y` or not. `tf.control_dependencies` is a mechanism provided to add annotations at graph construction time to influence graph execution. The TensorFlow user, a Python programmer, is thus forced to think about two execution models - TensorFlow graphs and the Python interpreter. To eliminate this cognitive load, `defun` will automatically insert control dependencies to ensure that operations that produce or consume a given `DT_RESOURCE` tensor and operations that are marked stateful (`REGISTER_OP(...).SetIsStateful()`) follow graph construction order. Thus: -``` +```python v = tf.Variable(1.0) @tf.defun def f(): @@ -226,7 +230,7 @@ In the above code, no `tf.Variable` objects are created inside a `tf.defun` deco Note that if the function naturally creates state only on the first trace, all is well: -``` +```python v = None @tf.defun @@ -267,7 +271,7 @@ We've introduced a single new symbol: `defun` that consumes a Python function an Since new graphs are traced when new input signatures are encountered, a `defun` can encapsulate multiple graphs. For example, consider the following snippet: -``` +```python @tf.defun def f(x): return tf.square(x) @@ -283,7 +287,7 @@ There are two graphs created here - one which corresponds to the `Square` operat The same holds for the case where arguments are not `Tensor`s, for example: -``` +```python @tf.defun def f(x, use_multiply): return tf.multiply(x, x) if use_multiply else tf.square(x) @@ -298,7 +302,7 @@ will result in 2 graphs being created. Note that the "type" of `Tensor` inputs to the function also incorporates the shape. For example: -``` +```python @tf.defun def f(x): return tf.add(x, 1.) f([2.0]) @@ -321,7 +325,7 @@ will result in 3 graphs being created: Tracing the decorated function to create a new graph on each input shape is a conservative choice (allowing for `f` to create graphs dependent on the shape), which may be unnecessary. Users can explicitly specify an input signature to ensure that the same graph is used for multiple inputs. For example: -``` +```python @tf.defun(input_signature=((tf.float32, [None])) def f(x): return tf.add(x, 1.) @@ -342,7 +346,7 @@ f([[2.0]]) # Raises an error as the arguments don't match the input signature If a member function of a class does not create variables, it may be decorated with `@tf.defun` and it will work: -``` +```python class ScalarModel(object): def __init__(self): self.v = tf.Variable(0) @@ -366,7 +370,7 @@ assert int(model2.v) == 5 This works since `increment()` has `self` as a non-tensor argument, and a new trace will be created for each value of `self`. However, if variables are created in a method, we want to allow a new set of variables for every instantiation of `self`. You get this behavior by using `@tf.method`: -``` +```python class AnyShapeModel(object): def __init__(self): self.v = None @@ -398,7 +402,7 @@ In addition, as long as all variable creation/initialization happens while we ar The definition of `tf.defun` above is careful to check that invoking a decorated Python function would have the same behavior as invoking an undecorated function. This is to guard against it being passed code from TensorFlow v1.x that expects to only be called once (and relies on things like graph collections to track which variables are created), for example: -``` +```python def f(x, do_add): v = tf.Variable(5.0) if do_add: @@ -412,7 +416,7 @@ def f(x, do_add): For this case, we use a different API, `tf.compat.v1.wrap_function`, that treats any created variables as static local state: -``` +```python f_add = tf.compat.v1.wrap_function(f, tf.TensorSpec(tf.float32, ()), True) assert float(f_add(1.0)) == 6.0 @@ -446,7 +450,7 @@ We recognize that code written for TensorFlow 1.x commonly does not encapsulate With the `tf.compat.v1.wrap_function` proposed above, most graph construction library functions written against TensorFlow 1.x can be incorporated into TensorFlow 2.x programs. -``` +```python def f(x): W = tf.compat.v1.get_variable(name="weight", shape=[10, 10]) b = tf.compat.v1.get_variable(name="bias", shape=[10], @@ -458,7 +462,7 @@ def f(x): -``` +```python f = tf.compat.v1.wrap_function(f, tf.placeholder(tf.float32, None)) print(f(make_input_value())) assert len(f.variables) == 3 @@ -789,7 +793,7 @@ The author realizes that this section can do with more detail. However, to keep `defun` decorates a graph construction function and transparently recreates graphs if needed. However, this does mean that if the function has data-dependent control flow then though the function will execute fine with eager execution enabled, `defun` decorating it will fail. For example: -``` +```python def f(x, y): if tf.equal(y, 0.0): return y @@ -809,7 +813,7 @@ df(x, y) # Will raise an error complaining about the data-dependent control flo To fix this, one would have to use the graph construction APIs for control flow (`tf.cond`, `tf.while_loop`): -``` +```python def f(x, y): return tf.cond(tf.equal(y, 0.0), lambda: y, lambda: x/y) @@ -833,7 +837,7 @@ This situation can be improved with the help of go/tf-autograph to allow expres The summary writing operations ([tb.summary.scalar](https://www.tensorflow.org/api_docs/python/tf/contrib/summary/scalar), [tb.summary.image](https://www.tensorflow.org/api_docs/python/tf/contrib/summary/image) etc.) can be naturally placed in the graph by using them in a defun-decorated function. These operations require two "external" inputs - the summary writer resource and the condition, that will be picked up from the context (e.g., [tb.summary.create_file_writer](https://www.tensorflow.org/api_docs/python/tf/contrib/summary/create_file_writer) and [tb.summary.record_summary_every_n_global_steps](https://www.tensorflow.org/api_docs/python/tf/contrib/summary/record_summary_every_n_global_steps)). When defining the graph, these inputs are converted to placeholders, which are then resolved at function invocation time. Thus, something like this: -``` +```python writer = tf.contrib.summary.create_file_writer('/tmp/test') with writer.as_default(), tf.contrib.summary.always_record_summaries(): f() @@ -845,7 +849,7 @@ with writer.as_default(), tf.contrib.summary.never_record_summaries(): Will write one summary to `writer` whether `f` is defined as: -``` +```python def f(): tb.summary.scalar("loss", compute_loss()) ``` @@ -854,7 +858,7 @@ def f(): Or -``` +```python f = tf.contrib.eager.defun(f) ``` @@ -919,7 +923,7 @@ How state (`DT_RESOURCE` tensors) created inside a `defun` should be handled is `tf.contrib.eager.defun` today treats state as function-static variables, which allows for code like: -``` +```python def f(x): v = tf.Variable(1, dtype=x.dtype) v.assign_add(x) @@ -937,7 +941,7 @@ print(df(x)) # 3.0 However, the one major issue with this approach is that it behaves differently from how an undecorated function would: -``` +```python print(f(1.0), df(1.0)) # 2.0, 2.0 print(f(1.0), df(1.0)) # 2.0, 3.0 ``` @@ -957,7 +961,7 @@ To be conservative, we propose some restrictions on `defun`, such as: Another option would be to match typical Python functions, where state is created and destroyed during the call to the function. So: -``` +```python def f(x): v = tf.Variable(1.0) v.assign_add(x) @@ -982,10 +986,4 @@ For now, we propose that `defun` continue with the restricted abilities proposed * `tf.defun` or `tf.function`? * `tf.compat.v1.wrap_function` or `tf.compat.v1.defun` or `tf.compat.v1.function` or `tf.compat.v1.wrap_graph_as_function`? - -## Notes - -[^1]: - - We say "encourage" instead of "require" since removing the Session API from the Python frontend within a year may be an unrealistic aspiration. Particularly given the use in Estimators and the use of MonitoredSession and hooks. The `Session` API may have to stick around in `tf.compat.v1`. From 7b7d8cb8b6b21f5fb850356cb1d748e3f0803933 Mon Sep 17 00:00:00 2001 From: Asim Shankar Date: Wed, 19 Sep 2018 19:12:20 -0700 Subject: [PATCH 03/15] Formatting tweaks --- rfcs/20180918-functions-not-sessions-20.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/rfcs/20180918-functions-not-sessions-20.md b/rfcs/20180918-functions-not-sessions-20.md index f9b836eb3..9949d2591 100644 --- a/rfcs/20180918-functions-not-sessions-20.md +++ b/rfcs/20180918-functions-not-sessions-20.md @@ -659,9 +659,9 @@ class Model(tf.train.Checkpointable): m = Model() -tf.export_saved_model(m, "/tmp/model") +tf.saved_model.export("/tmp/model", m) -m = tf.import_saved_model("/tmp/model") +m = tf.saved_model.import("/tmp/model") From 72a8314c40db0db9a831d8581f5fcece1f74ac40 Mon Sep 17 00:00:00 2001 From: Asim Shankar Date: Thu, 20 Sep 2018 12:40:42 -0700 Subject: [PATCH 04/15] Fix some links --- rfcs/20180918-functions-not-sessions-20.md | 79 +++++++++++++--------- 1 file changed, 46 insertions(+), 33 deletions(-) diff --git a/rfcs/20180918-functions-not-sessions-20.md b/rfcs/20180918-functions-not-sessions-20.md index 9949d2591..8bff1ed3a 100644 --- a/rfcs/20180918-functions-not-sessions-20.md +++ b/rfcs/20180918-functions-not-sessions-20.md @@ -8,11 +8,11 @@ ## Objective -This document elaborates on the proposal to make TensorFlow be more "Pythonic" in 2.0. In four bullet points, the proposal is to: +This document presents a proposal to make TensorFlow be more "Pythonic" in 2.0. In five bullet points, the proposal is to: -* Encourage[^1] the encapsulation of graph computation as Python functions \ +* Encourage the encapsulation of graph computation as Python functions \ (where the graph is executed when the function is invoked, instead of via `Session`) * Align "state" in the TensorFlow runtime (e.g., resource tensors like those that back `tf.Variable` objects) with state in the Python program (e.g., Python objects corresponding to the runtime state with lifetimes attached to each other). * Make it easy to export these encapsulations to a `GraphDef`+Checkpoint and/or `SavedModel`. @@ -21,14 +21,12 @@ This document elaborates on the proposal to make TensorFlow be more "Pythonic" i This document liberally employs the use of sample code to describe the end-user effect of proposed changes. -(We say "encourage" instead of "require" since removing the Session API from the -Python frontend within a year may be an unrealistic aspiration. Particularly -given the use in Estimators and the use of MonitoredSession and hooks. The -`Session` API may have to stick around in `tf.compat.v1`). +(We say "encourage" instead of "require" since removing the Session API from the Python frontend within a year may be an unrealistic aspiration. Particularly given the use in Estimators and the use of MonitoredSession and hooks. The `Session` API may have to stick around in `tf.compat.v1`.) ## Design Proposal + ### Basic idea: Python functions as Graphs Today, the TensorFlow graph defines the union of all computation that the author of the graph may be interested in. The actual computation to execute is defined by the arguments to `tf.Session.run`. Once this subgraph is defined, the runtime can optimize and execute. For example, consider the following: @@ -91,12 +89,13 @@ Having the Python function correspond to what the runtime will execute reduces c +* `f` is a Python function that returns zero or more `Tensor`s * `defun(f)` is a Python function that returns a Python callable, `C` * When the `C` is invoked it: 1. Determines an "input signature" \ -If one was not explicitly specified by the user (as an argument to `defun`), the signature is computed from the types of the input arguments (including `dtype` and `shape` for `Tensor` arguments) - 1. If a new input signature is encountered, then it invokes `f` to create a TensorFlow graph, `G`. If the input signature has been seen before, it looks up `G` from a cache keyed by the input signature. - 1. It executes the graph defined by `G` and feeding each argument as a value of the corresponding `Placeholder` node in the graph. +If an input signature was not explicitly specified by the user (as an argument to `defun`), the signature is computed from the types of the input arguments (including `dtype` and `shape` for `Tensor` arguments) + 1. Every time a new input signature is encountered, it invokes `f` to create a TensorFlow graph, `G`. If the input signature has been seen before, it looks up `G` from a cache keyed by the input signature. + 1. It executes the graph defined by `G,` feeding each argument as a value of the corresponding `Placeholder` node in the graph. Changes in input signature result in a new graph being traced. For example: @@ -134,15 +133,11 @@ Comparing TensorFlow code today with how we propose it looks in 2.x: -
W = tf.get_variable(
-  "weight", shape=[10, 10])
-b = tf.get_variable(
-  "bias", shape=[10],
-  initializer=tf.zeros_initializer())
-c = tf.get_variable(
-  "counter", shape=[],
-  dtype=tf.int32,
-  initializer=tf.zeros_initializer())
+
W = tf.Variable(
+  tf.glorot_uniform_initializer()(
+    (10, 10)))
+b = tf.Variable(tf.zeros(10))
+c = tf.Variable(0)
 
 x = tf.placeholder(tf.float32)
 ctr = c.assign_add(1)
@@ -165,7 +160,7 @@ with tf.Session() as sess:
 
 
W = tf.Variable(
   tf.glorot_uniform_initializer()(
-    (10, 10))
+    (10, 10)))
 b = tf.Variable(tf.zeros(10))
 c = tf.Variable(0)
 
@@ -206,7 +201,7 @@ with tf.Session() as sess:
 ```
 
 
-The output here is not deterministic, since `val` may evaluate to either 1.0 or 2.0 depending on whether the runtime happened to execute `assign_op` before `y` or not. `tf.control_dependencies` is a mechanism provided to add annotations at graph construction time to influence graph execution. The TensorFlow user, a Python programmer, is thus forced to think about two execution models - TensorFlow graphs and the Python interpreter. To eliminate this cognitive load, `defun` will automatically insert control dependencies to ensure that operations that produce or consume a given `DT_RESOURCE` tensor and operations that are marked stateful (`REGISTER_OP(...).SetIsStateful()`) follow graph construction order. Thus:
+The output here is not deterministic, since `val` may evaluate to either 1.0 or 2.0 depending on whether the runtime happened to execute `assign_op` before `y` or not. `tf.control_dependencies` is a mechanism provided to add annotations at graph construction time to influence graph execution. The TensorFlow user, a Python programmer, is thus forced to think about two execution models - TensorFlow graphs and the Python interpreter. To eliminate this cognitive load, `defun` will automatically insert control dependencies to ensure that (1) operations that produce or consume a given `DT_RESOURCE` tensor and (2) operations that are marked stateful (`REGISTER_OP(...).SetIsStateful()`) follow graph construction order. Thus:
 
 
 ```python
@@ -220,6 +215,23 @@ print(f()) # Always prints 2.0.
 ```
 
 
+Note that the intention here is to avoid _observable_ differences from program order. For example:
+
+
+```python
+a = tf.Variable(1.0)
+b = tf.Variable(1.0)
+@tf.defun
+def f():
+  a.assign(2.0)
+  b.assign(3.0)
+  return a + b
+print(f())
+```
+
+
+Will always print 5.0 since the assignments will occur before the read. However, there is no guaranteed ordering between the assignment of `a` and `b` (as any difference in that is not observable).
+
 A preview of this implemented in `tf.contrib.eager.defun` today (using [AutomaticControlDependencies](https://github.com/tensorflow/tensorflow/blob/2f886d17f1990da418366bd093a09fb01fe5e777/tensorflow/python/eager/function.py#L1800)).
 
 
@@ -260,7 +272,7 @@ In the future we may want to allow for function local `tf.Variable`s, which are
 
 ### API for `defun`
 
-We've introduced a single new symbol: `defun` that consumes a Python function and returns a callable Python object. The precise API of the object needs some iteration, but at a high level it will have methods to:
+We've introduced a single new symbol: `defun` that consumes a Python function and returns a callable Python object. The precise API of the object is being iterated on, but at a high level it will have methods to:
 
 
 
@@ -360,7 +372,7 @@ model1.increment(tf.constant(3))
 assert int(model1.v) == 3
 model1.increment(tf.constant(4))
 assert int(model1.v) == 7
-model2 = MyModel()
+model2 = ScalarModel()
 model2.increment(tf.constant(5))
 assert int(model2.v) == 5
 ```
@@ -386,13 +398,15 @@ model1.increment(tf.constant(3))
 assert int(model1.v) == 3
 model1.increment(tf.constant(4))
 assert int(model1.v) == 7
-model2 = MyModel()
+model2 = AnyShapeModel()
 model2.increment(tf.constant([4, 5]))
 assert model2.v.numpy() == [4, 5]
 ```
 
 
-The semantics here are that each new instance is allowed to create variables in each `@tf.method` once. The simple recommendation would be "always use `@tf.method` on methods, use `@tf.defun` for functions outside of a class". In the above example, if `increment` was decorated with `@tf.defun` instead, then the `model2.increment()` call would raise an exception (as per `defun`s stated behavior of disallowing state creation on anything but the first trace).
+The semantics here are that each new instance is allowed to create variables in each `@tf.method` once. The simple recommendation would be "always use `@tf.method` on methods, use `@tf.defun` for functions outside of a class".
+
+In the above example, if `increment` was decorated with `@tf.defun` instead, then the `model2.increment()` call would raise an exception (as per `defun`s stated behavior of disallowing state creation on anything but the first trace). However, if the method didn't create any state then `@tf.defun` or `@tf.method` would both have the same effect.
 
 In addition, as long as all variable creation/initialization happens while we are tracing, we should be able to support exporting the initialization graph when exporting a `SavedModel` or `MetaGraphDef`.
 
@@ -661,7 +675,8 @@ m = Model()
 
 tf.saved_model.export("/tmp/model", m)
 
-m = tf.saved_model.import("/tmp/model")
+m =
+  tf.saved_model.import("/tmp/model")
 
@@ -783,7 +798,7 @@ g(2.0) # 16.0
At the lowest level of the API, distributed execution continues to work with `tf.device` annotations, where the device name can reference remote devices as well, just like they do today. -The `DistributionStrategy` API, typically aimed at synchronous training will continue to be the method of choice (where the API can be used inside a `defun`). Other APIs such as go/tf-replicator should also be usable. +The `DistributionStrategy` API, typically aimed at synchronous training will continue to be the method of choice (where the API can be used inside a `defun`). The author realizes that this section can do with more detail. However, to keep this document more focused, these details will be discussed separately. In particular, usage of `MonitoredSession` and session hooks today needs additional thought. @@ -827,13 +842,11 @@ df(x, y) # Will be 1.0 ``` -This situation can be improved with the help of go/tf-autograph to allow expression of control flow in Python. Whether autograph will be enabled by default or not is still under debate, but the option will be there as a flag on defun. +This situation can be improved with the help of [autograph](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/autograph) to allow expression of control flow in Python. Whether autograph will be enabled by default or not is still under debate, but the option will be there as a flag on defun. ### Summaries -**Background**: See the updated 2.0 summaries design (go/tf-summaries-2.0) and plan (go/tf-2.0-summaries). Support for TensorFlow 1.x summaries is a non-goal. - The summary writing operations ([tb.summary.scalar](https://www.tensorflow.org/api_docs/python/tf/contrib/summary/scalar), [tb.summary.image](https://www.tensorflow.org/api_docs/python/tf/contrib/summary/image) etc.) can be naturally placed in the graph by using them in a defun-decorated function. These operations require two "external" inputs - the summary writer resource and the condition, that will be picked up from the context (e.g., [tb.summary.create_file_writer](https://www.tensorflow.org/api_docs/python/tf/contrib/summary/create_file_writer) and [tb.summary.record_summary_every_n_global_steps](https://www.tensorflow.org/api_docs/python/tf/contrib/summary/record_summary_every_n_global_steps)). When defining the graph, these inputs are converted to placeholders, which are then resolved at function invocation time. Thus, something like this: @@ -896,7 +909,7 @@ NOTE: In TensorFlow 1.x, eager execution is enabled by [tf.enable_eager_ex * This change **only** applies to the TensorFlow **Python** frontend * [TensorFlow.js](https://js.tensorflow.org/) is already "eager by default". * [Switf4TensorFlow](https://github.com/tensorflow/swift) has [similar design goals](https://github.com/tensorflow/swift/blob/master/docs/DesignOverview.md#swift-for-tensorflow-design-overview), doing away with the define-then-run style of TensorFlow graphs. - * Most other language bindings ([Java](https://www.tensorflow.org/api_docs/java/reference/org/tensorflow/package-summary), [C++](https://www.tensorflow.org/api_guides/cc/guide), [Go](https://godoc.org/github.com/tensorflow/tensorflow/tensorflow/go), others) are mostly targeting deployment of defined models in applications. While an imperative style might help simplify model development and training in these languages, doing so is explicitly out of scope for TensorFlow 2.0. The notion of graphs and sessions will remain in them, as well as in the stable [C API](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/c/c_api.h). + * Most other language bindings ([Java](https://www.tensorflow.org/api_docs/java/reference/org/tensorflow/package-summary), [C++](https://www.tensorflow.org/api_guides/cc/guide), [Go](https://godoc.org/github.com/tensorflow/tensorflow/tensorflow/go), others) are mostly targeting deployment of defined models in applications. While an imperative style might help simplify model development and training in these languages, doing so is explicitly out of scope for TensorFlow 2.0. The notion of graphs and sessions will remain in them, as well as in the stable [C API](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/c/c_api.h). In these APIs, the lifetime of program state (like variables) will continue to be tied to the lifetime of the `Session`. * Users of **Estimator** will see no change * Canned Estimators are black boxes that create and train models. Enabling eager execution will have no effect on their usage. This is true today. * The model_fn of a regular (non-canned) Estimator will remain as a graph construction function. @@ -980,10 +993,10 @@ This seems like an avenue definitely worth pursuing, but requires careful consid For now, we propose that `defun` continue with the restricted abilities proposed in this document and a "maintain Python semantics" decorator be investigated independently. -## Questions and Discussion Topics +## Open Questions + + * Naming: * `tf.defun` or `tf.function`? * `tf.compat.v1.wrap_function` or `tf.compat.v1.defun` or `tf.compat.v1.function` or `tf.compat.v1.wrap_graph_as_function`? - - From 2b712e59cf572ccf4c463519b0e062ad3c48bbe8 Mon Sep 17 00:00:00 2001 From: Asim Shankar Date: Thu, 20 Sep 2018 13:42:46 -0700 Subject: [PATCH 05/15] Fix the "Status" column --- rfcs/20180918-functions-not-sessions-20.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/rfcs/20180918-functions-not-sessions-20.md b/rfcs/20180918-functions-not-sessions-20.md index 8bff1ed3a..8df14431a 100644 --- a/rfcs/20180918-functions-not-sessions-20.md +++ b/rfcs/20180918-functions-not-sessions-20.md @@ -1,6 +1,6 @@ # TensorFlow 2.0: Functions, not Sessions. -| Status | (Proposed / Accepted / Implemented / Obsolete) | +| Status | Proposed | :-------------- |:---------------------------------------------------- | | **Author(s)** | ashankar@google.com, joshl@google.com | | **Sponsor** | apassos@google.com | From bad3fdbd1776518c5b7cda8db0a67d8b21e43257 Mon Sep 17 00:00:00 2001 From: Asim Shankar Date: Thu, 20 Sep 2018 17:06:14 -0700 Subject: [PATCH 06/15] Incorporate some suggestions. --- rfcs/20180918-functions-not-sessions-20.md | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/rfcs/20180918-functions-not-sessions-20.md b/rfcs/20180918-functions-not-sessions-20.md index 8df14431a..874390619 100644 --- a/rfcs/20180918-functions-not-sessions-20.md +++ b/rfcs/20180918-functions-not-sessions-20.md @@ -842,7 +842,17 @@ df(x, y) # Will be 1.0 ``` -This situation can be improved with the help of [autograph](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/autograph) to allow expression of control flow in Python. Whether autograph will be enabled by default or not is still under debate, but the option will be there as a flag on defun. +This situation can be improved with the help of [autograph](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/autograph) to allow expression of control flow in Python. Whether autograph will be enabled by default or not is still under debate, but the option will be there as a flag on defun. For example: + +```python +@tf.defun(autograph=True) +def f(x, y): +  if tf.equal(y, 0.0): +    return y +  return x / y + +f(tf.constant(2.0), tf.constant(2.0)) # Will be 1.0 +``` ### Summaries @@ -915,7 +925,8 @@ NOTE: In TensorFlow 1.x, eager execution is enabled by [tf.enable_eager_ex * The model_fn of a regular (non-canned) Estimator will remain as a graph construction function. * [SavedModel](https://www.tensorflow.org/guide/saved_model#save_and_restore_models) will continue to be the format encouraged for exporting trained models * Crudely speaking, a SavedModel encapsulates a Graph, a checkpoint of variable values, and some metadata like signature information (names of input and output tensors). - * A path will be provided to easily export models in this format (e.g., via tf.keras.Model.save()). There may be instances where converting the Python code to a graph is not trivial, in which case, exporting to a SavedModel (and thus a Graph) will fail. + * A path will be provided to easily export models in this format (e.g., via tf.keras.Model.save()). There may be instances where converting the Python code to a graph is not trivial (e.g., it uses the subset of Python that [autograph](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/autograph) does not support), in which case, exporting to a SavedModel (and thus a Graph) will fail. + ## Alternatives Considered From 10f4e62cefa4aa41f59899694162a42cefbe81ad Mon Sep 17 00:00:00 2001 From: Asim Shankar Date: Thu, 20 Sep 2018 17:07:18 -0700 Subject: [PATCH 07/15] Shorten the autograph example --- rfcs/20180918-functions-not-sessions-20.md | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/rfcs/20180918-functions-not-sessions-20.md b/rfcs/20180918-functions-not-sessions-20.md index 874390619..2e84f88dd 100644 --- a/rfcs/20180918-functions-not-sessions-20.md +++ b/rfcs/20180918-functions-not-sessions-20.md @@ -845,13 +845,8 @@ df(x, y) # Will be 1.0 This situation can be improved with the help of [autograph](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/autograph) to allow expression of control flow in Python. Whether autograph will be enabled by default or not is still under debate, but the option will be there as a flag on defun. For example: ```python -@tf.defun(autograph=True) -def f(x, y): -  if tf.equal(y, 0.0): -    return y -  return x / y - -f(tf.constant(2.0), tf.constant(2.0)) # Will be 1.0 +df = tf.defun(f, autograph=True) +f(x, y) # Will be 1.0 ``` From 161decbdb4a16f953da83554ab02e6be91e9d4b0 Mon Sep 17 00:00:00 2001 From: Asim Shankar Date: Fri, 21 Sep 2018 14:53:11 -0700 Subject: [PATCH 08/15] Correct autograph link. --- rfcs/20180918-functions-not-sessions-20.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/rfcs/20180918-functions-not-sessions-20.md b/rfcs/20180918-functions-not-sessions-20.md index 2e84f88dd..c549f9758 100644 --- a/rfcs/20180918-functions-not-sessions-20.md +++ b/rfcs/20180918-functions-not-sessions-20.md @@ -842,7 +842,7 @@ df(x, y) # Will be 1.0 ``` -This situation can be improved with the help of [autograph](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/autograph) to allow expression of control flow in Python. Whether autograph will be enabled by default or not is still under debate, but the option will be there as a flag on defun. For example: +This situation can be improved with the help of [autograph](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/python/autograph) to allow expression of control flow in Python. Whether autograph will be enabled by default or not is still under debate, but the option will be there as a flag on defun. For example: ```python df = tf.defun(f, autograph=True) @@ -920,7 +920,7 @@ NOTE: In TensorFlow 1.x, eager execution is enabled by [tf.enable_eager_ex * The model_fn of a regular (non-canned) Estimator will remain as a graph construction function. * [SavedModel](https://www.tensorflow.org/guide/saved_model#save_and_restore_models) will continue to be the format encouraged for exporting trained models * Crudely speaking, a SavedModel encapsulates a Graph, a checkpoint of variable values, and some metadata like signature information (names of input and output tensors). - * A path will be provided to easily export models in this format (e.g., via tf.keras.Model.save()). There may be instances where converting the Python code to a graph is not trivial (e.g., it uses the subset of Python that [autograph](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/autograph) does not support), in which case, exporting to a SavedModel (and thus a Graph) will fail. + * A path will be provided to easily export models in this format (e.g., via tf.keras.Model.save()). There may be instances where converting the Python code to a graph is not trivial (e.g., it uses the subset of Python that [autograph](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/python/autograph) does not support), in which case, exporting to a SavedModel (and thus a Graph) will fail. From c3fe0351a14a763998a3d052e9dd3efc4ce4ba9b Mon Sep 17 00:00:00 2001 From: Asim Shankar Date: Tue, 2 Oct 2018 13:10:39 -0700 Subject: [PATCH 09/15] Additional details on Trace Caches and Input Signatures. --- rfcs/20180918-functions-not-sessions-20.md | 228 ++++++++++++++++----- 1 file changed, 174 insertions(+), 54 deletions(-) diff --git a/rfcs/20180918-functions-not-sessions-20.md b/rfcs/20180918-functions-not-sessions-20.md index c549f9758..bda9abed9 100644 --- a/rfcs/20180918-functions-not-sessions-20.md +++ b/rfcs/20180918-functions-not-sessions-20.md @@ -4,14 +4,12 @@ :-------------- |:---------------------------------------------------- | | **Author(s)** | ashankar@google.com, joshl@google.com | | **Sponsor** | apassos@google.com | -| **Updated** | 2018-09-18 | +| **Updated** | 2018-10-02 | ## Objective This document presents a proposal to make TensorFlow be more "Pythonic" in 2.0. In five bullet points, the proposal is to: - - * Encourage the encapsulation of graph computation as Python functions \ (where the graph is executed when the function is invoked, instead of via `Session`) * Align "state" in the TensorFlow runtime (e.g., resource tensors like those that back `tf.Variable` objects) with state in the Python program (e.g., Python objects corresponding to the runtime state with lifetimes attached to each other). @@ -90,28 +88,12 @@ Having the Python function correspond to what the runtime will execute reduces c * `f` is a Python function that returns zero or more `Tensor`s -* `defun(f)` is a Python function that returns a Python callable, `C` -* When the `C` is invoked it: - 1. Determines an "input signature" \ -If an input signature was not explicitly specified by the user (as an argument to `defun`), the signature is computed from the types of the input arguments (including `dtype` and `shape` for `Tensor` arguments) - 1. Every time a new input signature is encountered, it invokes `f` to create a TensorFlow graph, `G`. If the input signature has been seen before, it looks up `G` from a cache keyed by the input signature. - 1. It executes the graph defined by `G,` feeding each argument as a value of the corresponding `Placeholder` node in the graph. - -Changes in input signature result in a new graph being traced. For example: - - -```python -@tf.defun -def f(x): - one = tf.constant(1, dtype=x.dtype) - return tf.add(x, one) - -# Traces a graph with int32 operations and executes it -f(tf.constant(1, dtype=tf.int32)) -# Traces a graph with float32 operations and executes it. -f(tf.constant(1, dtype=tf.float32)) -``` - +* `defun(f)` is a Python function that returns a Python callable, `F` +* When `F` is invoked it: + 1. Potentially casts inputs to tensors if an input signature was specified, see the "Input Signatures" section below. + 1. Determines a "trace_cache_key" (based on the types and/or values of the arguments). + 1. Every time a new trace_cache_key is encountered, it invokes `f` to create a TensorFlow graph, `G`. If the trace_cache_key has been seen before, it looks up `G` from a cache. + 1. It executes the graph defined by `G,` feeding each argument as a value of the corresponding node in the graph, and returns a tuple of `Tensor`s (or list of `Tensor`s). ### Referencing state: Variables, tables etc. @@ -270,17 +252,30 @@ If any variables are created in the first execution of `f`, then `@tf.defun` wil In the future we may want to allow for function local `tf.Variable`s, which are created and destroyed each time the decorated function is invoked. -### API for `defun` +### Trace Caches -We've introduced a single new symbol: `defun` that consumes a Python function and returns a callable Python object. The precise API of the object is being iterated on, but at a high level it will have methods to: +Every argument to a `defun` decorated Python function (`F`) must be either: -* List out all captured state (`tf.Variable` objects, other `DT_RESOURCE` tensors used by the computation and provided as implicit inputs). -* Access the `tf.Graph` that corresponds to the graph executed by the `__call__` method of the object. -* Execute the function with custom `RunOptions` and retrieve `RunMetadata`. +* A `Tensor` object (NumPy `ndarray`s are converted to the equivalent `Tensor`), or +* A list of `Tensor` objects, or +* An arbitrary Python value. + +(There seems to be some interest expressed in supporting structured inputs using [nest.flatten](https://github.com/tensorflow/tensorflow/blob/ed7ae86228c58e0a32f0dc21aedc9dad62db97c7/tensorflow/python/util/util.i#L77) and nest.pack_sequence_as. This will be considered as follow-up work.) + +Every time `F` is invoked in the Python program, a `trace_cache_key` is computed as a function of: + -Since new graphs are traced when new input signatures are encountered, a `defun` can encapsulate multiple graphs. For example, consider the following snippet: + +1. The element datatype and shape of every `Tensor` argument +1. The length of the list, and (dtype, shape) of every element in the list of `Tensor` argument +1. The concrete value of non-`Tensor` (and list of `Tensor`) Python object arguments +1. The "context" in which `F` is invoked (e.g., the device prescribed by the `tf.device()` scope in which `F` is invoked). + +This key is used to determine if a new graph needs to be created or if a previously created graph can be invoked. + +Since new graphs are traced when new input signatures are encountered, a `defun` can encapsulate multiple graphs. For example, consider the following: ```python @@ -288,15 +283,24 @@ Since new graphs are traced when new input signatures are encountered, a `defun` def f(x): return tf.square(x) -f(int(1)) -f(float(1.0)) +f(tf.constant(1, dtype=tf.int32)) +f(tf.constant(1.0, dtype=tf.float32)) ``` \ There are two graphs created here - one which corresponds to the `Square` operation applied to `DT_INT32` tensors, and one with the `Square` operation applied to `DT_FLOAT32` tensors. The object returned by `defun` encapsulates multiple graphs (lazily generated based on the type and shape of input arguments), multiplexing between them in `__call__`. -The same holds for the case where arguments are not `Tensor`s, for example: +Note the use of `tf.constant` to ensure that the argument is a `Tensor`. If the argument were a Python value, then additional graphs will be traced for each such value. For example, the following two calls will result in two additional graphs being traced: + + +```python +f(1.0) +f(2.0) +``` + + +Where arguments are not `Tensor`s, the "value" of the argument is used to compute the `trace_cache_key`. For example: ```python @@ -304,12 +308,12 @@ The same holds for the case where arguments are not `Tensor`s, for example: def f(x, use_multiply): return tf.multiply(x, x) if use_multiply else tf.square(x) -f(2.0, True) -f(2.0, False) +f(tf.constant(2.0), True) +f(tf.constant(2.0), False) ``` -will result in 2 graphs being created. +will result in 2 graphs being created, since the two calls result in two different cache keys because the value of the Python object (the second argument) changes between the two. Note that the "type" of `Tensor` inputs to the function also incorporates the shape. For example: @@ -317,9 +321,11 @@ Note that the "type" of `Tensor` inputs to the function also incorporates the sh ```python @tf.defun def f(x): return tf.add(x, 1.) -f([2.0]) -f([2.0, 3.0]) -f([[2.0]]) +f(tf.constant([2.0])) +f(tf.constant([2.0, 3.0])) +f(tf.constant([[2.0]])) +f(tf.constant([3.0])) +f(tf.constant([4.0, 5.0])) ``` @@ -327,30 +333,132 @@ will result in 3 graphs being created: -1. One for when the first argument is a `tf.float32` vector with 1 element \ -(input signature: `((tf.float32, [1]))`) -1. One for when the first argument is a `tf.float32` vector with 2 elements \ -(input signature: `((tf.float32, [2]]))`) -1. One for when the first argument is a `tf.float32` 1x1 matrix \ -(input signature: `((tf.float32, [1, 1]))`) +1. One for when the first argument is a `tf.float32` vector with 1 element +1. One for when the first argument is a `tf.float32` vector with 2 elements +1. One for when the first argument is a `tf.float32` 1x1 matrix + +The trace_cache_key also incorporates the "context" in which the call was made. For example: + + +```python +@tf.defun +def f(x): return tf.add(x, 1.) + +with tf.device("/device:CPU:0"): + f(tf.constant(2.0)) +with tf.device("/device:GPU:0"): + f(tf.constant(2.0)) +``` + + +Will create 2 graphs, one where the operations are pinned to the CPU device and one where they are pinned to the GPU device. + + +#### CAUTION: Too many traces + +Since new traces are generated on demand, the object returned by `defun` may hold on to more resources than the user may realize. Possible mitigations: + + + +* Garbage collect the graphs when the weak reference to any component of the `trace_cache_key` is no longer alive. +* Use input signatures to prevent unnecessary retraces (see "Input Signatures" section below) +* Raise / log an error when the ratio of calls to traces is greater than some threshold (e.g., if every 2 calls to a `defun` decorated function generates a new graph). + + +#### CAUTION: Mutable non-`Tensor` arguments + +The trace_cache_key includes the Python object for non-`Tensor` arguments. Mutations of these arguments might not be detected. For example: + + +```python +class Params(object): + multiply = True + +p = Params() +@tf.defun +def f(x, y): + return tf.multiply(x, 2.) if y.multiply else tf.add(x, 2.) + +f(3., p) # Returns 6.0 +y.multiply = False +f(3., y) # Mutations to `y` may not trigger a retrace, so might still return 6.0 +``` + + + +### Input Signatures + +Tracing the decorated function to create a new graph on each input shape is a conservative choice. Often the same graph suffices for `Tensor`s of multiple shapes. As a trivial example, consider: + + +```python +@tf.defun +def f(x): return tf.add(x, 1.) + +f(tf.constant(1.0)) # Scalar argument +f(tf.constant([1.0, 2.0])) # Vector argument +f(tf.constant([[3.0]])) # Matrix +``` -Tracing the decorated function to create a new graph on each input shape is a conservative choice (allowing for `f` to create graphs dependent on the shape), which may be unnecessary. Users can explicitly specify an input signature to ensure that the same graph is used for multiple inputs. For example: + + \ +This snippet would result in 3 graphs being traced. An "input signature" can be explicitly specified to control the `trace_cache_key` computation based on the type and shape of `Tensor` (and list of `Tensor`) arguments to `f`. + +For example: ```python @tf.defun(input_signature=((tf.float32, [None])) def f(x): return tf.add(x, 1.) -f([2.0]) # Returns [3.0] -f([2.0, 3.0]) # Matches the input signature as [None] matches the actual shape [2] -f([[2.0]]) # Raises an error as the arguments don't match the input signature. +f(tf.constant([2.0])) # Returns [3.0] +f(tf.constant([2.0, 3.0])) # Matches the input signature as [None] + # matches the actual shape [2] +f(tf.constant([[2.0]])) # Raises an error as the arguments don't match the + # input signature. +f(tf.constant([2], dtype=tf.int32)) # Raises an error as the dtype of the argument + # does not match the input signature # f is backed by a single Graph since the input signature specification allowed # for the same graph to be used when the input shape is (1,) or (2,). ``` - +An "input signature" specifies a pattern for each of the arguments that may be accepted by the `defun`-decorated function. Specifically: + + + +* For a `Tensor` argument, it specifies a (dtype, shape pattern). \ +For example: + * `(tf.float32, [None])` means the argument must be a float32 vector (with any number of elements). + * `(tf.int32, [])` means that the argument must be an int32 scalar. \ + \ +In this case, non-`Tensor` Python values provided at call time are automatically converted (using `tf.convert_to_tensor`) to a `Tensor` matching this signature. +* For a list of `Tensor` objects, it specifies an optional list length and the signature for elements in the list (i.e., the dtype and shape pattern for all elements in the list). +* For non-`Tensor` arguments: `tf.PYTHON_VALUE` + +When an input signature is specified, new graphs are traced only when the value of the Python argument or the context in which the function is invoked changes. If this is considered to be too restrictive, a possible future extension would be to annotate signature of an argument so that new traces can be created. For example: + + +```python +@tf.defun(input_signature=((tf.TRACE_ON_NEW_VALUE, [None])) +def f(x): return tf.square(x) + +f(tf.constant([2.0])) # Returns 4.0 +f(tf.constant([2, 2], dtype=tf.int32) # Returns [4, 4] after tracing a new graph +``` + + + +### API for `defun` + +We've introduced a single new symbol: `defun` that consumes a Python function and returns a callable Python object. The precise API of the object is being iterated on in go/tf-2.0-function-api, but at a high level it will have methods to: + + + +* List out all captured state (`tf.Variable` objects, other `DT_RESOURCE` tensors used by the computation and provided as implicit inputs). +* Access the `tf.Graph` that corresponds to the graph executed by the `__call__` method of the object. +* Execute the function with custom `RunOptions` and retrieve `RunMetadata`. ### Classes @@ -798,7 +906,7 @@ g(2.0) # 16.0
At the lowest level of the API, distributed execution continues to work with `tf.device` annotations, where the device name can reference remote devices as well, just like they do today. -The `DistributionStrategy` API, typically aimed at synchronous training will continue to be the method of choice (where the API can be used inside a `defun`). +The `DistributionStrategy` API, typically aimed at synchronous training will continue to be the method of choice (where the API can be used inside a `defun`). Other APIs such as go/tf-replicator will also be usable. The author realizes that this section can do with more detail. However, to keep this document more focused, these details will be discussed separately. In particular, usage of `MonitoredSession` and session hooks today needs additional thought. @@ -844,12 +952,14 @@ df(x, y) # Will be 1.0 This situation can be improved with the help of [autograph](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/python/autograph) to allow expression of control flow in Python. Whether autograph will be enabled by default or not is still under debate, but the option will be there as a flag on defun. For example: + ```python df = tf.defun(f, autograph=True) f(x, y) # Will be 1.0 ``` + ### Summaries The summary writing operations ([tb.summary.scalar](https://www.tensorflow.org/api_docs/python/tf/contrib/summary/scalar), [tb.summary.image](https://www.tensorflow.org/api_docs/python/tf/contrib/summary/image) etc.) can be naturally placed in the graph by using them in a defun-decorated function. These operations require two "external" inputs - the summary writer resource and the condition, that will be picked up from the context (e.g., [tb.summary.create_file_writer](https://www.tensorflow.org/api_docs/python/tf/contrib/summary/create_file_writer) and [tb.summary.record_summary_every_n_global_steps](https://www.tensorflow.org/api_docs/python/tf/contrib/summary/record_summary_every_n_global_steps)). When defining the graph, these inputs are converted to placeholders, which are then resolved at function invocation time. Thus, something like this: @@ -923,7 +1033,6 @@ NOTE: In TensorFlow 1.x, eager execution is enabled by [tf.enable_eager_ex * A path will be provided to easily export models in this format (e.g., via tf.keras.Model.save()). There may be instances where converting the Python code to a graph is not trivial (e.g., it uses the subset of Python that [autograph](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/python/autograph) does not support), in which case, exporting to a SavedModel (and thus a Graph) will fail. - ## Alternatives Considered @@ -999,10 +1108,21 @@ This seems like an avenue definitely worth pursuing, but requires careful consid For now, we propose that `defun` continue with the restricted abilities proposed in this document and a "maintain Python semantics" decorator be investigated independently. -## Open Questions +## Open Questions/Ideas * Naming: * `tf.defun` or `tf.function`? * `tf.compat.v1.wrap_function` or `tf.compat.v1.defun` or `tf.compat.v1.function` or `tf.compat.v1.wrap_graph_as_function`? +* Signatures in Python 3? ([From ngc92](https://github.com/tensorflow/community/pull/20#issuecomment-423345326)) +* Can `tf.defun` and `tf.method` be combined into a single decorator (where `tf.defun` has the behavior of `tf.method` when applied to a class method)? Or is it okay to have two separate decorators? + * (How do you detect if the decorated function is a method or a function? Rely on the convention of first argument being called `self`?) +* Supporting structured inputs: \ +As proposed, arguments to `defun` must be either `Tensor` objects, or objects that can be converted to a `Tensor` (`tf.convert_to_tensor`), or opaque Python objects. \ + \ +Perhaps we can support nested structures of `Tensor`s (using `nest.flatten` and `nest.pack_sequence_as`), or even arbitrary Python objects? \ + \ +If this is supported, then specifying an `input_signature` may become cumbersome, but perhaps we can have a `defun(infer_signature_from_first_call=True)` to make that easier. \ + \ + From 33f380d6192592aef4803480cdc68df83e36b062 Mon Sep 17 00:00:00 2001 From: Asim Shankar Date: Tue, 9 Oct 2018 14:12:58 -0700 Subject: [PATCH 10/15] s/tf.defun/tf.function/ --- rfcs/20180918-functions-not-sessions-20.md | 143 ++++++++++----------- 1 file changed, 70 insertions(+), 73 deletions(-) diff --git a/rfcs/20180918-functions-not-sessions-20.md b/rfcs/20180918-functions-not-sessions-20.md index bda9abed9..b3d61b46c 100644 --- a/rfcs/20180918-functions-not-sessions-20.md +++ b/rfcs/20180918-functions-not-sessions-20.md @@ -63,11 +63,11 @@ The core proposal of this document is the alignment between computation expresse ```python import tensorflow as tf -@tf.defun +@tf.function def compute_z1(x, y): return tf.add(x, y) -@tf.defun +@tf.function def compute_z0(x): return compute_z1(x, tf.square(x)) @@ -76,19 +76,19 @@ z1 = compute_z1(2., 2.) ``` -Where `tf.defun` is a decorator that "**de**fines a TensorFlow **fun**ction". A "TensorFlow function" defines a computation as a graph of TensorFlow operations, with named arguments and explicit return values. Users define the function they want TensorFlow to "accelerate" as a Python function and integrate it into their Python program like any other Python function call. +Where `tf.function` is a decorator that "**de**fines a TensorFlow **fun**ction". A "TensorFlow function" defines a computation as a graph of TensorFlow operations, with named arguments and explicit return values. Users define the function they want TensorFlow to "accelerate" as a Python function and integrate it into their Python program like any other Python function call. Having the Python function correspond to what the runtime will execute reduces conceptual complexity in translating between the two domains. It also affords an opportunity to provide more helpful stacktraces on errors. More advanced features available today (e.g., carving sub-graphs, feeding intermediate values) will still be possible (discussed later), though most users should not need to think in terms of graphs, feeds, and fetches. The constructed graph also provides a natural point for accelerators/acceleration libraries (NVIDIA TensorRT, Google Cloud TPUs etc.) to hook in for rewrites. -### `defun`: A brief specification +### `function`: A brief specification -`defun` constructs a TensorFlow graph by "tracing" the TensorFlow operations executed by the Python function. Specifically: +`function` constructs a TensorFlow graph by "tracing" the TensorFlow operations executed by the Python function. Specifically: * `f` is a Python function that returns zero or more `Tensor`s -* `defun(f)` is a Python function that returns a Python callable, `F` +* `function(f)` is a Python function that returns a Python callable, `F` * When `F` is invoked it: 1. Potentially casts inputs to tensors if an input signature was specified, see the "Input Signatures" section below. 1. Determines a "trace_cache_key" (based on the types and/or values of the arguments). @@ -98,7 +98,7 @@ Having the Python function correspond to what the runtime will execute reduces c ### Referencing state: Variables, tables etc. -A `defun` decorated Python function encapsulates a graph and its execution. The Python function may reference stateful objects (i.e., state backed by `DT_RESOURCE` tensors in the runtime, e.g., `tf.Variable`) by referencing the corresponding Python object, and these will be captured as implicit inputs to the function. +A `function` decorated Python function encapsulates a graph and its execution. The Python function may reference stateful objects (i.e., state backed by `DT_RESOURCE` tensors in the runtime, e.g., `tf.Variable`) by referencing the corresponding Python object, and these will be captured as implicit inputs to the function. Comparing TensorFlow code today with how we propose it looks in 2.x: @@ -146,7 +146,7 @@ with tf.Session() as sess: b = tf.Variable(tf.zeros(10)) c = tf.Variable(0) -@tf.defun +@tf.function def f(x): c.assign_add(1) return tf.matmul(x, W) + b @@ -183,12 +183,12 @@ with tf.Session() as sess: ``` -The output here is not deterministic, since `val` may evaluate to either 1.0 or 2.0 depending on whether the runtime happened to execute `assign_op` before `y` or not. `tf.control_dependencies` is a mechanism provided to add annotations at graph construction time to influence graph execution. The TensorFlow user, a Python programmer, is thus forced to think about two execution models - TensorFlow graphs and the Python interpreter. To eliminate this cognitive load, `defun` will automatically insert control dependencies to ensure that (1) operations that produce or consume a given `DT_RESOURCE` tensor and (2) operations that are marked stateful (`REGISTER_OP(...).SetIsStateful()`) follow graph construction order. Thus: +The output here is not deterministic, since `val` may evaluate to either 1.0 or 2.0 depending on whether the runtime happened to execute `assign_op` before `y` or not. `tf.control_dependencies` is a mechanism provided to add annotations at graph construction time to influence graph execution. The TensorFlow user, a Python programmer, is thus forced to think about two execution models - TensorFlow graphs and the Python interpreter. To eliminate this cognitive load, `function` will automatically insert control dependencies to ensure that (1) operations that produce or consume a given `DT_RESOURCE` tensor and (2) operations that are marked stateful (`REGISTER_OP(...).SetIsStateful()`) follow graph construction order. Thus: ```python v = tf.Variable(1.0) -@tf.defun +@tf.function def f():   v.assign(2.0)   return v.read_value() @@ -203,7 +203,7 @@ Note that the intention here is to avoid _observable_ differences from program o ```python a = tf.Variable(1.0) b = tf.Variable(1.0) -@tf.defun +@tf.function def f(): a.assign(2.0) b.assign(3.0) @@ -219,7 +219,7 @@ A preview of this implemented in `tf.contrib.eager.defun` today (using [Au ### Functions that create state -In the above code, no `tf.Variable` objects are created inside a `tf.defun` decorated function. This is makes it clear that the code will have the same semantics once wrapped. +In the above code, no `tf.Variable` objects are created inside a `tf.function` decorated function. This is makes it clear that the code will have the same semantics once wrapped. Note that if the function naturally creates state only on the first trace, all is well: @@ -227,7 +227,7 @@ Note that if the function naturally creates state only on the first trace, all i ```python v = None -@tf.defun +@tf.function def f(x): global v if v is None: @@ -239,22 +239,22 @@ f(tf.constant(2, dtype=tf.int32)) # Reuses the variable, returns 3.0 ``` -To support this `defun` imposes some requirements on the decorated function: +To support this `function` imposes some requirements on the decorated function: 1. State (like `tf.Variable` objects) are only created the first time the function `f` is called. \ How that is accomplished is left up to the implementation of `f`. \ -If any variables are created in the first execution of `f`, then `@tf.defun` will trace `f` a second time in order to record the behavior that will be used from then on. No variables may be created during that second trace, or any other trace after that (due to different dtypes, shapes, or non-tensor arguments). +If any variables are created in the first execution of `f`, then `@tf.function` will trace `f` a second time in order to record the behavior that will be used from then on. No variables may be created during that second trace, or any other trace after that (due to different dtypes, shapes, or non-tensor arguments). 1. The caller must make sure that any variable referenced by the function still exists whenever the function is evaluated. \ -`@tf.defun` itself will keep only weak references to these created variables. Thus, if the referenced state does not exist when the decorated function is invoked, an exception will be raised. +`@tf.function` itself will keep only weak references to these created variables. Thus, if the referenced state does not exist when the decorated function is invoked, an exception will be raised. In the future we may want to allow for function local `tf.Variable`s, which are created and destroyed each time the decorated function is invoked. ### Trace Caches -Every argument to a `defun` decorated Python function (`F`) must be either: +Every argument to a `function` decorated Python function (`F`) must be either: @@ -275,11 +275,11 @@ Every time `F` is invoked in the Python program, a `trace_cache_key` is computed This key is used to determine if a new graph needs to be created or if a previously created graph can be invoked. -Since new graphs are traced when new input signatures are encountered, a `defun` can encapsulate multiple graphs. For example, consider the following: +Since new graphs are traced when new input signatures are encountered, a `function` can encapsulate multiple graphs. For example, consider the following: ```python -@tf.defun +@tf.function def f(x): return tf.square(x) @@ -289,7 +289,7 @@ f(tf.constant(1.0, dtype=tf.float32)) \ -There are two graphs created here - one which corresponds to the `Square` operation applied to `DT_INT32` tensors, and one with the `Square` operation applied to `DT_FLOAT32` tensors. The object returned by `defun` encapsulates multiple graphs (lazily generated based on the type and shape of input arguments), multiplexing between them in `__call__`. +There are two graphs created here - one which corresponds to the `Square` operation applied to `DT_INT32` tensors, and one with the `Square` operation applied to `DT_FLOAT32` tensors. The object returned by `function` encapsulates multiple graphs (lazily generated based on the type and shape of input arguments), multiplexing between them in `__call__`. Note the use of `tf.constant` to ensure that the argument is a `Tensor`. If the argument were a Python value, then additional graphs will be traced for each such value. For example, the following two calls will result in two additional graphs being traced: @@ -304,7 +304,7 @@ Where arguments are not `Tensor`s, the "value" of the argument is used to comput ```python -@tf.defun +@tf.function def f(x, use_multiply): return tf.multiply(x, x) if use_multiply else tf.square(x) @@ -319,7 +319,7 @@ Note that the "type" of `Tensor` inputs to the function also incorporates the sh ```python -@tf.defun +@tf.function def f(x): return tf.add(x, 1.) f(tf.constant([2.0])) f(tf.constant([2.0, 3.0])) @@ -341,7 +341,7 @@ The trace_cache_key also incorporates the "context" in which the call was made. ```python -@tf.defun +@tf.function def f(x): return tf.add(x, 1.) with tf.device("/device:CPU:0"): @@ -356,13 +356,13 @@ Will create 2 graphs, one where the operations are pinned to the CPU device and #### CAUTION: Too many traces -Since new traces are generated on demand, the object returned by `defun` may hold on to more resources than the user may realize. Possible mitigations: +Since new traces are generated on demand, the object returned by `function` may hold on to more resources than the user may realize. Possible mitigations: * Garbage collect the graphs when the weak reference to any component of the `trace_cache_key` is no longer alive. * Use input signatures to prevent unnecessary retraces (see "Input Signatures" section below) -* Raise / log an error when the ratio of calls to traces is greater than some threshold (e.g., if every 2 calls to a `defun` decorated function generates a new graph). +* Raise / log an error when the ratio of calls to traces is greater than some threshold (e.g., if every 2 calls to a `function` decorated function generates a new graph). #### CAUTION: Mutable non-`Tensor` arguments @@ -375,7 +375,7 @@ class Params(object): multiply = True p = Params() -@tf.defun +@tf.function def f(x, y): return tf.multiply(x, 2.) if y.multiply else tf.add(x, 2.) @@ -392,7 +392,7 @@ Tracing the decorated function to create a new graph on each input shape is a co ```python -@tf.defun +@tf.function def f(x): return tf.add(x, 1.) f(tf.constant(1.0)) # Scalar argument @@ -408,7 +408,7 @@ For example: ```python -@tf.defun(input_signature=((tf.float32, [None])) +@tf.function(input_signature=((tf.float32, [None])) def f(x): return tf.add(x, 1.) f(tf.constant([2.0])) # Returns [3.0] @@ -424,7 +424,7 @@ f(tf.constant([2], dtype=tf.int32)) # Raises an error as the dtype of the argume ``` -An "input signature" specifies a pattern for each of the arguments that may be accepted by the `defun`-decorated function. Specifically: +An "input signature" specifies a pattern for each of the arguments that may be accepted by the `function`-decorated function. Specifically: @@ -441,7 +441,7 @@ When an input signature is specified, new graphs are traced only when the value ```python -@tf.defun(input_signature=((tf.TRACE_ON_NEW_VALUE, [None])) +@tf.function(input_signature=((tf.TRACE_ON_NEW_VALUE, [None])) def f(x): return tf.square(x) f(tf.constant([2.0])) # Returns 4.0 @@ -450,9 +450,9 @@ f(tf.constant([2, 2], dtype=tf.int32) # Returns [4, 4] after tracing a new graph -### API for `defun` +### API for `function` -We've introduced a single new symbol: `defun` that consumes a Python function and returns a callable Python object. The precise API of the object is being iterated on in go/tf-2.0-function-api, but at a high level it will have methods to: +We've introduced a single new symbol: `function` that consumes a Python function and returns a callable Python object. The precise API of the object is being iterated on in go/tf-2.0-function-api, but at a high level it will have methods to: @@ -463,7 +463,7 @@ We've introduced a single new symbol: `defun` that consumes a Python function an ### Classes -If a member function of a class does not create variables, it may be decorated with `@tf.defun` and it will work: +If a member function of a class does not create variables, it may be decorated with `@tf.function` and it will work: ```python @@ -471,7 +471,7 @@ class ScalarModel(object): def __init__(self): self.v = tf.Variable(0) - @tf.defun + @tf.function def increment(self, amount): self.v.assign_add(amount) @@ -512,16 +512,16 @@ assert model2.v.numpy() == [4, 5] ``` -The semantics here are that each new instance is allowed to create variables in each `@tf.method` once. The simple recommendation would be "always use `@tf.method` on methods, use `@tf.defun` for functions outside of a class". +The semantics here are that each new instance is allowed to create variables in each `@tf.method` once. The simple recommendation would be "always use `@tf.method` on methods, use `@tf.function` for functions outside of a class". -In the above example, if `increment` was decorated with `@tf.defun` instead, then the `model2.increment()` call would raise an exception (as per `defun`s stated behavior of disallowing state creation on anything but the first trace). However, if the method didn't create any state then `@tf.defun` or `@tf.method` would both have the same effect. +In the above example, if `increment` was decorated with `@tf.function` instead, then the `model2.increment()` call would raise an exception (as per `function`s stated behavior of disallowing state creation on anything but the first trace). However, if the method didn't create any state then `@tf.function` or `@tf.method` would both have the same effect. In addition, as long as all variable creation/initialization happens while we are tracing, we should be able to support exporting the initialization graph when exporting a `SavedModel` or `MetaGraphDef`. ### Transitioning from 1.x -The definition of `tf.defun` above is careful to check that invoking a decorated Python function would have the same behavior as invoking an undecorated function. This is to guard against it being passed code from TensorFlow v1.x that expects to only be called once (and relies on things like graph collections to track which variables are created), for example: +The definition of `tf.function` above is careful to check that invoking a decorated Python function would have the same behavior as invoking an undecorated function. This is to guard against it being passed code from TensorFlow v1.x that expects to only be called once (and relies on things like graph collections to track which variables are created), for example: ```python @@ -553,7 +553,7 @@ assert float(f_sub(1.0)) == 3.0 ``` -Note these differences from `tf.defun`: +Note these differences from `tf.function`: @@ -565,7 +565,7 @@ Note these differences from `tf.defun`: * Keeps strong references to variables created in f, weak references to other variables accessed by f. This is to match the v1.x graph behavior that variables have the lifetime of the graph they are created, and can generally be accessed through graph collections. Some common patterns of writing v1.x code don't leave any references to those variables around. Keeping references to those variables extends their lifetime to match that of the object returned by `tf.compat.v1.wrap_function`. * Typically won't be used as a decorator. Calling `tf.compat.v1.wrap_function` takes some arguments, traces the function, and creates an object with state. The lifetime of the return value should be tracked explicitly by saving it in a variable. -Treating state (like `tf.Variable`) as static local does mean that the behavior of a `tf.compat.v1.wrap_function`-decorated Python function differs from that of an undecorated one. In the above example, `f(1.0, True)` will always return 6.0 (as a scalar `Tensor`), while each call to `f_add(1.0)` will return a different value. We propose this separate `tf.compat.v1.wrap_function` endpoint specifically to make it easy to migrate TensorFlow 1.x libraries to the TensorFlow 2.0. The behavior of 2.0 `tf.defun` is restricted to cases where we can say that the behavior will match. +Treating state (like `tf.Variable`) as static local does mean that the behavior of a `tf.compat.v1.wrap_function`-decorated Python function differs from that of an undecorated one. In the above example, `f(1.0, True)` will always return 6.0 (as a scalar `Tensor`), while each call to `f_add(1.0)` will return a different value. We propose this separate `tf.compat.v1.wrap_function` endpoint specifically to make it easy to migrate TensorFlow 1.x libraries to the TensorFlow 2.0. The behavior of 2.0 `tf.function` is restricted to cases where we can say that the behavior will match. We recognize that code written for TensorFlow 1.x commonly does not encapsulate state in Python objects, instead adding to hidden (graph-)global collections. We will support code that accesses collections inside a `tf.compat.v1.wrap_function`, though those collections will be local to a single trace. @@ -595,7 +595,7 @@ assert f.variables[0].name == "weight" \ In this case, the object returned by `tf.compat.v1.wrap_function` owns the state created within `f`, and the `__call__` method on it invokes the corresponding computation. -Long story short, `tf.compat.v1.wrap_function` helps in incorporating graph construction code written against TensorFlow 1.x into TensorFlow 2.x programs. `wrap_function` constructs the same object as a `defun` decorated function, which provides the conceptual equivalent of graph construction and `Session.run`. +Long story short, `tf.compat.v1.wrap_function` helps in incorporating graph construction code written against TensorFlow 1.x into TensorFlow 2.x programs. `wrap_function` constructs the same object as a `function` decorated function, which provides the conceptual equivalent of graph construction and `Session.run`. ### Serialization: Exporting SavedModel/GraphDefs @@ -619,7 +619,7 @@ At a high level, the SavedModel format packages the `MetaGraphDef`, checkpoint, (`tf.saved_model.simple_save` / `tf.saved_model.builder.SavedModelBuilder`) \ This is the format preferred for exporting for serving via TensorFlow Serving or to other languages (e.g., `SavedModelBundle.load()` in Java, `LoadSavedModel` in Go) -The objects created by `defun` encapsulate (1) the computation expressed as a `GraphDef`, (2) the state used by it. Thus, these objects are naturally suited for import/export in any of the above formats, using something like the following: +The objects created by `function` encapsulate (1) the computation expressed as a `GraphDef`, (2) the state used by it. Thus, these objects are naturally suited for import/export in any of the above formats, using something like the following: @@ -665,7 +665,7 @@ with tf.Session() as sess: tf.glorot_uniform_initializer()( (10, 10))) -@tf.defun +@tf.function def train(): W.assign_add(1.) @@ -715,7 +715,7 @@ tf.import_graph_def(graph_def) tf.glorot_uniform_initializer()( (10, 10))) -@tf.defun +@tf.function def f(x): return tf.matmul(x, W) @@ -799,7 +799,7 @@ m = One reservation expressed by TensorFlow graph/session enthusiasts today is that the ability to write generic analysis/inspection tooling on graphs, precluding the need to understand or modify the Python code that constructed the graph, is important to them. To put it differently, some find it easier to navigate the `GraphDef` program than navigating the Python program. \ -This ability will be maintained. `defun`-decorated Python functions have an associated graph, and new functions can be created by specifying the sub-graph of interest. For example: +This ability will be maintained. `function`-decorated Python functions have an associated graph, and new functions can be created by specifying the sub-graph of interest. For example:
@@ -834,7 +834,7 @@ with tf.Session() as sess: -
@tf.defun
+
@tf.function
 def f(x):
   return tf.square(tf.square(x))
 
@@ -885,11 +885,11 @@ with tf.Session() as sess:
 
 
 
-
@tf.defun
+
@tf.function
 def f(x):
   return tf.square(x)
 
-@tf.defun
+@tf.function
 def g(x):
   return tf.square(f(x))
 
@@ -906,14 +906,14 @@ g(2.0) # 16.0
At the lowest level of the API, distributed execution continues to work with `tf.device` annotations, where the device name can reference remote devices as well, just like they do today. -The `DistributionStrategy` API, typically aimed at synchronous training will continue to be the method of choice (where the API can be used inside a `defun`). Other APIs such as go/tf-replicator will also be usable. +The `DistributionStrategy` API, typically aimed at synchronous training will continue to be the method of choice (where the API can be used inside a `function`). Other APIs such as go/tf-replicator will also be usable. The author realizes that this section can do with more detail. However, to keep this document more focused, these details will be discussed separately. In particular, usage of `MonitoredSession` and session hooks today needs additional thought. -### `defun`-ing Python control flow +### `function`-ing Python control flow -`defun` decorates a graph construction function and transparently recreates graphs if needed. However, this does mean that if the function has data-dependent control flow then though the function will execute fine with eager execution enabled, `defun` decorating it will fail. For example: +`function` decorates a graph construction function and transparently recreates graphs if needed. However, this does mean that if the function has data-dependent control flow then though the function will execute fine with eager execution enabled, `function` decorating it will fail. For example: ```python @@ -927,7 +927,7 @@ y = tf.constant(2.0) f(x, y) # Will be 1.0 -df = tf.defun(f) +df = tf.function(f) df(x, y) # Will raise an error complaining about the data-dependent control flow ``` @@ -945,16 +945,16 @@ y = tf.constant(2.0) f(x, y) # Will be 1.0 -df = tf.defun(f) +df = tf.function(f) df(x, y) # Will be 1.0 ``` -This situation can be improved with the help of [autograph](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/python/autograph) to allow expression of control flow in Python. Whether autograph will be enabled by default or not is still under debate, but the option will be there as a flag on defun. For example: +This situation can be improved with the help of [autograph](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/python/autograph) to allow expression of control flow in Python. Whether autograph will be enabled by default or not is still under debate, but the option will be there as a flag on function. For example: ```python -df = tf.defun(f, autograph=True) +df = tf.function(f, autograph=True) f(x, y) # Will be 1.0 ``` @@ -962,7 +962,7 @@ f(x, y) # Will be 1.0 ### Summaries -The summary writing operations ([tb.summary.scalar](https://www.tensorflow.org/api_docs/python/tf/contrib/summary/scalar), [tb.summary.image](https://www.tensorflow.org/api_docs/python/tf/contrib/summary/image) etc.) can be naturally placed in the graph by using them in a defun-decorated function. These operations require two "external" inputs - the summary writer resource and the condition, that will be picked up from the context (e.g., [tb.summary.create_file_writer](https://www.tensorflow.org/api_docs/python/tf/contrib/summary/create_file_writer) and [tb.summary.record_summary_every_n_global_steps](https://www.tensorflow.org/api_docs/python/tf/contrib/summary/record_summary_every_n_global_steps)). When defining the graph, these inputs are converted to placeholders, which are then resolved at function invocation time. Thus, something like this: +The summary writing operations ([tb.summary.scalar](https://www.tensorflow.org/api_docs/python/tf/contrib/summary/scalar), [tb.summary.image](https://www.tensorflow.org/api_docs/python/tf/contrib/summary/image) etc.) can be naturally placed in the graph by using them in a function-decorated function. These operations require two "external" inputs - the summary writer resource and the condition, that will be picked up from the context (e.g., [tb.summary.create_file_writer](https://www.tensorflow.org/api_docs/python/tf/contrib/summary/create_file_writer) and [tb.summary.record_summary_every_n_global_steps](https://www.tensorflow.org/api_docs/python/tf/contrib/summary/record_summary_every_n_global_steps)). When defining the graph, these inputs are converted to placeholders, which are then resolved at function invocation time. Thus, something like this: ```python @@ -1000,19 +1000,19 @@ Note that the runtime is free to prune away the summary writing operations when So far this proposal has dealt with the encapsulation of TensorFlow graphs in Python functions with the intention of making it easier to integrate TensorFlow-accelerated computation in Python programs. -_Additionally_, this proposal suggests enabling eager execution by default in TensorFlow 2.0. Keeping `defun` in mind, this basically means: +_Additionally_, this proposal suggests enabling eager execution by default in TensorFlow 2.0. Keeping `function` in mind, this basically means: -* Inside the context of defining a TensorFlow function (i.e., within a `defun` decorated function) `tf.Tensor` objects created refer to symbolic tensors. +* Inside the context of defining a TensorFlow function (i.e., within a `function` decorated function) `tf.Tensor` objects created refer to symbolic tensors. * Outside this context, `tf.Tensor` objects created are backed by concrete values and TensorFlow API. The underlying memory of the tensor can be backed by any device (i.e., CPU/GPU) and is not restricted to host-memory (like numpy arrays). See the [docstring for tf.contrib.eager.defun](https://www.tensorflow.org/api_docs/python/tf/contrib/eager/defun) - the evolving playground for the implementation of the proposal in this document. The basic takeaway is that: -* For users that embrace symbolic tensors and graphs, continue doing so with your code placed inside a `defun` decorated Python function. -* We believe most users (new ones in particular) will find it more convenient to deal with `Tensor` objects backed by concrete values and then selectively "compiling" portions of their Python program into TensorFlow graphs rather than being exposed to graph metaprogramming in Python upfront. In spirit, this is similar to Swift4TensorFlow with the obvious glaring difference that[ graph program extraction](https://github.com/tensorflow/swift/blob/master/docs/DesignOverview.md#graph-program-extraction) here is manually specified (with the `defun` decoration). +* For users that embrace symbolic tensors and graphs, continue doing so with your code placed inside a `function` decorated Python function. +* We believe most users (new ones in particular) will find it more convenient to deal with `Tensor` objects backed by concrete values and then selectively "compiling" portions of their Python program into TensorFlow graphs rather than being exposed to graph metaprogramming in Python upfront. In spirit, this is similar to Swift4TensorFlow with the obvious glaring difference that[ graph program extraction](https://github.com/tensorflow/swift/blob/master/docs/DesignOverview.md#graph-program-extraction) here is manually specified (with the `function` decoration). NOTE: In TensorFlow 1.x, eager execution is enabled by [tf.enable_eager_execution()](https://www.tensorflow.org/api_docs/python/tf/enable_eager_execution). Once invoked, all public API endpoints that consume or produce symbolic Tensor objects begin to produce and consume Tensor objects that are backed by a concrete value. See the "Research and Experimentation" section at [www.tensorflow.org/tutorials](http://www.tensorflow.org/tutorials) for an introduction. @@ -1036,9 +1036,9 @@ NOTE: In TensorFlow 1.x, eager execution is enabled by [tf.enable_eager_ex ## Alternatives Considered -### Creating state inside a `defun` +### Creating state inside a `function` -How state (`DT_RESOURCE` tensors) created inside a `defun` should be handled is actively being debated. Options include: +How state (`DT_RESOURCE` tensors) created inside a `function` should be handled is actively being debated. Options include: @@ -1048,7 +1048,7 @@ How state (`DT_RESOURCE` tensors) created inside a `defun` should be handled is #### "Static-local" state -`tf.contrib.eager.defun` today treats state as function-static variables, which allows for code like: +`tf.contrib.eager.function` today treats state as function-static variables, which allows for code like: ```python @@ -1058,7 +1058,7 @@ def f(x): return v df = tf.contrib.eager.defun(f) -# tf.defun(f) proposed in this document will raise an exception on first use +# tf.function(f) proposed in this document will raise an exception on first use x = tf.constant(1, dtype=tf.float32)) print(df(x)) # 2.0 print(df(x)) # 3.0 @@ -1075,12 +1075,12 @@ print(f(1.0), df(1.0)) # 2.0, 3.0 ``` -To be conservative, we propose some restrictions on `defun`, such as: +To be conservative, we propose some restrictions on `function`, such as: -1. State is created only once, i.e., `defun` will fail if calling `f` a second time results in new state being created. -1. `defun` decorated functions can only produce `Tensor` return values. +1. State is created only once, i.e., `function` will fail if calling `f` a second time results in new state being created. +1. `function` decorated functions can only produce `Tensor` return values. 1. If you want to convert TF v1.x code like `f` above, you may use `tf.compat.v1.wrap_function` which guarantees it will only trace `f` once. @@ -1095,7 +1095,7 @@ def f(x): v.assign_add(x) return v -df = tf.defun(f) +df = tf.function(f) assert f(1.0) == df(1.0) # Both will be 2.0 assert f(1.0) == df(1.0) # Still 2.0, since 'v' would be recreated. @@ -1105,7 +1105,7 @@ assert f(1.0) == df(1.0) # Still 2.0, since 'v' would be recreated. \ This seems like an avenue definitely worth pursuing, but requires careful consideration of some additional design points such as escape analysis of return values (e.g. the lifetime of `tf.Variable` objects that are returned from a decorated function). -For now, we propose that `defun` continue with the restricted abilities proposed in this document and a "maintain Python semantics" decorator be investigated independently. +For now, we propose that `function` continue with the restricted abilities proposed in this document and a "maintain Python semantics" decorator be investigated independently. ## Open Questions/Ideas @@ -1113,16 +1113,13 @@ For now, we propose that `defun` continue with the restricted abilities proposed * Naming: - * `tf.defun` or `tf.function`? * `tf.compat.v1.wrap_function` or `tf.compat.v1.defun` or `tf.compat.v1.function` or `tf.compat.v1.wrap_graph_as_function`? * Signatures in Python 3? ([From ngc92](https://github.com/tensorflow/community/pull/20#issuecomment-423345326)) -* Can `tf.defun` and `tf.method` be combined into a single decorator (where `tf.defun` has the behavior of `tf.method` when applied to a class method)? Or is it okay to have two separate decorators? - * (How do you detect if the decorated function is a method or a function? Rely on the convention of first argument being called `self`?) * Supporting structured inputs: \ -As proposed, arguments to `defun` must be either `Tensor` objects, or objects that can be converted to a `Tensor` (`tf.convert_to_tensor`), or opaque Python objects. \ +As proposed, arguments to `function` must be either `Tensor` objects, or objects that can be converted to a `Tensor` (`tf.convert_to_tensor`), or opaque Python objects. \ \ Perhaps we can support nested structures of `Tensor`s (using `nest.flatten` and `nest.pack_sequence_as`), or even arbitrary Python objects? \ \ -If this is supported, then specifying an `input_signature` may become cumbersome, but perhaps we can have a `defun(infer_signature_from_first_call=True)` to make that easier. \ +If this is supported, then specifying an `input_signature` may become cumbersome, but perhaps we can have a `function(infer_signature_from_first_call=True)` to make that easier. \ \ From f9b64c0f505e15844dbc9d8c926e2058e4b9daa2 Mon Sep 17 00:00:00 2001 From: Asim Shankar Date: Tue, 9 Oct 2018 14:17:48 -0700 Subject: [PATCH 11/15] Collapse tf.method and tf.function. (Prototyped in https://github.com/tensorflow/tensorflow/commit/84ace0358526bb51c04a3bef4b3072b93b9d1bec) --- rfcs/20180918-functions-not-sessions-20.md | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/rfcs/20180918-functions-not-sessions-20.md b/rfcs/20180918-functions-not-sessions-20.md index b3d61b46c..89fb39b49 100644 --- a/rfcs/20180918-functions-not-sessions-20.md +++ b/rfcs/20180918-functions-not-sessions-20.md @@ -487,7 +487,7 @@ assert int(model2.v) == 5 \ -This works since `increment()` has `self` as a non-tensor argument, and a new trace will be created for each value of `self`. However, if variables are created in a method, we want to allow a new set of variables for every instantiation of `self`. You get this behavior by using `@tf.method`: +This works since `increment()` has `self` as a non-tensor argument, and a new trace will be created for each value of `self`. However, if variables are created in a method, we want to allow a new set of variables for every instantiation of `self`. ```python @@ -495,7 +495,7 @@ class AnyShapeModel(object): def __init__(self): self.v = None - @tf.method + @tf.function def increment(self, amount): if self.v is None: self.v = tf.Variable(tf.zeros_like(amount)) @@ -512,9 +512,7 @@ assert model2.v.numpy() == [4, 5] ``` -The semantics here are that each new instance is allowed to create variables in each `@tf.method` once. The simple recommendation would be "always use `@tf.method` on methods, use `@tf.function` for functions outside of a class". - -In the above example, if `increment` was decorated with `@tf.function` instead, then the `model2.increment()` call would raise an exception (as per `function`s stated behavior of disallowing state creation on anything but the first trace). However, if the method didn't create any state then `@tf.function` or `@tf.method` would both have the same effect. +The semantics here are that each new instance is allowed to create variables in each `@tf.function` once. In addition, as long as all variable creation/initialization happens while we are tracing, we should be able to support exporting the initialization graph when exporting a `SavedModel` or `MetaGraphDef`. @@ -775,7 +773,7 @@ class Model(tf.train.Checkpointable): def __init__(self): self.W = tf.Variable(...) - @tf.method + @tf.function def f(self, x): return tf.matmul(x, self.W) From 868b1106502b54825bcf477194d23d4f2ffed486 Mon Sep 17 00:00:00 2001 From: Asim Shankar Date: Mon, 15 Oct 2018 11:53:52 -0700 Subject: [PATCH 12/15] Fix typo --- rfcs/20180918-functions-not-sessions-20.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/rfcs/20180918-functions-not-sessions-20.md b/rfcs/20180918-functions-not-sessions-20.md index 89fb39b49..a54a366c3 100644 --- a/rfcs/20180918-functions-not-sessions-20.md +++ b/rfcs/20180918-functions-not-sessions-20.md @@ -380,8 +380,8 @@ def f(x, y): return tf.multiply(x, 2.) if y.multiply else tf.add(x, 2.) f(3., p) # Returns 6.0 -y.multiply = False -f(3., y) # Mutations to `y` may not trigger a retrace, so might still return 6.0 +p.multiply = False +f(3., p) # Mutations to `p` may not trigger a retrace, so might still return 6.0 ``` From 765263d70eac7919fb07098c1c7e05791376b70c Mon Sep 17 00:00:00 2001 From: Asim Shankar Date: Tue, 30 Oct 2018 20:05:30 -0700 Subject: [PATCH 13/15] Formatting tweak --- rfcs/20180918-functions-not-sessions-20.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/rfcs/20180918-functions-not-sessions-20.md b/rfcs/20180918-functions-not-sessions-20.md index a54a366c3..2b1374c95 100644 --- a/rfcs/20180918-functions-not-sessions-20.md +++ b/rfcs/20180918-functions-not-sessions-20.md @@ -76,7 +76,7 @@ z1 = compute_z1(2., 2.) ``` -Where `tf.function` is a decorator that "**de**fines a TensorFlow **fun**ction". A "TensorFlow function" defines a computation as a graph of TensorFlow operations, with named arguments and explicit return values. Users define the function they want TensorFlow to "accelerate" as a Python function and integrate it into their Python program like any other Python function call. +Where `tf.function` is a decorator that "defines a TensorFlow function". A "TensorFlow function" defines a computation as a graph of TensorFlow operations, with named arguments and explicit return values. Users define the function they want TensorFlow to "accelerate" as a Python function and integrate it into their Python program like any other Python function call. Having the Python function correspond to what the runtime will execute reduces conceptual complexity in translating between the two domains. It also affords an opportunity to provide more helpful stacktraces on errors. More advanced features available today (e.g., carving sub-graphs, feeding intermediate values) will still be possible (discussed later), though most users should not need to think in terms of graphs, feeds, and fetches. The constructed graph also provides a natural point for accelerators/acceleration libraries (NVIDIA TensorRT, Google Cloud TPUs etc.) to hook in for rewrites. @@ -952,7 +952,7 @@ This situation can be improved with the help of [autograph](https://github.com/ ```python -df = tf.function(f, autograph=True) +df = tf.function(autograph=True)(f) f(x, y) # Will be 1.0 ``` From f6bab47e858d7bd4015b92dff1f0a6b7f7c752ca Mon Sep 17 00:00:00 2001 From: Asim Shankar Date: Mon, 5 Nov 2018 23:51:14 -0800 Subject: [PATCH 14/15] Fix typo --- rfcs/20180918-functions-not-sessions-20.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/rfcs/20180918-functions-not-sessions-20.md b/rfcs/20180918-functions-not-sessions-20.md index 2b1374c95..c64df0d11 100644 --- a/rfcs/20180918-functions-not-sessions-20.md +++ b/rfcs/20180918-functions-not-sessions-20.md @@ -183,7 +183,7 @@ with tf.Session() as sess: ``` -The output here is not deterministic, since `val` may evaluate to either 1.0 or 2.0 depending on whether the runtime happened to execute `assign_op` before `y` or not. `tf.control_dependencies` is a mechanism provided to add annotations at graph construction time to influence graph execution. The TensorFlow user, a Python programmer, is thus forced to think about two execution models - TensorFlow graphs and the Python interpreter. To eliminate this cognitive load, `function` will automatically insert control dependencies to ensure that (1) operations that produce or consume a given `DT_RESOURCE` tensor and (2) operations that are marked stateful (`REGISTER_OP(...).SetIsStateful()`) follow graph construction order. Thus: +The output here is not deterministic, since `val` may evaluate to either 1.0 or 2.0 depending on whether the runtime happened to execute `assign_op` before `read` or not. `tf.control_dependencies` is a mechanism provided to add annotations at graph construction time to influence graph execution. The TensorFlow user, a Python programmer, is thus forced to think about two execution models - TensorFlow graphs and the Python interpreter. To eliminate this cognitive load, `function` will automatically insert control dependencies to ensure that (1) operations that produce or consume a given `DT_RESOURCE` tensor and (2) operations that are marked stateful (`REGISTER_OP(...).SetIsStateful()`) follow graph construction order. Thus: ```python From ba6c79ea8cd265367c4dc8b65a12ed2d2a37b92c Mon Sep 17 00:00:00 2001 From: Asim Shankar Date: Thu, 15 Nov 2018 18:53:27 -0800 Subject: [PATCH 15/15] Update note on tracing twice. --- rfcs/20180918-functions-not-sessions-20.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/rfcs/20180918-functions-not-sessions-20.md b/rfcs/20180918-functions-not-sessions-20.md index c64df0d11..ed11a3bcd 100644 --- a/rfcs/20180918-functions-not-sessions-20.md +++ b/rfcs/20180918-functions-not-sessions-20.md @@ -245,7 +245,7 @@ To support this `function` imposes some requirements on the decorated function: 1. State (like `tf.Variable` objects) are only created the first time the function `f` is called. \ How that is accomplished is left up to the implementation of `f`. \ -If any variables are created in the first execution of `f`, then `@tf.function` will trace `f` a second time in order to record the behavior that will be used from then on. No variables may be created during that second trace, or any other trace after that (due to different dtypes, shapes, or non-tensor arguments). +If any variables are created in the first execution of `f`, then `@tf.function` will trace `f` again the second time it is invoked in order to record the behavior that will be used from then on. No variables may be created during that second trace, or any other trace after that (due to different dtypes, shapes, or non-tensor arguments). 1. The caller must make sure that any variable referenced by the function still exists whenever the function is evaluated. \ `@tf.function` itself will keep only weak references to these created variables. Thus, if the referenced state does not exist when the decorated function is invoked, an exception will be raised.