Home | History | Annotate | Download | only in learn
      1 #  Copyright 2016 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 """Example of using Estimator with multiple GPUs to distribute one model.
     15 
     16 This example only runs if you have multiple GPUs to assign to.
     17 """
     18 
     19 from __future__ import absolute_import
     20 from __future__ import division
     21 from __future__ import print_function
     22 
     23 import numpy as np
     24 from sklearn import datasets
     25 from sklearn import metrics
     26 from sklearn import model_selection
     27 import tensorflow as tf
     28 
     29 
     30 X_FEATURE = 'x'  # Name of the input feature.
     31 
     32 
     33 def my_model(features, labels, mode):
     34   """DNN with three hidden layers, and dropout of 0.1 probability.
     35 
     36   Note: If you want to run this example with multiple GPUs, Cuda Toolkit 7.0 and
     37   CUDNN 6.5 V2 from NVIDIA need to be installed beforehand.
     38 
     39   Args:
     40     features: Dict of input `Tensor`.
     41     labels: Label `Tensor`.
     42     mode: One of `ModeKeys`.
     43 
     44   Returns:
     45     `EstimatorSpec`.
     46   """
     47   # Create three fully connected layers respectively of size 10, 20, and 10 with
     48   # each layer having a dropout probability of 0.1.
     49   net = features[X_FEATURE]
     50   with tf.device('/device:GPU:1'):
     51     for units in [10, 20, 10]:
     52       net = tf.layers.dense(net, units=units, activation=tf.nn.relu)
     53       net = tf.layers.dropout(net, rate=0.1)
     54 
     55   with tf.device('/device:GPU:2'):
     56     # Compute logits (1 per class).
     57     logits = tf.layers.dense(net, 3, activation=None)
     58 
     59     # Compute predictions.
     60     predicted_classes = tf.argmax(logits, 1)
     61     if mode == tf.estimator.ModeKeys.PREDICT:
     62       predictions = {
     63           'class': predicted_classes,
     64           'prob': tf.nn.softmax(logits)
     65       }
     66       return tf.estimator.EstimatorSpec(mode, predictions=predictions)
     67 
     68     # Compute loss.
     69     loss = tf.losses.sparse_softmax_cross_entropy(labels=labels, logits=logits)
     70 
     71     # Create training op.
     72     if mode == tf.estimator.ModeKeys.TRAIN:
     73       optimizer = tf.train.AdagradOptimizer(learning_rate=0.1)
     74       train_op = optimizer.minimize(
     75           loss, global_step=tf.train.get_global_step())
     76       return tf.estimator.EstimatorSpec(mode, loss=loss, train_op=train_op)
     77 
     78     # Compute evaluation metrics.
     79     eval_metric_ops = {
     80         'accuracy': tf.metrics.accuracy(
     81             labels=labels, predictions=predicted_classes)
     82     }
     83     return tf.estimator.EstimatorSpec(
     84         mode, loss=loss, eval_metric_ops=eval_metric_ops)
     85 
     86 
     87 def main(unused_argv):
     88   iris = datasets.load_iris()
     89   x_train, x_test, y_train, y_test = model_selection.train_test_split(
     90       iris.data, iris.target, test_size=0.2, random_state=42)
     91 
     92   classifier = tf.estimator.Estimator(model_fn=my_model)
     93 
     94   # Train.
     95   train_input_fn = tf.estimator.inputs.numpy_input_fn(
     96       x={X_FEATURE: x_train}, y=y_train, num_epochs=None, shuffle=True)
     97   classifier.train(input_fn=train_input_fn, steps=100)
     98 
     99   # Predict.
    100   test_input_fn = tf.estimator.inputs.numpy_input_fn(
    101       x={X_FEATURE: x_test}, y=y_test, num_epochs=1, shuffle=False)
    102   predictions = classifier.predict(input_fn=test_input_fn)
    103   y_predicted = np.array(list(p['class'] for p in predictions))
    104   y_predicted = y_predicted.reshape(np.array(y_test).shape)
    105 
    106   # Score with sklearn.
    107   score = metrics.accuracy_score(y_test, y_predicted)
    108   print('Accuracy (sklearn): {0:f}'.format(score))
    109 
    110   # Score with tensorflow.
    111   scores = classifier.evaluate(input_fn=test_input_fn)
    112   print('Accuracy (tensorflow): {0:f}'.format(scores['accuracy']))
    113 
    114 
    115 if __name__ == '__main__':
    116   tf.app.run()
    117