I think it was supposed to be used with with tf.device("/gpu:0")
, but where do I put it? I don't think it's:
with tf.device("/gpu:0"):
tf.app.run()
So should I put it in the main()
function of the tf.app
, or the model function I use for the estimator?
EDIT: If this helps, this is my main()
function:
def main(unused_argv):
"""Code to load training folds data pickle or generate one if not present"""
# Create the Estimator
mnist_classifier = tf.estimator.Estimator(
model_fn=cnn_model_fn2, model_dir="F:/python_machine_learning_codes/tmp/custom_age_adience_1")
# Set up logging for predictions
# Log the values in the "Softmax" tensor with label "probabilities"
tensors_to_log = {"probabilities": "softmax_tensor"}
logging_hook = tf.train.LoggingTensorHook(
tensors=tensors_to_log, every_n_iter=100)
# Train the model
train_input_fn = tf.estimator.inputs.numpy_input_fn(
x={"x": train_data},
y=train_labels,
batch_size=64,
num_epochs=None,
shuffle=True)
mnist_classifier.train(
input_fn=train_input_fn,
steps=500,
hooks=[logging_hook])
# Evaluate the model and print results
"""Code to load eval fold data pickle or generate one if not present"""
eval_logs = {"probabilities": "softmax_tensor"}
eval_hook = tf.train.LoggingTensorHook(
tensors=eval_logs, every_n_iter=100)
eval_input_fn = tf.estimator.inputs.numpy_input_fn(
x={"x": eval_data},
y=eval_labels,
num_epochs=1,
shuffle=False)
eval_results = mnist_classifier.evaluate(input_fn=eval_input_fn, hooks=[eval_hook])
As you can see, I have no explicit declaration of a session anywhere here, so where exactly do I put the with tf.device("/gpu:0")
?
You can put it at the beginning of your model function, i.e., when you define your model, you should write:
def cnn_model_fn2(...):
with tf.device('/gpu:0'):
...
However, I would expect tensorflow to automatically use the gpu for your model. You may want to check whether it is properly detected:
from tensorflow.python.client import device_lib
device_lib.list_local_devices()