1 # Copyright 2017 The TensorFlow Authors. All Rights Reserved. 2 # 3 # Licensed under the Apache License, Version 2.0 (the "License"); 4 # you may not use this file except in compliance with the License. 5 # You may obtain a copy of the License at 6 # 7 # http://www.apache.org/licenses/LICENSE-2.0 8 # 9 # Unless required by applicable law or agreed to in writing, software 10 # distributed under the License is distributed on an "AS IS" BASIS, 11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 # See the License for the specific language governing permissions and 13 # limitations under the License. 14 # ============================================================================== 15 r"""TensorFlow Eager Execution Example: Linear Regression. 16 17 This example shows how to use TensorFlow Eager Execution to fit a simple linear 18 regression model using some synthesized data. Specifically, it illustrates how 19 to define the forward path of the linear model and the loss function, as well 20 as how to obtain the gradients of the loss function with respect to the 21 variables and update the variables with the gradients. 22 """ 23 24 from __future__ import absolute_import 25 from __future__ import division 26 from __future__ import print_function 27 28 import argparse 29 import sys 30 31 import tensorflow as tf 32 33 import tensorflow.contrib.eager as tfe 34 35 36 class LinearModel(tfe.Network): 37 """A TensorFlow linear regression model. 38 39 Uses TensorFlow's eager execution. 40 41 For those familiar with TensorFlow graphs, notice the absence of 42 `tf.Session`. The `forward()` method here immediately executes and 43 returns output values. The `loss()` method immediately compares the 44 output of `forward()` with the target and returns the MSE loss value. 45 The `fit()` performs gradient-descent training on the model's weights 46 and bias. 47 """ 48 49 def __init__(self): 50 """Constructs a LinearModel object.""" 51 super(LinearModel, self).__init__() 52 self._hidden_layer = self.track_layer(tf.layers.Dense(1)) 53 54 def call(self, xs): 55 """Invoke the linear model. 56 57 Args: 58 xs: input features, as a tensor of size [batch_size, ndims]. 59 60 Returns: 61 ys: the predictions of the linear mode, as a tensor of size [batch_size] 62 """ 63 return self._hidden_layer(xs) 64 65 66 def mean_square_loss(model, xs, ys): 67 return tf.reduce_mean(tf.square(model(xs) - ys)) 68 69 70 def fit(model, dataset, optimizer, verbose=False, logdir=None): 71 """Fit the linear-regression model. 72 73 Args: 74 model: The LinearModel to fit. 75 dataset: The tf.data.Dataset to use for training data. 76 optimizer: The TensorFlow Optimizer object to be used. 77 verbose: If true, will print out loss values at every iteration. 78 logdir: The directory in which summaries will be written for TensorBoard 79 (optional). 80 """ 81 82 # The loss function to optimize. 83 mse = lambda xs, ys: mean_square_loss(model, xs, ys) 84 loss_and_grads = tfe.implicit_value_and_gradients(mse) 85 86 tf.train.get_or_create_global_step() 87 if logdir: 88 # Support for TensorBoard summaries. Once training has started, use: 89 # tensorboard --logdir=<logdir> 90 summary_writer = tf.contrib.summary.create_file_writer(logdir) 91 92 # Training loop. 93 for i, (xs, ys) in enumerate(tfe.Iterator(dataset)): 94 loss, grads = loss_and_grads(xs, ys) 95 if verbose: 96 print("Iteration %d: loss = %s" % (i, loss.numpy())) 97 98 optimizer.apply_gradients(grads, global_step=tf.train.get_global_step()) 99 100 if logdir: 101 with summary_writer.as_default(): 102 with tf.contrib.summary.always_record_summaries(): 103 tf.contrib.summary.scalar("loss", loss) 104 105 106 def synthetic_dataset(w, b, noise_level, batch_size, num_batches): 107 """tf.data.Dataset that yields synthetic data for linear regression.""" 108 return synthetic_dataset_helper(w, b, 109 tf.shape(w)[0], noise_level, batch_size, 110 num_batches) 111 112 113 def synthetic_dataset_helper(w, b, num_features, noise_level, batch_size, 114 num_batches): 115 # w is a matrix with shape [N, M] 116 # b is a vector with shape [M] 117 # So: 118 # - Generate x's as vectors with shape [batch_size N] 119 # - y = tf.matmul(x, W) + b + noise 120 def batch(_): 121 x = tf.random_normal([batch_size, num_features]) 122 y = tf.matmul(x, w) + b + noise_level * tf.random_normal([]) 123 return x, y 124 125 with tf.device("/device:CPU:0"): 126 return tf.data.Dataset.range(num_batches).map(batch) 127 128 129 def main(_): 130 tfe.enable_eager_execution() 131 # Ground-truth constants. 132 true_w = [[-2.0], [4.0], [1.0]] 133 true_b = [0.5] 134 noise_level = 0.01 135 136 # Training constants. 137 batch_size = 64 138 learning_rate = 0.1 139 140 print("True w: %s" % true_w) 141 print("True b: %s\n" % true_b) 142 143 model = LinearModel() 144 dataset = synthetic_dataset(true_w, true_b, noise_level, batch_size, 20) 145 146 device = "gpu:0" if tfe.num_gpus() else "cpu:0" 147 print("Using device: %s" % device) 148 with tf.device(device): 149 optimizer = tf.train.GradientDescentOptimizer(learning_rate) 150 fit(model, dataset, optimizer, verbose=True, logdir=FLAGS.logdir) 151 152 print("\nAfter training: w = %s" % model.variables[0].numpy()) 153 print("\nAfter training: b = %s" % model.variables[1].numpy()) 154 155 156 if __name__ == "__main__": 157 parser = argparse.ArgumentParser() 158 parser.add_argument( 159 "--logdir", 160 type=str, 161 default=None, 162 help="logdir in which TensorBoard summaries will be written (optional).") 163 FLAGS, unparsed = parser.parse_known_args() 164 165 tf.app.run(main=main, argv=[sys.argv[0]] + unparsed) 166
