class TFBenchModel(Model):
def __init__(self, batch=64, use_cpus=False):
image_shape = [batch, 224, 224, 3]
labels_shape = [batch]
# Synthetic image should be within [0, 255].
images = tf.truncated_normal(
image_shape,
dtype=tf.float32,
mean=127,
stddev=60,
name='synthetic_images')
# Minor hack to avoid H2D copy when using synthetic data
inputs = tf.contrib.framework.local_variable(
images, name='gpu_cached_images')
labels = tf.random_uniform(
labels_shape,
minval=0,
maxval=999,
dtype=tf.int32,
name='synthetic_labels')
model = model_config.get_model_config("resnet101", MockDataset())
logits, aux = model.build_network(
inputs, data_format=use_cpus and "NHWC" or "NCHW")
loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
logits=logits, labels=labels)
# Implement model interface
self.loss = tf.reduce_mean(loss, name='xentropy-loss')
self.optimizer = tf.train.GradientDescentOptimizer(1e-6)
self.variables = TensorFlowVariables(self.loss,
tf.get_default_session())
def get_loss(self):
return self.loss
def get_optimizer(self):
return self.optimizer
def get_feed_dict(self):
return {}
def get_weights(self):
return self.variables.get_flat()
def set_weights(self, weights):
self.variables.set_flat(weights)
class MNISTModel(Model):
def __init__(self):
# Import data
error = None
for _ in range(10):
try:
self.mnist = input_data.read_data_sets(
"/tmp/tensorflow/mnist/input_data", one_hot=True)
error = None
break
except Exception as e:
error = e
time.sleep(5)
if error:
raise ValueError("Failed to import data", error)
# Set seed and build layers
tf.set_random_seed(0)
self.x = tf.placeholder(tf.float32, [None, 784], name="x")
self.y_ = tf.placeholder(tf.float32, [None, 10], name="y_")
y_conv, self.keep_prob = deepnn(self.x)
# Need to define loss and optimizer attributes
self.loss = tf.reduce_mean(
tf.nn.softmax_cross_entropy_with_logits(labels=self.y_,
logits=y_conv))
self.optimizer = tf.train.AdamOptimizer(1e-4)
self.variables = TensorFlowVariables(self.loss,
tf.get_default_session())
# For evaluating test accuracy
correct_prediction = tf.equal(tf.argmax(y_conv, 1),
tf.argmax(self.y_, 1))
self.accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
def get_loss(self):
return self.loss
def get_optimizer(self):
return self.optimizer
def get_variables(self):
return self.variables
def get_feed_dict(self):
batch = self.mnist.train.next_batch(50)
return {
self.x: batch[0],
self.y_: batch[1],
self.keep_prob: 0.5,
}
def get_metrics(self):
accuracy = self.accuracy.eval(
feed_dict={
self.x: self.mnist.test.images,
self.y_: self.mnist.test.labels,
self.keep_prob: 1.0,
})
return {"accuracy": accuracy}
def get_weights(self):
return self.variables.get_flat()
def set_weights(self, weights):
self.variables.set_flat(weights)