def fit(self, X, y): with tf.Graph().as_default() as graph: tf.set_random_seed(self.tf_random_seed) self._global_step = tf.Variable(0, name="global_step", trainable=False) # Setting up input and output placeholders. input_shape = [None] + list(X.shape[1:]) self._inp = tf.placeholder(tf.float32, input_shape, name="input") self._out = tf.placeholder(tf.float32, [None] if self.n_classes < 2 else [None, self.n_classes], name="output") # Create model's graph. self._model_predictions, self._model_loss = self.model_fn(self._inp, self._out) # Create data feeder, to sample inputs from dataset. self._data_feeder = DataFeeder(X, y, self.n_classes, self.batch_size) # Create trainer and augment graph with gradients and optimizer. self._trainer = TensorFlowTrainer(self._model_loss, self._global_step, self.optimizer, self.learning_rate) self._session = tf.Session(self.tf_master) # Initialize and train model. self._trainer.initialize(self._session) self._trainer.train(self._session, self._data_feeder.get_feed_dict_fn(self._inp, self._out), self.steps)
class TensorFlowEstimator(BaseEstimator): """Base class for all TensorFlow estimators. Parameters: model_fn: Model function, that takes input X, y tensors and outputs prediction and loss tensors. n_classes: Number of classes in the target. tf_master: TensorFlow master. Empty string is default for local. batch_size: Mini batch size. steps: Number of steps to run over data. optimizer: Optimizer name (or class), for example "SGD", "Adam", "Adagrad". learning_rate: Learning rate for optimizer. tf_random_seed: Random seed for TensorFlow initializers. Setting this value, allows consistency between reruns. """ def __init__(self, model_fn, n_classes, tf_master="", batch_size=32, steps=50, optimizer="SGD", learning_rate=0.1, tf_random_seed=42): self.n_classes = n_classes self.tf_master = tf_master self.batch_size = batch_size self.steps = steps self.optimizer = optimizer self.learning_rate = learning_rate self.tf_random_seed = tf_random_seed self.model_fn = model_fn def fit(self, X, y): with tf.Graph().as_default() as graph: tf.set_random_seed(self.tf_random_seed) self._global_step = tf.Variable(0, name="global_step", trainable=False) # Setting up input and output placeholders. input_shape = [None] + list(X.shape[1:]) self._inp = tf.placeholder(tf.float32, input_shape, name="input") self._out = tf.placeholder(tf.float32, [None] if self.n_classes < 2 else [None, self.n_classes], name="output") # Create model's graph. self._model_predictions, self._model_loss = self.model_fn(self._inp, self._out) # Create data feeder, to sample inputs from dataset. self._data_feeder = DataFeeder(X, y, self.n_classes, self.batch_size) # Create trainer and augment graph with gradients and optimizer. self._trainer = TensorFlowTrainer(self._model_loss, self._global_step, self.optimizer, self.learning_rate) self._session = tf.Session(self.tf_master) # Initialize and train model. self._trainer.initialize(self._session) self._trainer.train(self._session, self._data_feeder.get_feed_dict_fn(self._inp, self._out), self.steps) def predict(self, X): pred = self._session.run(self._model_predictions, feed_dict={ self._inp.name: X }) if self.n_classes < 2: return pred return pred.argmax(axis=1)