From 252a5d13ed129107584a26766ac934d336e4b755 Mon Sep 17 00:00:00 2001 From: Richard Liaw Date: Mon, 21 Oct 2019 16:46:16 -0700 Subject: [PATCH] [sgd/tune][minor] more tf ports (#5953) --- .../sgd/examples/tensorflow_train_example.py | 6 +++--- python/ray/tune/logger.py | 13 +++++++++++++ 2 files changed, 16 insertions(+), 3 deletions(-) diff --git a/python/ray/experimental/sgd/examples/tensorflow_train_example.py b/python/ray/experimental/sgd/examples/tensorflow_train_example.py index 2b18a23bd..0ebd95e1c 100644 --- a/python/ray/experimental/sgd/examples/tensorflow_train_example.py +++ b/python/ray/experimental/sgd/examples/tensorflow_train_example.py @@ -3,7 +3,7 @@ from __future__ import division from __future__ import print_function import argparse -from tensorflow.data import Dataset +import tensorflow as tf from tensorflow.keras.models import Sequential from tensorflow.keras.layers import Dense import numpy as np @@ -43,8 +43,8 @@ def simple_dataset(config): x_train, y_train = linear_dataset(size=NUM_TRAIN_SAMPLES) x_test, y_test = linear_dataset(size=NUM_TEST_SAMPLES) - train_dataset = Dataset.from_tensor_slices((x_train, y_train)) - test_dataset = Dataset.from_tensor_slices((x_test, y_test)) + train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train)) + test_dataset = tf.data.Dataset.from_tensor_slices((x_test, y_test)) train_dataset = train_dataset.shuffle(NUM_TRAIN_SAMPLES).repeat().batch( batch_size) test_dataset = test_dataset.repeat().batch(batch_size) diff --git a/python/ray/tune/logger.py b/python/ray/tune/logger.py index 53a09207e..c672bdbcc 100644 --- a/python/ray/tune/logger.py +++ b/python/ray/tune/logger.py @@ -148,6 +148,11 @@ def tf2_compat_logger(config, logdir, trial=None): use_tf2_api = (distutils.version.LooseVersion(tf.__version__) >= distutils.version.LooseVersion("1.15.0")) if use_tf2_api: + # This is temporarily for RLlib because it disables v2 behavior... + from tensorflow.python import tf2 + if not tf2.enabled(): + tf = tf.compat.v1 + return TFLogger(config, logdir, trial) tf = tf.compat.v2 # setting this for TF2.0 return TF2Logger(config, logdir, trial) else: @@ -166,6 +171,10 @@ class TF2Logger(Logger): """ def _init(self): + global tf + if tf is None: + import tensorflow as tf + tf = tf.compat.v2 # setting this for TF2.0 self._file_writer = None self._hp_logged = False @@ -237,6 +246,10 @@ class TFLogger(Logger): """ def _init(self): + global tf + if tf is None: + import tensorflow as tf + tf = tf.compat.v1 # setting this for regular TF logger logger.debug("Initializing TFLogger instead of TF2Logger.") self._file_writer = tf.summary.FileWriter(self.logdir)