Article Directory

Effective TensorFlow 2

Effective TensorFlow 2

Several changes have been made in TensorFlow 2.0 to make TensorFlow users more efficient. TensorFlow 2.0 removes the redundant API , makes the API more consistent ( unified RNN , unified optimizer ), and better integrates with the Python runtime through Eager execution .

Many RFCs explain the changes in TensorFlow 2.0. This guide presents the appearance of TensorFlow 2.0 development. Suppose you have some knowledge of TensorFlow 1.x.

Brief summary of major changes

API cleanup

In TF 2.0, many APIs have disappeared or moved . Some of the major changes include deletion tf.app, tf.flagsand tf.loggingto support the now open source absl-py , relocation resides in the tf.contribproject, and tf.*moves less used functions into tf.mathsub-packages by cleaning up the main namespace . Certain APIs have been replaced by their 2.0 equivalents tf.summary, tf.keras.metricsand tf.keras.optimizers. The easiest way to automatically apply these renamings is to use the v2 upgrade script .

Eager to execute

TensorFlow 1.X requires users to tf.*manually stitch together abstract syntax trees (graphs) by calling API . Then, it asks the user session.run()to manually compile the abstract syntax tree by passing a set of output tensors and input tensors to the call. TensorFlow 2.0 executes eagerly (as Python usually does), and in 2.0, graphs and sessions should feel like implementation details.

A significant by-product of eager execution is that it is no longer needed tf.control_dependencies(), because all lines of code tf.functionare executed sequentially (in tf.function, code with side effects is executed in the order in which it was written).

No more global variables

TensorFlow 1.X relies heavily on the implicit global namespace. When you call it tf.Variable(), it will be placed in the default graph, even if you lose the Python variable pointing to it, it will remain in the default graph. Then, you can restore that tf.Variable, but tf.Variableyou know the name of the file to create. If you cannot control the creation of variables, it is difficult to do this. As a result, various mechanisms surge again tried to help users find their variable and framework to find user-created variables: variable scope, global collection, such as tf.get_global_step(), tf.global_variables_initializer()helper class method to optimize the gradient program implicitly calculate all the variables of training ,So on and so forth. TensorFlow 2.0 eliminates all these mechanisms ( Variables 2.0 RFC ), and instead uses the default mechanism: tracking variables! If you forget tf.Variable, garbage will be collected.

The requirement to track variables brings some extra work to the user, but for Keras objects (see below), the burden is minimized.

Function instead of session

session.run()A call is almost like a function call: specify the input and the function to be called, and then get a set of outputs. In TensorFlow 2.0, you can tf.function()decorate a Python function to mark it as JIT compilation so that TensorFlow will run it as a single graph ( Function 2.0 RFC ). This mechanism enables TensorFlow 2.0 to obtain all the advantages of the graphics mode:

Performance: functions can be optimized (node pruning, kernel fusion, etc.)
Portability: This function can be exported/re-imported ( SavedModel 2.0 RFC ), allowing users to reuse and share modular TensorFlow functions.

# TensorFlow 1.X
outputs = session.run(f(placeholder), feed_dict={
    
    placeholder: input})
# TensorFlow 2.0
outputs = f(input)

With the free distribution of Python and TensorFlow code, users can take advantage of Python's expressive power. But portable TensorFlow executes in a context without a Python interpreter, such as mobile, C++, and JavaScript. To help users avoid @tf.functionhaving to rewrite code when adding , AutoGraph converts a subset of Python constructs into their TensorFlow equivalents:

for/ while-> tf.while_loop( breakAnd continuesupport)
if -> tf.cond
for _ in dataset dataset.reduce > dataset.reduce

AutoGraph supports arbitrary nesting of control flow, which makes it possible to efficiently and concisely implement many complex ML programs, such as sequence models, reinforcement learning, custom training loops, etc.

Idiomatic TensorFlow 2.0 recommendations

Refactor the code into smaller functions

A common usage pattern in TensorFlow 1.X is the "kitchen sink" strategy, which first accounts for the union of all possible calculations, and then session.run()evaluates the selected tensor. In TensorFlow 2.0, users should refactor their code into smaller functions and call these functions as needed. Generally, it is not necessary to tf.functiondecorate every smaller function; only tf.functiondecorate advanced calculations-for example, a step of training or the forward pass of the model.

Use Keras layers and models to manage variables

Keras models and layers provide convenience variablesand trainable_variablesattributes, which collect all dependent variables in a recursive manner. This makes it easy to manage variables locally where they are used.

Compared:

def dense(x, W, b):
  return tf.nn.sigmoid(tf.matmul(x, W) + b)

@tf.function
def multilayer_perceptron(x, w0, b0, w1, b1, w2, b2 ...):
  x = dense(x, w0, b0)
  x = dense(x, w1, b1)
  x = dense(x, w2, b2)
  ...

# You still have to manage w_i and b_i, and their shapes are defined far away from the code.

Use Keras version:

# Each layer can be called, with a signature equivalent to linear(x)
layers = [tf.keras.layers.Dense(hidden_size, activation=tf.nn.sigmoid) for _ in range(n)]
perceptron = tf.keras.Sequential(layers)

# layers[3].trainable_variables => returns [w3, b3]
# perceptron.trainable_variables => returns [w0, b0, ...]

Keras layers/models are inherited from tf.train.Checkpointableand @tf.functionintegrated with, which makes it possible to checkpoint or export SavedModels directly from Keras objects. You don't necessarily have to use .fit()the .fit()API to take advantage of these integrations.

This is a transfer learning example that demonstrates how Keras makes it easy to collect a subset of related variables. Suppose you are using a shared backbone to train a multi-head model:

trunk = tf.keras.Sequential([...])
head1 = tf.keras.Sequential([...])
head2 = tf.keras.Sequential([...])

path1 = tf.keras.Sequential([trunk, head1])
path2 = tf.keras.Sequential([trunk, head2])

# Train on primary dataset
for x, y in main_dataset:
  with tf.GradientTape() as tape:
    # training=True is only needed if there are layers with different
    # behavior during training versus inference (e.g. Dropout).
    prediction = path1(x, training=True)
    loss = loss_fn_head1(prediction, y)
  # Simultaneously optimize trunk and head1 weights.
  gradients = tape.gradient(loss, path1.trainable_variables)
  optimizer.apply_gradients(zip(gradients, path1.trainable_variables))

# Fine-tune second head, reusing the trunk
for x, y in small_dataset:
  with tf.GradientTape() as tape:
    # training=True is only needed if there are layers with different
    # behavior during training versus inference (e.g. Dropout).
    prediction = path2(x, training=True)
    loss = loss_fn_head2(prediction, y)
  # Only optimize head2 weights, not trunk weights
  gradients = tape.gradient(loss, head2.trainable_variables)
  optimizer.apply_gradients(zip(gradients, head2.trainable_variables))

# You can publish just the trunk computation for other people to reuse.
tf.saved_model.save(trunk, output_path)

Combine tf.data.Datasets and @tf.function

When iterating through training data that fits in memory, please feel free to use regular Python iterations. Otherwise, it tf.data.Datasetis the best way to stream training data from disk. Data sets are iterable (not iterators) , just like other Python iterables in Eager mode. By wrapping the code in tf.function(), you can take full advantage of the data set asynchronous prefetch/streaming feature, which replaces Python iteration with the equivalent graph operation using AutoGraph.

@tf.function
def train(model, dataset, optimizer):
  for x, y in dataset:
    with tf.GradientTape() as tape:
      # training=True is only needed if there are layers with different
      # behavior during training versus inference (e.g. Dropout).
      prediction = model(x, training=True)
      loss = loss_fn(prediction, y)
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))

If you use the .fit()API, you don't have to worry about data set iteration.

model.compile(optimizer=optimizer, loss=loss_fn)
model.fit(dataset)

Utilize AutoGraph with Python control flow

AutoGraph provides a way to convert data-dependent control flow into an equivalent tf.condand tf.while_loopgraph mode tf.while_loop.

A common place where data-dependent control flow appears is the sequence model. tf.keras.layers.RNNPackage an RNN unit so that you can expand the repetition statically or dynamically. For demonstration purposes, you can re-implement the dynamic expansion as follows:

class DynamicRNN(tf.keras.Model):

  def __init__(self, rnn_cell):
    super(DynamicRNN, self).__init__(self)
    self.cell = rnn_cell

  def call(self, input_data):
    # [batch, time, features] -> [time, batch, features]
    input_data = tf.transpose(input_data, [1, 0, 2])
    outputs = tf.TensorArray(tf.float32, input_data.shape[0])
    state = self.cell.zero_state(input_data.shape[1], dtype=tf.float32)
    for i in tf.range(input_data.shape[0]):
      output, state = self.cell(input_data[i], state)
      outputs = outputs.write(i, output)
    return tf.transpose(outputs.stack(), [1, 0, 2]), state

For a more detailed overview of AutoGraph features, see the guide .

tf.metrics summarizes data, and tf.summary records them

To record a summary, use tf.summary.(scalar|histogram|...)and then use the context manager to redirect it to the author. (If the context manager is omitted, nothing happens.) Unlike TF 1.x, the summary is sent directly to the writer; there is no separate "merge" operation, and no separate add_summary()call, which means it must be on the calling site Provide stepvalue.

summary_writer = tf.summary.create_file_writer('/tmp/summaries')
with summary_writer.as_default():
  tf.summary.scalar('loss', 0.1, step=42)

To aggregate data before recording it as a summary, use tf.metrics. Indicators are stateful: when you call .result()them, they accumulate values and return the accumulated result. Use to .reset_states()clear the accumulated value.

def train(model, optimizer, dataset, log_freq=10):
  avg_loss = tf.keras.metrics.Mean(name='loss', dtype=tf.float32)
  for images, labels in dataset:
    loss = train_step(model, optimizer, images, labels)
    avg_loss.update_state(loss)
    if tf.equal(optimizer.iterations % log_freq, 0):
      tf.summary.scalar('loss', avg_loss.result(), step=optimizer.iterations)
      avg_loss.reset_states()

def test(model, test_x, test_y, step_num):
  # training=False is only needed if there are layers with different
  # behavior during training versus inference (e.g. Dropout).
  loss = loss_fn(model(test_x, training=False), test_y)
  tf.summary.scalar('loss', loss, step=step_num)

train_summary_writer = tf.summary.create_file_writer('/tmp/summaries/train')
test_summary_writer = tf.summary.create_file_writer('/tmp/summaries/test')

with train_summary_writer.as_default():
  train(model, optimizer, dataset)

with test_summary_writer.as_default():
  test(model, test_x, test_y, optimizer.iterations)

Visualize the generated summary by pointing TensorBoard to the summary log directory:

tensorboard --logdir /tmp/summaries

Use tf.config.experimental_run_functions_eagerly() when debugging

In TensorFlow 2.0, Eager execution allows you to run code step by step to check shapes, data types, and values. Certain APIs (for example tf.function, tf.kerasetc.) are designed to be executed using Graph to improve performance and portability. When debugging, tf.config.experimental_run_functions_eagerly(True)use Eager to execute within this code.

E.g:

@tf.function 
def f(x): 
	if x > 0:  
		import pdb  
		pdb.set_trace()  
		x = x + 1 
		return x
tf.config.experimental_run_functions_eagerly(True)
f(tf.constant(1))

>>> f()
-> x = x + 1
(Pdb) l 
6   @tf.function 
7   def f(x): 
8    if x > 0: 
9     import pdb 
10     pdb.set_trace() 
11  ->   x = x + 1 
12    return x 
13 
14   tf.config.experimental_run_functions_eagerly(True) 
15   f(tf.constant(1))[EOF]

This can also be used in Keras models and other APIs that support Eager execution:

class CustomModel(tf.keras.models.Model): 
	@tf.function 
	def call(self, input_data):  
		if tf.reduce_mean(input_data) > 0:   
			return input_data  
		else:   
			import pdb   
			pdb.set_trace()   
			return input_data // 2
tf.config.experimental_run_functions_eagerly(True)
model = CustomModel()
model(tf.constant([-2, -4]))

>>> call()-> return input_data // 2
(Pdb) l 
10     if tf.reduce_mean(input_data) > 0: 
11      return input_data 
12     else: 
13      import pdb 
14      pdb.set_trace() 
15  ->    return input_data // 2 
16 
17 
18   tf.config.experimental_run_functions_eagerly(True) 
19   model = CustomModel() 
20   model(tf.constant([-2, -4]))

Source: Effective TensorFlow 2

Introduction to the new features of Tensorflow 2.0-the most detailed version