Commit fc7342bf authored by derekjchow's avatar derekjchow Committed by Sergio Guadarrama
Browse files

Update model_deploy. (#1557)

Fix slow down with only 1 GPU
parent 001a2602
...@@ -103,8 +103,6 @@ import collections ...@@ -103,8 +103,6 @@ import collections
import tensorflow as tf import tensorflow as tf
from tensorflow.python.ops import control_flow_ops
slim = tf.contrib.slim slim = tf.contrib.slim
...@@ -594,8 +592,7 @@ class DeploymentConfig(object): ...@@ -594,8 +592,7 @@ class DeploymentConfig(object):
if self._clone_on_cpu: if self._clone_on_cpu:
device += '/device:CPU:0' device += '/device:CPU:0'
else: else:
if self._num_clones > 1: device += '/device:GPU:%d' % clone_index
device += '/device:GPU:%d' % clone_index
return device return device
def clone_scope(self, clone_index): def clone_scope(self, clone_index):
......
...@@ -33,7 +33,7 @@ class DeploymentConfigTest(tf.test.TestCase): ...@@ -33,7 +33,7 @@ class DeploymentConfigTest(tf.test.TestCase):
self.assertEqual(slim.get_variables(), []) self.assertEqual(slim.get_variables(), [])
self.assertEqual(deploy_config.caching_device(), None) self.assertEqual(deploy_config.caching_device(), None)
self.assertDeviceEqual(deploy_config.clone_device(0), '') self.assertDeviceEqual(deploy_config.clone_device(0), 'GPU:0')
self.assertEqual(deploy_config.clone_scope(0), '') self.assertEqual(deploy_config.clone_scope(0), '')
self.assertDeviceEqual(deploy_config.optimizer_device(), 'CPU:0') self.assertDeviceEqual(deploy_config.optimizer_device(), 'CPU:0')
self.assertDeviceEqual(deploy_config.inputs_device(), 'CPU:0') self.assertDeviceEqual(deploy_config.inputs_device(), 'CPU:0')
...@@ -65,7 +65,7 @@ class DeploymentConfigTest(tf.test.TestCase): ...@@ -65,7 +65,7 @@ class DeploymentConfigTest(tf.test.TestCase):
deploy_config = model_deploy.DeploymentConfig(num_clones=1, num_ps_tasks=1) deploy_config = model_deploy.DeploymentConfig(num_clones=1, num_ps_tasks=1)
self.assertDeviceEqual(deploy_config.clone_device(0), self.assertDeviceEqual(deploy_config.clone_device(0),
'/job:worker') '/job:worker/device:GPU:0')
self.assertEqual(deploy_config.clone_scope(0), '') self.assertEqual(deploy_config.clone_scope(0), '')
self.assertDeviceEqual(deploy_config.optimizer_device(), self.assertDeviceEqual(deploy_config.optimizer_device(),
'/job:worker/device:CPU:0') '/job:worker/device:CPU:0')
...@@ -105,7 +105,7 @@ class DeploymentConfigTest(tf.test.TestCase): ...@@ -105,7 +105,7 @@ class DeploymentConfigTest(tf.test.TestCase):
num_ps_tasks=2) num_ps_tasks=2)
self.assertDeviceEqual(deploy_config.clone_device(0), self.assertDeviceEqual(deploy_config.clone_device(0),
'/job:worker') '/job:worker/device:GPU:0')
self.assertEqual(deploy_config.clone_scope(0), '') self.assertEqual(deploy_config.clone_scope(0), '')
self.assertDeviceEqual(deploy_config.optimizer_device(), self.assertDeviceEqual(deploy_config.optimizer_device(),
'/job:worker/device:CPU:0') '/job:worker/device:CPU:0')
...@@ -201,7 +201,7 @@ class CreatecloneTest(tf.test.TestCase): ...@@ -201,7 +201,7 @@ class CreatecloneTest(tf.test.TestCase):
self.assertEqual(clone.outputs.op.name, self.assertEqual(clone.outputs.op.name,
'LogisticClassifier/fully_connected/Sigmoid') 'LogisticClassifier/fully_connected/Sigmoid')
self.assertEqual(clone.scope, '') self.assertEqual(clone.scope, '')
self.assertDeviceEqual(clone.device, '') self.assertDeviceEqual(clone.device, 'GPU:0')
self.assertEqual(len(slim.losses.get_losses()), 1) self.assertEqual(len(slim.losses.get_losses()), 1)
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
self.assertEqual(update_ops, []) self.assertEqual(update_ops, [])
...@@ -227,7 +227,7 @@ class CreatecloneTest(tf.test.TestCase): ...@@ -227,7 +227,7 @@ class CreatecloneTest(tf.test.TestCase):
self.assertEqual(clone.outputs.op.name, self.assertEqual(clone.outputs.op.name,
'BatchNormClassifier/fully_connected/Sigmoid') 'BatchNormClassifier/fully_connected/Sigmoid')
self.assertEqual(clone.scope, '') self.assertEqual(clone.scope, '')
self.assertDeviceEqual(clone.device, '') self.assertDeviceEqual(clone.device, 'GPU:0')
self.assertEqual(len(slim.losses.get_losses()), 1) self.assertEqual(len(slim.losses.get_losses()), 1)
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
self.assertEqual(len(update_ops), 2) self.assertEqual(len(update_ops), 2)
...@@ -278,7 +278,7 @@ class CreatecloneTest(tf.test.TestCase): ...@@ -278,7 +278,7 @@ class CreatecloneTest(tf.test.TestCase):
clone = clones[0] clone = clones[0]
self.assertEqual(clone.outputs.op.name, self.assertEqual(clone.outputs.op.name,
'BatchNormClassifier/fully_connected/Sigmoid') 'BatchNormClassifier/fully_connected/Sigmoid')
self.assertDeviceEqual(clone.device, '/job:worker') self.assertDeviceEqual(clone.device, '/job:worker/device:GPU:0')
self.assertEqual(clone.scope, '') self.assertEqual(clone.scope, '')
self.assertEqual(len(slim.get_variables()), 5) self.assertEqual(len(slim.get_variables()), 5)
for v in slim.get_variables(): for v in slim.get_variables():
...@@ -350,7 +350,7 @@ class OptimizeclonesTest(tf.test.TestCase): ...@@ -350,7 +350,7 @@ class OptimizeclonesTest(tf.test.TestCase):
self.assertEqual(len(grads_and_vars), len(tf.trainable_variables())) self.assertEqual(len(grads_and_vars), len(tf.trainable_variables()))
self.assertEqual(total_loss.op.name, 'total_loss') self.assertEqual(total_loss.op.name, 'total_loss')
for g, v in grads_and_vars: for g, v in grads_and_vars:
self.assertDeviceEqual(g.device, '') self.assertDeviceEqual(g.device, 'GPU:0')
self.assertDeviceEqual(v.device, 'CPU:0') self.assertDeviceEqual(v.device, 'CPU:0')
def testCreateSingleclone(self): def testCreateSingleclone(self):
...@@ -376,7 +376,7 @@ class OptimizeclonesTest(tf.test.TestCase): ...@@ -376,7 +376,7 @@ class OptimizeclonesTest(tf.test.TestCase):
self.assertEqual(len(grads_and_vars), len(tf.trainable_variables())) self.assertEqual(len(grads_and_vars), len(tf.trainable_variables()))
self.assertEqual(total_loss.op.name, 'total_loss') self.assertEqual(total_loss.op.name, 'total_loss')
for g, v in grads_and_vars: for g, v in grads_and_vars:
self.assertDeviceEqual(g.device, '') self.assertDeviceEqual(g.device, 'GPU:0')
self.assertDeviceEqual(v.device, 'CPU:0') self.assertDeviceEqual(v.device, 'CPU:0')
def testCreateMulticlone(self): def testCreateMulticlone(self):
...@@ -458,7 +458,7 @@ class OptimizeclonesTest(tf.test.TestCase): ...@@ -458,7 +458,7 @@ class OptimizeclonesTest(tf.test.TestCase):
self.assertEqual(len(grads_and_vars), len(tf.trainable_variables())) self.assertEqual(len(grads_and_vars), len(tf.trainable_variables()))
self.assertEqual(total_loss.op.name, 'total_loss') self.assertEqual(total_loss.op.name, 'total_loss')
for g, v in grads_and_vars: for g, v in grads_and_vars:
self.assertDeviceEqual(g.device, '/job:worker') self.assertDeviceEqual(g.device, '/job:worker/device:GPU:0')
self.assertDeviceEqual(v.device, '/job:ps/task:0/CPU:0') self.assertDeviceEqual(v.device, '/job:ps/task:0/CPU:0')
...@@ -515,7 +515,7 @@ class DeployTest(tf.test.TestCase): ...@@ -515,7 +515,7 @@ class DeployTest(tf.test.TestCase):
for _ in range(10): for _ in range(10):
sess.run(model.train_op) sess.run(model.train_op)
final_loss = sess.run(model.total_loss) final_loss = sess.run(model.total_loss)
self.assertLess(final_loss, initial_loss / 10.0) self.assertLess(final_loss, initial_loss / 5.0)
final_mean, final_variance = sess.run([moving_mean, final_mean, final_variance = sess.run([moving_mean,
moving_variance]) moving_variance])
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment