def model_fn(features, labels, mode):
from nets import lenet
slim = tf.contrib.slim
with slim.arg_scope(lenet.lenet_arg_scope()):
logits, end_points = lenet.lenet(features,
num_classes=10,
is_training=True)
if mode == tf.estimator.ModeKeys.EVAL or mode == tf.estimator.ModeKeys.TRAIN:
loss = tf.reduce_mean(
tf.losses.sparse_softmax_cross_entropy(logits=logits,
labels=labels))
optimizer = ZooOptimizer(tf.train.AdamOptimizer())
train_op = optimizer.minimize(loss)
return tf.estimator.EstimatorSpec(mode,
loss=loss,
train_op=train_op)
else:
return tf.estimator.EstimatorSpec(mode, predictions=logits)
class TFOptimizerWrapper(Estimator):
def __init__(self, *, inputs, outputs, labels, loss, optimizer, metrics,
updates, sess, model_dir):
self.inputs = inputs
self.outputs = outputs
self.labels = labels
self.loss = loss
if optimizer is not None:
assert isinstance(optimizer, tf.train.Optimizer), \
"optimizer is of type {}, ".format(type(self.optimizer)) + \
"it should be an instance of tf.train.Optimizer"
self.optimizer = ZooOptimizer(optimizer)
self.train_op = self.optimizer.minimize(self.loss)
else:
self.optimizer = None
self.train_op = None
self.metrics = metrics
self.updates = updates
if sess is None:
self.sess = tf.Session()
self.sess.run(tf.global_variables_initializer())
else:
self.sess = sess
self.model_dir = model_dir
def fit(self,
data,
steps,
batch_size=32,
validation_data=None,
feed_dict=None,
session_config=None):
assert self.labels is not None, \
"labels is None; it should not be None in training"
assert self.loss is not None, \
"loss is None; it should not be None in training"
assert self.optimizer is not None, \
"optimizer is None; it not None in training"
dataset = _to_dataset(data, batch_size=batch_size, batch_per_thread=-1)
if feed_dict is not None:
tensor_with_value = {
key: (value, value)
for key, value in feed_dict.items()
}
else:
tensor_with_value = None
optimizer = TFOptimizer.from_train_op(
train_op=self.train_op,
loss=self.loss,
inputs=self.inputs,
labels=self.labels,
dataset=dataset,
metrics=self.metrics,
updates=self.updates,
sess=self.sess,
tensor_with_value=tensor_with_value,
session_config=session_config,
model_dir=self.model_dir)
optimizer.optimize(end_trigger=MaxIteration(steps))
return self
def predict(self, data, batch_size=32):
assert self.outputs is not None, \
"output is None, it should not be None in prediction"
dataset = _to_dataset(data, batch_size=-1, batch_per_thread=batch_size)
flat_inputs = nest.flatten(self.inputs)
flat_outputs = nest.flatten(self.outputs)
tfnet = TFNet.from_session(sess=self.sess,
inputs=flat_inputs,
outputs=flat_outputs)
return tfnet.predict(dataset)
def evaluate(self, data, batch_size=32):
assert self.metrics is not None, \
"metrics is None, it should not be None in evaluate"
dataset = _to_dataset(data, batch_size=-1, batch_per_thread=batch_size)
flat_inputs = nest.flatten(self.inputs)
flat_labels = nest.flatten(self.labels)
return evaluate_metrics(flat_inputs + flat_labels,
sess=self.sess,
dataset=dataset,
metrics=self.metrics)
class TensorFlowEstimator(Estimator):
def __init__(self, *, inputs, outputs, labels, loss, optimizer, clip_norm,
clip_value, metrics, updates, sess, model_dir):
self.inputs = inputs
self.outputs = outputs
self.labels = labels
self.loss = loss
self.use_bigdl_optim = False
self.clip_norm = clip_norm
self.clip_value = clip_value
if optimizer is not None:
from zoo.orca.learn.optimizers import Optimizer
if isinstance(optimizer, Optimizer):
self.train_op = None
self.optimizer = optimizer.get_optimizer()
self.use_bigdl_optim = True
else:
assert isinstance(optimizer, tf.train.Optimizer), \
"optimizer is of type {}, ".format(type(optimizer)) + \
"it should be an instance of tf.train.Optimizer"
self.optimizer = ZooOptimizer(optimizer)
if clip_norm or clip_value:
gvs = self.optimizer.compute_gradients(self.loss)
if clip_norm:
gvs = [(tf.clip_by_norm(g_v[0], clip_norm), g_v[1])
for g_v in gvs]
if clip_value:
if isinstance(clip_value, tuple):
assert len(clip_value) == 2 and clip_value[0] < clip_value[1], \
"clip value should be (clip_min, clip_max)"
gvs = [(tf.clip_by_value(g_v[0], clip_value[0],
clip_value[1]), g_v[1])
for g_v in gvs]
if isinstance(clip_value, (int, float)):
assert clip_value > 0, "clip value should be larger than 0"
gvs = [(tf.clip_by_value(g_v[0], -clip_value,
clip_value), g_v[1])
for g_v in gvs]
else:
raise Exception(
"clip_value should be a tuple or one number")
self.train_op = self.optimizer.apply_gradients(gvs)
else:
self.train_op = self.optimizer.minimize(self.loss)
else:
self.optimizer = None
self.train_op = None
self.metrics = metrics
self.updates = updates
if sess is None:
self.sess = tf.Session()
self.sess.run(tf.global_variables_initializer())
else:
self.sess = sess
self.model_dir = model_dir
self.load_checkpoint = False
self.tf_optimizer = None
self.log_dir = None
self.app_name = None
def fit(self,
data,
epochs=1,
batch_size=32,
feature_cols=None,
label_cols=None,
validation_data=None,
session_config=None,
checkpoint_trigger=None,
auto_shard_files=False,
feed_dict=None):
"""
Train this graph model with train data.
:param data: train data. It can be XShards, Spark DataFrame, tf.data.Dataset.
If data is XShards, each partition can be Pandas Dataframe or a dictionary of
{'x': feature, 'y': label}, where feature(label) is a numpy array or a tuple of
numpy arrays.
If data is tf.data.Dataset, each element is a tuple of input tensors.
:param epochs: number of epochs to train.
:param batch_size: total batch size for each iteration.
:param feature_cols: feature column names if train data is Spark DataFrame or XShards
of Pandas Dataframe.
:param label_cols: label column names if train data is Spark DataFrame or XShards of
Pandas Dataframe.
:param validation_data: validation data. Validation data type should be the same
as train data.
:param auto_shard_files: whether to automatically detect if the dataset is file-based and
and apply sharding on files, otherwise sharding on records. Default is False.
:param session_config: tensorflow session configuration for training.
Should be object of tf.ConfigProto
:param feed_dict: a dictionary. The key is TensorFlow tensor, usually a
placeholder, the value of the dictionary is a tuple of two elements. The first one of
the tuple is the value to feed to the tensor in training phase and the second one
is the value to feed to the tensor in validation phase.
:param checkpoint_trigger: when to trigger checkpoint during training.
Should be a zoo.orca.learn.trigger, like EveryEpoch(), SeveralIteration(
num_iterations),etc.
"""
assert self.labels is not None, \
"labels is None; it should not be None in training"
assert self.loss is not None, \
"loss is None; it should not be None in training"
assert self.optimizer is not None, \
"optimizer is None; it should not be None in training"
if isinstance(data, DataFrame):
assert feature_cols is not None, \
"feature columns is None; it should not be None in training"
assert label_cols is not None, \
"label columns is None; it should not be None in training"
if isinstance(data, SparkXShards):
if data._get_class_name() == 'pandas.core.frame.DataFrame':
assert feature_cols is not None, \
"feature columns is None; it should not be None in training"
assert label_cols is not None, \
"label columns is None; it should not be None in training"
data, validation_data = process_xshards_of_pandas_dataframe(
data, feature_cols, label_cols, validation_data, "fit")
if checkpoint_trigger is not None:
checkpoint_trigger = Trigger.convert_trigger(checkpoint_trigger)
memory_type = OrcaContext.train_data_store
dataset = to_dataset(data,
batch_size=batch_size,
batch_per_thread=-1,
validation_data=validation_data,
feature_cols=feature_cols,
label_cols=label_cols,
hard_code_batch_size=False,
sequential_order=False,
shuffle=True,
auto_shard_files=auto_shard_files,
memory_type=memory_type)
if feed_dict is not None:
tensor_with_value = {
key: (value[0], value[1])
for key, value in feed_dict.items()
}
else:
tensor_with_value = None
if self.use_bigdl_optim:
self.tf_optimizer = TFOptimizer.from_loss(
self.loss,
self.optimizer,
session=self.sess,
inputs=(self.inputs, self.labels),
dataset=dataset,
clip_norm=self.clip_norm,
clip_value=self.clip_value,
metrics=self.metrics,
tensor_with_value=tensor_with_value,
session_config=session_config,
model_dir=self.model_dir,
updates=self.updates)
else:
self.tf_optimizer = TFOptimizer.from_train_op(
train_op=self.train_op,
loss=self.loss,
inputs=self.inputs,
labels=self.labels,
dataset=dataset,
metrics=self.metrics,
updates=self.updates,
sess=self.sess,
tensor_with_value=tensor_with_value,
session_config=session_config,
model_dir=self.model_dir)
if self.load_checkpoint:
self.tf_optimizer.load_checkpoint(self.checkpoint_path,
self.checkpoint_version)
if self.log_dir and self.app_name:
self.tf_optimizer.estimator.set_tensorboard(
self.log_dir, self.app_name)
self.tf_optimizer.optimize(end_trigger=MaxEpoch(epochs),
checkpoint_trigger=checkpoint_trigger)
return self
def predict(
self,
data,
batch_size=4,
feature_cols=None,
auto_shard_files=False,
):
"""
Predict input data
:param data: data to be predicted. It can be XShards, Spark DataFrame.
If data is XShards, each partition can be Pandas Dataframe or a dictionary of
{'x': feature}, where feature is a numpy array or a tuple of numpy arrays.
:param batch_size: batch size per thread
:param feature_cols: list of feature column names if input data is Spark DataFrame
or XShards of Pandas DataFrame.
:param auto_shard_files: whether to automatically detect if the dataset is file-based and
and apply sharding on files, otherwise sharding on records. Default is False.
:return: predicted result.
If input data is XShards or tf.data.Dataset, the predict result is a XShards, each
partition of the XShards is a dictionary of {'prediction': result}, where the
result is a numpy array or a list of numpy arrays.
If input data is Spark DataFrame, the predict result is a DataFrame which includes
original columns plus 'prediction' column. The 'prediction' column can be
FloatType, VectorUDT or Array of VectorUDT depending on model outputs shape.
"""
assert self.outputs is not None, \
"output is None, it should not be None in prediction"
if isinstance(data, DataFrame):
assert feature_cols is not None, \
"feature columns is None; it should not be None in prediction"
if isinstance(data, SparkXShards):
if data._get_class_name() == 'pandas.core.frame.DataFrame':
assert feature_cols is not None, \
"feature columns is None; it should not be None in prediction"
data = process_xshards_of_pandas_dataframe(data, feature_cols)
assert not is_tf_data_dataset(data), "tf.data.Dataset currently cannot be used for" \
"estimator prediction"
dataset = to_dataset(
data,
batch_size=-1,
batch_per_thread=batch_size,
validation_data=None,
feature_cols=feature_cols,
label_cols=None,
hard_code_batch_size=False,
sequential_order=True,
shuffle=False,
auto_shard_files=auto_shard_files,
)
flat_inputs = nest.flatten(self.inputs)
flat_outputs = nest.flatten(self.outputs)
tfnet = TFNet.from_session(sess=self.sess,
inputs=flat_inputs,
outputs=flat_outputs)
predicted_rdd = tfnet.predict(dataset)
if isinstance(data, DataFrame):
return convert_predict_rdd_to_dataframe(data, predicted_rdd)
elif isinstance(data, SparkXShards):
return convert_predict_rdd_to_xshard(data, predicted_rdd)
else:
return predicted_rdd
def evaluate(
self,
data,
batch_size=32,
feature_cols=None,
label_cols=None,
auto_shard_files=False,
):
"""
Evaluate model.
:param data: evaluation data. It can be XShards, Spark DataFrame, tf.data.Dataset.
If data is XShards, each partition can be Pandas Dataframe or a dictionary of
{'x': feature, 'y': label}, where feature(label) is a numpy array or a tuple of
numpy arrays.
If data is tf.data.Dataset, each element is a tuple of input tensors.
:param batch_size: batch size per thread.
:param feature_cols: feature_cols: feature column names if train data is Spark DataFrame
or XShards of Pandas DataFrame.
:param label_cols: label column names if train data is Spark DataFrame or XShards
of Pandas DataFrame.
:param auto_shard_files: whether to automatically detect if the dataset is file-based and
and apply sharding on files, otherwise sharding on records. Default is False.
:return: evaluation result as a dictionary of {'metric name': metric value}
"""
assert self.metrics is not None, \
"metrics is None, it should not be None in evaluate"
if isinstance(data, DataFrame):
assert feature_cols is not None, \
"feature columns is None; it should not be None in evaluation"
assert label_cols is not None, \
"label columns is None; it should not be None in evaluation"
if isinstance(data, SparkXShards):
if data._get_class_name() == 'pandas.core.frame.DataFrame':
assert feature_cols is not None, \
"feature columns is None; it should not be None in evaluation"
assert label_cols is not None, \
"label columns is None; it should not be None in evaluation"
data = process_xshards_of_pandas_dataframe(
data, feature_cols, label_cols)
dataset = to_dataset(
data,
batch_size=-1,
batch_per_thread=batch_size,
validation_data=None,
feature_cols=feature_cols,
label_cols=label_cols,
hard_code_batch_size=False,
sequential_order=True,
shuffle=False,
auto_shard_files=auto_shard_files,
)
flat_inputs = nest.flatten(self.inputs)
flat_labels = nest.flatten(self.labels)
return evaluate_metrics(flat_inputs + flat_labels,
sess=self.sess,
dataset=dataset,
metrics=self.metrics)
def save_tf_checkpoint(self, path):
"""
Save tensorflow checkpoint in this estimator.
:param path: tensorflow checkpoint path.
"""
save_tf_checkpoint(self.sess, path)
def get_model(self):
"""
Get_model is not supported in tensorflow graph estimator
"""
raise NotImplementedError
def save(self, model_path):
"""
Save model to model_path
:param model_path: path to save the trained model.
:return:
"""
self.save_tf_checkpoint(model_path)
def clear_gradient_clipping(self):
"""
Clear gradient clipping is not supported in TensorFlowEstimator.
"""
raise NotImplementedError
def set_constant_gradient_clipping(self, min, max):
"""
Set constant gradient clipping is not supported in TensorFlowEstimator. Please pass the
clip_value to Estimator.from_graph.
"""
raise NotImplementedError
def set_l2_norm_gradient_clipping(self, clip_norm):
"""
Set l2 norm gradient clipping is not supported in TensorFlowEstimator. Please pass the
clip_norm to Estimator.from_graph.
"""
raise NotImplementedError
def shutdown(self):
"""
Close TensorFlow session and release resources.
"""
self.sess.close()
class TFOptimizerWrapper(Estimator):
def __init__(self, *, inputs, outputs, labels, loss, optimizer, clip_norm,
clip_value, metrics, updates, sess, model_dir):
self.inputs = inputs
self.outputs = outputs
self.labels = labels
self.loss = loss
if optimizer is not None:
assert isinstance(optimizer, tf.train.Optimizer), \
"optimizer is of type {}, ".format(type(optimizer)) + \
"it should be an instance of tf.train.Optimizer"
self.optimizer = ZooOptimizer(optimizer)
if clip_norm or clip_value:
gvs = self.optimizer.compute_gradients(self.loss)
if clip_norm:
gvs = [(tf.clip_by_norm(g_v[0], clip_norm), g_v[1])
for g_v in gvs]
if clip_value:
if isinstance(clip_value, tuple):
assert len(clip_value) == 2 and clip_value[0] < clip_value[1], \
"clip value should be (clip_min, clip_max)"
gvs = [(tf.clip_by_value(g_v[0], clip_value[0],
clip_value[1]), g_v[1])
for g_v in gvs]
if isinstance(clip_value, (int, float)):
assert clip_value > 0, "clip value should be larger than 0"
gvs = [(tf.clip_by_value(g_v[0], -clip_value,
clip_value), g_v[1])
for g_v in gvs]
else:
raise Exception(
"clip_value should be a tuple or one number")
self.train_op = self.optimizer.apply_gradients(gvs)
else:
self.train_op = self.optimizer.minimize(self.loss)
else:
self.optimizer = None
self.train_op = None
self.metrics = metrics
self.updates = updates
if sess is None:
self.sess = tf.Session()
self.sess.run(tf.global_variables_initializer())
else:
self.sess = sess
self.model_dir = model_dir
self.load_checkpoint = False
self.tf_optimizer = None
self.log_dir = None
self.app_name = None
def fit(self,
data,
epochs=1,
batch_size=32,
feature_cols=None,
labels_cols=None,
validation_data=None,
hard_code_batch_size=False,
session_config=None,
feed_dict=None,
checkpoint_trigger=None):
"""
Train this graph model with train data.
:param data: train data. It can be XShards, Spark DataFrame, tf.data.Dataset.
If data is XShards, each element needs to be {'x': a feature numpy array
or a tuple of feature numpy arrays, 'y': a label numpy array or a tuple of
label numpy arrays}
If data is tf.data.Dataset, each element is a tuple of input tensors.
:param epochs: number of epochs to train.
:param batch_size: total batch size for each iteration.
:param feature_cols: feature column names if train data is Spark DataFrame.
:param labels_cols: label column names if train data is Spark DataFrame.
:param validation_data: validation data. Validation data type should be the same
as train data.
:param hard_code_batch_size: whether hard code batch size for training. Default is False.
:param session_config: tensorflow session configuration for training.
Should be object of tf.ConfigProto
:param feed_dict: a dictionary. The key is TensorFlow tensor, usually a
placeholder, the value of the dictionary is a tuple of two elements. The first one of
the tuple is the value to feed to the tensor in training phase and the second one
is the value to feed to the tensor in validation phase.
:param checkpoint_trigger: when to trigger checkpoint during training.
Should be bigdl optimzer trigger, like EveryEpoch(), SeveralIteration(num_iterations),etc.
"""
assert self.labels is not None, \
"labels is None; it should not be None in training"
assert self.loss is not None, \
"loss is None; it should not be None in training"
assert self.optimizer is not None, \
"optimizer is None; it should not be None in training"
if isinstance(data, DataFrame):
assert feature_cols is not None, \
"feature columns is None; it should not be None in training"
assert labels_cols is not None, \
"label columns is None; it should not be None in training"
dataset = to_dataset(data,
batch_size=batch_size,
batch_per_thread=-1,
validation_data=validation_data,
feature_cols=feature_cols,
labels_cols=labels_cols,
hard_code_batch_size=hard_code_batch_size,
sequential_order=False,
shuffle=True)
if feed_dict is not None:
tensor_with_value = {
key: (value[0], value[1])
for key, value in feed_dict.items()
}
else:
tensor_with_value = None
self.tf_optimizer = TFOptimizer.from_train_op(
train_op=self.train_op,
loss=self.loss,
inputs=self.inputs,
labels=self.labels,
dataset=dataset,
metrics=self.metrics,
updates=self.updates,
sess=self.sess,
tensor_with_value=tensor_with_value,
session_config=session_config,
model_dir=self.model_dir)
if self.load_checkpoint:
self.tf_optimizer.load_checkpoint(self.checkpoint_path,
self.checkpoint_version)
if self.log_dir and self.app_name:
self.tf_optimizer.estimator.set_tensorboad(self.log_dir,
self.app_name)
self.tf_optimizer.optimize(end_trigger=MaxEpoch(epochs),
checkpoint_trigger=checkpoint_trigger)
return self
def predict(self,
data,
batch_size=4,
feature_cols=None,
hard_code_batch_size=False):
"""
Predict input data
:param data: data to be predicted. It can be XShards, Spark DataFrame, or tf.data.Dataset.
If data is XShards, each element needs to be {'x': a feature numpy array
or a tuple of feature numpy arrays}.
If data is tf.data.Dataset, each element is a tuple of input tensors.
:param batch_size: batch size per thread
:param feature_cols: list of feature column names if input data is Spark DataFrame.
:param hard_code_batch_size: whether to hard code batch size for prediction.
The default value is False.
:return: predicted result.
If input data is XShards or tf.data.Dataset, the predict result is a XShards,
and the schema for each result is: {'prediction': predicted numpy array or
list of predicted numpy arrays}.
If input data is Spark DataFrame, the predict result is a DataFrame which includes original
columns plus 'prediction' column. The 'prediction' column can be FloatType, VectorUDT
or Array of VectorUDT depending on model outputs shape.
"""
assert self.outputs is not None, \
"output is None, it should not be None in prediction"
if isinstance(data, DataFrame):
assert feature_cols is not None, \
"feature columns is None; it should not be None in prediction"
dataset = to_dataset(data,
batch_size=-1,
batch_per_thread=batch_size,
validation_data=None,
feature_cols=feature_cols,
labels_cols=None,
hard_code_batch_size=hard_code_batch_size,
sequential_order=True,
shuffle=False)
flat_inputs = nest.flatten(self.inputs)
flat_outputs = nest.flatten(self.outputs)
tfnet = TFNet.from_session(sess=self.sess,
inputs=flat_inputs,
outputs=flat_outputs)
predicted_rdd = tfnet.predict(dataset)
if isinstance(data, DataFrame):
return convert_predict_to_dataframe(data, predicted_rdd)
elif isinstance(data, SparkXShards) or isinstance(
data, tf.data.Dataset):
return convert_predict_to_xshard(predicted_rdd)
else:
return predicted_rdd
def evaluate(self,
data,
batch_size=32,
feature_cols=None,
labels_cols=None,
hard_code_batch_size=False):
"""
Evaluate model.
:param data: evaluation data. It can be XShards, Spark DataFrame, tf.data.Dataset.
If data is XShards, each element needs to be {'x': a feature numpy array
or a tuple of feature numpy arrays, 'y': a label numpy array or a tuple of
label numpy arrays}
If data is tf.data.Dataset, each element is a tuple of input tensors.
:param batch_size: batch size per thread.
:param feature_cols: feature_cols: feature column names if train data is Spark DataFrame.
:param labels_cols: label column names if train data is Spark DataFrame.
:param hard_code_batch_size: whether to hard code batch size for evaluation.
:return: evaluation result as a dictionary of {'metric name': metric value}
"""
assert self.metrics is not None, \
"metrics is None, it should not be None in evaluate"
if isinstance(data, DataFrame):
assert feature_cols is not None, \
"feature columns is None; it should not be None in evaluation"
assert labels_cols is not None, \
"label columns is None; it should not be None in evaluation"
dataset = to_dataset(data,
batch_size=-1,
batch_per_thread=batch_size,
validation_data=None,
feature_cols=feature_cols,
labels_cols=labels_cols,
hard_code_batch_size=hard_code_batch_size,
sequential_order=True,
shuffle=False)
flat_inputs = nest.flatten(self.inputs)
flat_labels = nest.flatten(self.labels)
return evaluate_metrics(flat_inputs + flat_labels,
sess=self.sess,
dataset=dataset,
metrics=self.metrics)
def save_tf_checkpoint(self, path):
save_tf_checkpoint(self.sess, path)
class TFOptimizerWrapper(Estimator):
def __init__(self, *, inputs, outputs, labels, loss, optimizer, metrics,
updates, sess, model_dir):
self.inputs = inputs
self.outputs = outputs
self.labels = labels
self.loss = loss
if optimizer is not None:
assert isinstance(optimizer, tf.train.Optimizer), \
"optimizer is of type {}, ".format(type(self.optimizer)) + \
"it should be an instance of tf.train.Optimizer"
self.optimizer = ZooOptimizer(optimizer)
self.train_op = self.optimizer.minimize(self.loss)
else:
self.optimizer = None
self.train_op = None
self.metrics = metrics
self.updates = updates
if sess is None:
self.sess = tf.Session()
self.sess.run(tf.global_variables_initializer())
else:
self.sess = sess
self.model_dir = model_dir
def fit(self,
data,
epochs=1,
batch_size=32,
feature_cols=None,
labels_cols=None,
validation_data=None,
hard_code_batch_size=False,
session_config=None,
feed_dict=None):
assert self.labels is not None, \
"labels is None; it should not be None in training"
assert self.loss is not None, \
"loss is None; it should not be None in training"
assert self.optimizer is not None, \
"optimizer is None; it should not be None in training"
if isinstance(data, DataFrame):
assert feature_cols is not None, \
"feature columns is None; it should not be None in training"
assert labels_cols is not None, \
"label columns is None; it should not be None in training"
dataset = to_dataset(data,
batch_size=batch_size,
batch_per_thread=-1,
validation_data=validation_data,
feature_cols=feature_cols,
labels_cols=labels_cols,
hard_code_batch_size=hard_code_batch_size,
sequential_order=False,
shuffle=True)
if feed_dict is not None:
tensor_with_value = {
key: (value, value)
for key, value in feed_dict.items()
}
else:
tensor_with_value = None
optimizer = TFOptimizer.from_train_op(
train_op=self.train_op,
loss=self.loss,
inputs=self.inputs,
labels=self.labels,
dataset=dataset,
metrics=self.metrics,
updates=self.updates,
sess=self.sess,
tensor_with_value=tensor_with_value,
session_config=session_config,
model_dir=self.model_dir)
optimizer.optimize(end_trigger=MaxEpoch(epochs))
return self
def predict(self,
data,
batch_size=4,
feature_cols=None,
hard_code_batch_size=False):
assert self.outputs is not None, \
"output is None, it should not be None in prediction"
if isinstance(data, DataFrame):
assert feature_cols is not None, \
"feature columns is None; it should not be None in prediction"
dataset = to_dataset(data,
batch_size=-1,
batch_per_thread=batch_size,
validation_data=None,
feature_cols=feature_cols,
labels_cols=None,
hard_code_batch_size=hard_code_batch_size,
sequential_order=True,
shuffle=False)
flat_inputs = nest.flatten(self.inputs)
flat_outputs = nest.flatten(self.outputs)
tfnet = TFNet.from_session(sess=self.sess,
inputs=flat_inputs,
outputs=flat_outputs)
predicted_rdd = tfnet.predict(dataset)
if isinstance(data, DataFrame):
return convert_predict_to_dataframe(data, predicted_rdd)
else:
return predicted_rdd
def evaluate(self,
data,
batch_size=32,
feature_cols=None,
labels_cols=None,
hard_code_batch_size=False):
assert self.metrics is not None, \
"metrics is None, it should not be None in evaluate"
if isinstance(data, DataFrame):
assert feature_cols is not None, \
"feature columns is None; it should not be None in evaluation"
assert labels_cols is not None, \
"label columns is None; it should not be None in evaluation"
dataset = to_dataset(data,
batch_size=-1,
batch_per_thread=batch_size,
validation_data=None,
feature_cols=feature_cols,
labels_cols=labels_cols,
hard_code_batch_size=hard_code_batch_size,
sequential_order=True,
shuffle=False)
flat_inputs = nest.flatten(self.inputs)
flat_labels = nest.flatten(self.labels)
return evaluate_metrics(flat_inputs + flat_labels,
sess=self.sess,
dataset=dataset,
metrics=self.metrics)
class TFOptimizerWrapper(Estimator):
def __init__(self, *, inputs, outputs, labels, loss,
optimizer,
metrics,
updates, sess,
model_dir):
self.inputs = inputs
self.outputs = outputs
self.labels = labels
self.loss = loss
if optimizer is not None:
assert isinstance(optimizer, tf.train.Optimizer), \
"optimizer is of type {}, ".format(type(self.optimizer)) + \
"it should be an instance of tf.train.Optimizer"
self.optimizer = ZooOptimizer(optimizer)
self.train_op = self.optimizer.minimize(self.loss)
else:
self.optimizer = None
self.train_op = None
self.metrics = metrics
self.updates = updates
if sess is None:
self.sess = tf.Session()
self.sess.run(tf.global_variables_initializer())
else:
self.sess = sess
self.model_dir = model_dir
def fit(self, data, steps,
batch_size=32,
validation_data=None,
feed_dict=None,
session_config=None):
assert self.labels is not None, \
"labels is None; it should not be None in training"
assert self.loss is not None, \
"loss is None; it should not be None in training"
assert self.optimizer is not None, \
"optimizer is None; it not None in training"
if isinstance(data, SparkXShards):
dataset = _xshards_to_tf_dataset(data,
batch_size=batch_size,
validation_data_shard=validation_data)
elif isinstance(data, Dataset):
dataset = TFDataDataset2(data, batch_size=batch_size,
batch_per_thread=-1,
validation_dataset=validation_data)
else:
raise ValueError("data type {} is not supported; "
"it must be created by zoo.orca.data.package")
if feed_dict is not None:
tensor_with_value = {key: (value, value) for key, value in feed_dict.items()}
else:
tensor_with_value = None
optimizer = TFOptimizer.from_train_op(
train_op=self.train_op,
loss=self.loss,
inputs=self.inputs,
labels=self.labels,
dataset=dataset,
metrics=self.metrics,
updates=self.updates, sess=self.sess,
tensor_with_value=tensor_with_value,
session_config=session_config,
model_dir=self.model_dir)
optimizer.optimize(end_trigger=MaxIteration(steps))
return self
def predict(self, data, batch_size=32):
assert self.outputs is not None, \
"output is None, it should not be None in prediction"
if isinstance(data, SparkXShards):
dataset = _xshards_to_tf_dataset(data,
batch_per_thread=batch_size)
elif isinstance(data, Dataset):
dataset = TFDataDataset2(data, batch_size=-1,
batch_per_thread=batch_size)
else:
raise ValueError("data must be a SparkXShards or an orca.data.tf.Dataset")
flat_inputs = nest.flatten(self.inputs)
flat_outputs = nest.flatten(self.outputs)
tfnet = TFNet.from_session(sess=self.sess, inputs=flat_inputs, outputs=flat_outputs)
return tfnet.predict(dataset)
class TFOptimizerWrapper(Estimator):
def __init__(self, *, inputs, outputs, labels, loss,
optimizer, clip_norm, clip_value,
metrics,
updates, sess,
model_dir
):
self.inputs = inputs
self.outputs = outputs
self.labels = labels
self.loss = loss
if optimizer is not None:
assert isinstance(optimizer, tf.train.Optimizer), \
"optimizer is of type {}, ".format(type(self.optimizer)) + \
"it should be an instance of tf.train.Optimizer"
self.optimizer = ZooOptimizer(optimizer)
if clip_norm or clip_value:
gvs = self.optimizer.compute_gradients(self.loss)
if clip_norm:
gvs = [(tf.clip_by_norm(g_v[0], clip_norm), g_v[1]) for g_v in gvs]
if clip_value:
if isinstance(clip_value, tuple):
assert len(clip_value) == 2 and clip_value[0] < clip_value[1], \
"clip value should be (clip_min, clip_max)"
gvs = [(tf.clip_by_value(g_v[0], clip_value[0], clip_value[1]), g_v[1])
for g_v in gvs]
if isinstance(clip_value, (int, float)):
assert clip_value > 0, "clip value should be larger than 0"
gvs = [(tf.clip_by_value(g_v[0], -clip_value, clip_value), g_v[1])
for g_v in gvs]
else:
raise Exception("clip_value should be a tuple or one number")
self.train_op = self.optimizer.apply_gradients(gvs)
else:
self.train_op = self.optimizer.minimize(self.loss)
else:
self.optimizer = None
self.train_op = None
self.metrics = metrics
self.updates = updates
if sess is None:
self.sess = tf.Session()
self.sess.run(tf.global_variables_initializer())
else:
self.sess = sess
self.model_dir = model_dir
def fit(self, data,
epochs=1,
batch_size=32,
feature_cols=None,
labels_cols=None,
validation_data=None,
hard_code_batch_size=False,
session_config=None,
feed_dict=None
):
assert self.labels is not None, \
"labels is None; it should not be None in training"
assert self.loss is not None, \
"loss is None; it should not be None in training"
assert self.optimizer is not None, \
"optimizer is None; it should not be None in training"
if isinstance(data, DataFrame):
assert feature_cols is not None, \
"feature columns is None; it should not be None in training"
assert labels_cols is not None, \
"label columns is None; it should not be None in training"
dataset = to_dataset(data, batch_size=batch_size, batch_per_thread=-1,
validation_data=validation_data,
feature_cols=feature_cols, labels_cols=labels_cols,
hard_code_batch_size=hard_code_batch_size,
sequential_order=False, shuffle=True
)
if feed_dict is not None:
tensor_with_value = {key: (value, value) for key, value in feed_dict.items()}
else:
tensor_with_value = None
optimizer = TFOptimizer.from_train_op(
train_op=self.train_op,
loss=self.loss,
inputs=self.inputs,
labels=self.labels,
dataset=dataset,
metrics=self.metrics,
updates=self.updates, sess=self.sess,
tensor_with_value=tensor_with_value,
session_config=session_config,
model_dir=self.model_dir)
optimizer.optimize(end_trigger=MaxEpoch(epochs))
return self
def predict(self, data, batch_size=4,
feature_cols=None,
hard_code_batch_size=False
):
assert self.outputs is not None, \
"output is None, it should not be None in prediction"
if isinstance(data, DataFrame):
assert feature_cols is not None, \
"feature columns is None; it should not be None in prediction"
dataset = to_dataset(data, batch_size=-1, batch_per_thread=batch_size,
validation_data=None,
feature_cols=feature_cols, labels_cols=None,
hard_code_batch_size=hard_code_batch_size,
sequential_order=True,
shuffle=False
)
flat_inputs = nest.flatten(self.inputs)
flat_outputs = nest.flatten(self.outputs)
tfnet = TFNet.from_session(sess=self.sess, inputs=flat_inputs, outputs=flat_outputs)
predicted_rdd = tfnet.predict(dataset)
if isinstance(data, DataFrame):
return convert_predict_to_dataframe(data, predicted_rdd)
else:
return predicted_rdd
def evaluate(self, data, batch_size=32,
feature_cols=None,
labels_cols=None,
hard_code_batch_size=False
):
assert self.metrics is not None, \
"metrics is None, it should not be None in evaluate"
if isinstance(data, DataFrame):
assert feature_cols is not None, \
"feature columns is None; it should not be None in evaluation"
assert labels_cols is not None, \
"label columns is None; it should not be None in evaluation"
dataset = to_dataset(data, batch_size=-1, batch_per_thread=batch_size,
validation_data=None,
feature_cols=feature_cols, labels_cols=labels_cols,
hard_code_batch_size=hard_code_batch_size,
sequential_order=True,
shuffle=False
)
flat_inputs = nest.flatten(self.inputs)
flat_labels = nest.flatten(self.labels)
return evaluate_metrics(flat_inputs + flat_labels,
sess=self.sess,
dataset=dataset, metrics=self.metrics)