def test_estimator_for_imageset(self):
model_fn = self.create_model_fn()
input_fn = self.create_imageset_input_fn()
estimator = TFEstimator(model_fn, tf.train.AdamOptimizer())
estimator.train(input_fn, steps=1)
estimator.evaluate(input_fn, ["acc"])
results = estimator.predict(input_fn).get_predict().collect()
assert all(r[1] is not None for r in results)
def test_estimator_without_batch(self):
def model_fn(features, labels, mode):
assert features.shape.ndims == 1
if labels is not None:
assert labels.shape.ndims == 0
features = tf.expand_dims(features, axis=0)
h1 = tf.layers.dense(features, 64, activation=tf.nn.relu)
h2 = tf.layers.dense(h1, 64, activation=tf.nn.relu)
logits = tf.layers.dense(h2, 10)
if mode == tf.estimator.ModeKeys.EVAL or mode == tf.estimator.ModeKeys.TRAIN:
labels = tf.expand_dims(labels, axis=0)
loss = tf.reduce_mean(
tf.losses.sparse_softmax_cross_entropy(logits=logits,
labels=labels))
return TFEstimatorSpec(mode, predictions=logits, loss=loss)
else:
return TFEstimatorSpec(mode, predictions=logits)
def input_fn(mode):
np.random.seed(20)
x = np.random.rand(20, 10)
y = np.random.randint(0, 10, (20))
rdd_x = self.sc.parallelize(x)
rdd_y = self.sc.parallelize(y)
rdd = rdd_x.zip(rdd_y)
if mode == tf.estimator.ModeKeys.TRAIN or mode == tf.estimator.ModeKeys.EVAL:
dataset = TFDataset.from_rdd(rdd,
features=(tf.float32, [10]),
labels=(tf.int32, []))
else:
dataset = TFDataset.from_rdd(rdd_x,
features=(tf.float32, [10]))
return dataset
estimator = TFEstimator(model_fn, tf.train.AdamOptimizer())
self.intercept(
lambda: estimator.train(input_fn, steps=1),
"The batch_size of TFDataset must be specified when used for training."
)
estimator.evaluate(input_fn, ["acc"])
estimator.predict(input_fn).collect()
def test_init_TFDataset_from_ndarrays(self):
model_fn = self.create_model_fn()
def input_fn(mode):
x = np.random.rand(20, 10)
y = np.random.randint(0, 10, (20,))
if mode == tf.estimator.ModeKeys.TRAIN:
return TFDataset.from_ndarrays((x, y), batch_size=8)
elif mode == tf.estimator.ModeKeys.EVAL:
return TFDataset.from_ndarrays((x, y), batch_per_thread=1)
else:
return TFDataset.from_ndarrays(x, batch_per_thread=1)
estimator = TFEstimator(model_fn, tf.train.AdamOptimizer())
estimator.train(input_fn, 10)
estimator.evaluate(input_fn, ["acc"])
estimator.predict(input_fn)
def main():
sc = init_nncontext()
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))
return TFEstimatorSpec(mode, predictions=logits, loss=loss)
else:
return TFEstimatorSpec(mode, predictions=logits)
def input_fn(mode):
if mode == tf.estimator.ModeKeys.TRAIN:
training_rdd = get_data_rdd("train", sc)
dataset = TFDataset.from_rdd(training_rdd,
features=(tf.float32, [28, 28, 1]),
labels=(tf.int32, []),
batch_size=320)
elif mode == tf.estimator.ModeKeys.EVAL:
testing_rdd = get_data_rdd("test", sc)
dataset = TFDataset.from_rdd(testing_rdd,
features=(tf.float32, [28, 28, 1]),
labels=(tf.int32, []),
batch_size=320)
else:
testing_rdd = get_data_rdd("test", sc).map(lambda x: x[0])
dataset = TFDataset.from_rdd(testing_rdd,
features=(tf.float32, [28, 28, 1]),
batch_per_thread=80)
return dataset
estimator = TFEstimator(model_fn,
tf.train.AdamOptimizer(),
model_dir="/tmp/estimator")
estimator.train(input_fn, steps=60000 // 320)
metrics = estimator.evaluate(input_fn, ["acc"])
print(metrics)
predictions = estimator.predict(input_fn)
print(predictions.first())
def test_evaluating(self):
model_fn = self.create_model_fn()
input_fn = self.create_input_fn()
estimator = TFEstimator(model_fn, tf.train.AdamOptimizer())
eval_results = estimator.evaluate(input_fn, ["acc"])
assert len(eval_results) > 0