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