def test_get_optimizer():
optimizer_keys = ['adam', 'adagrad', 'momentum', 'ftrl']
invalid_optimizer_keys = ['adddam']
for optimizer_key in optimizer_keys:
get_optimizer(optimizer_key=optimizer_key, learning_rate=0.3)
for invalid_optimizer_key in invalid_optimizer_keys:
with pytest.raises(ValueError, match="Invalid optimizer_key :"):
get_optimizer(optimizer_key=invalid_optimizer_key, learning_rate=0.3)
def test_get_optimizer():
from submarine.ml.tensorflow.optimizer import get_optimizer
optimizer_keys = ["adam", "adagrad", "momentum", "ftrl"]
invalid_optimizer_keys = ["adddam"]
for optimizer_key in optimizer_keys:
get_optimizer(optimizer_key=optimizer_key, learning_rate=0.3)
for invalid_optimizer_key in invalid_optimizer_keys:
with pytest.raises(ValueError, match="Invalid optimizer_key :"):
get_optimizer(optimizer_key=invalid_optimizer_key,
learning_rate=0.3)
def get_estimator_spec(logit, labels, mode, params):
"""
Returns `EstimatorSpec` that a model_fn can return.
:param logit: logits `Tensor` to be used.
:param labels: Labels `Tensor`, or `dict` of same.
:param mode: Estimator's `ModeKeys`.
:param params: Optional dict of hyperparameters. Will receive what is passed to Estimator
in params parameter.
:return:
"""
learning_rate = params["training"]["learning_rate"]
optimizer = params["training"]["optimizer"]
metric = params["output"]["metric"]
output = tf.sigmoid(logit)
predictions = {"probabilities": output}
export_outputs = {
# https://github.com/psf/black/issues/2434
# fmt: off
tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
tf.estimator.export.PredictOutput(predictions)
# fmt: on
}
# Provide an estimator spec for `ModeKeys.PREDICT`
if mode == tf.estimator.ModeKeys.PREDICT:
return tf.estimator.EstimatorSpec(mode=mode,
predictions=predictions,
export_outputs=export_outputs)
with tf.name_scope("Loss"):
loss = tf.reduce_mean(
tf.nn.sigmoid_cross_entropy_with_logits(logits=logit,
labels=labels))
# Provide an estimator spec for `ModeKeys.EVAL`
eval_metric_ops = {}
if metric == "auc":
eval_metric_ops["auc"] = tf.metrics.auc(labels, output)
else:
raise TypeError("Invalid metric :", metric)
if mode == tf.estimator.ModeKeys.EVAL:
return tf.estimator.EstimatorSpec(mode=mode,
predictions=predictions,
loss=loss,
eval_metric_ops=eval_metric_ops)
with tf.name_scope("Train"):
op = get_optimizer(optimizer, learning_rate)
train_op = op.minimize(loss, global_step=tf.train.get_global_step())
# Provide an estimator spec for `ModeKeys.TRAIN` modes
if mode == tf.estimator.ModeKeys.TRAIN:
return tf.estimator.EstimatorSpec(mode=mode,
predictions=predictions,
loss=loss,
train_op=train_op)