Commit 739193ea authored by Amelie Royer's avatar Amelie Royer

Adding CPU mode

parent 88e50fc3
...@@ -44,7 +44,7 @@ parser.add_argument('-m', '--nr_logistic_mix', type=int, default=10, help='Numbe ...@@ -44,7 +44,7 @@ parser.add_argument('-m', '--nr_logistic_mix', type=int, default=10, help='Numbe
parser.add_argument('-lr', '--learning_rate', type=float, default=1e-3, help='Initial learning rate') parser.add_argument('-lr', '--learning_rate', type=float, default=1e-3, help='Initial learning rate')
parser.add_argument('-ld', '--lr_decay', type=float, default=0.99999, help='Learning rate decay, applied every decay epoch') parser.add_argument('-ld', '--lr_decay', type=float, default=0.99999, help='Learning rate decay, applied every decay epoch')
parser.add_argument('-g', '--nr_gpus', type=int, default=1, help='Number of GPUs') parser.add_argument('-g', '--nr_gpus', type=int, default=1, help='Number of GPUs. If 0, will run on one CPU.')
parser.add_argument('-b', '--batch_size', type=int, default=32, help='Batch size per GPU (train)') parser.add_argument('-b', '--batch_size', type=int, default=32, help='Batch size per GPU (train)')
parser.add_argument('-bt', '--test_batch_size', type=int, default=16, help='Batch size per GPU (test)') parser.add_argument('-bt', '--test_batch_size', type=int, default=16, help='Batch size per GPU (test)')
parser.add_argument('-bin', '--init_batch_size', type=int, default=300, help='Batch size per GPU for the data-dependent initialization.') parser.add_argument('-bin', '--init_batch_size', type=int, default=300, help='Batch size per GPU for the data-dependent initialization.')
...@@ -63,6 +63,11 @@ parser.add_argument('-nt', '--test_epochs', type=int, default=1, help='Log-likel ...@@ -63,6 +63,11 @@ parser.add_argument('-nt', '--test_epochs', type=int, default=1, help='Log-likel
args = parser.parse_args() args = parser.parse_args()
### Init ### Init
use_gpu = True
if args.nr_gpus == 0:
args.nr_gpus = 1
use_gpu = False
if args.color not in ['RGB', 'lab']: if args.color not in ['RGB', 'lab']:
raise ValueError("Unknown color mode", args.color) raise ValueError("Unknown color mode", args.color)
...@@ -177,7 +182,7 @@ with tf.name_scope('MAIN'): ...@@ -177,7 +182,7 @@ with tf.name_scope('MAIN'):
x_in = tf.identity(xs[i], name="color_image") # x_in RGB (BxWxHx3) x_in = tf.identity(xs[i], name="color_image") # x_in RGB (BxWxHx3)
x_gray = color_to_gray(xs_clr[i], colorspace=args.color) # x_gray colorspace (BxWxHx1) x_gray = color_to_gray(xs_clr[i], colorspace=args.color) # x_gray colorspace (BxWxHx1)
with tf.device('/gpu:%i' % i): with tf.device('/%s:%i' % ('gpu' if use_gpu else 'cpu', i)):
picolor_out, embedding_out = picolor_with_scope(downsample_tf(x_in, args.downsample), picolor_out, embedding_out = picolor_with_scope(downsample_tf(x_in, args.downsample),
x_gray, x_gray,
return_embedding=True) return_embedding=True)
...@@ -253,7 +258,7 @@ with tf.name_scope("sampling"): ...@@ -253,7 +258,7 @@ with tf.name_scope("sampling"):
embedding_cache = [None] * args.nr_gpus embedding_cache = [None] * args.nr_gpus
with tf.name_scope('picolor_sampling'): with tf.name_scope('picolor_sampling'):
for i in range(args.nr_gpus): for i in range(args.nr_gpus):
with tf.device('/gpu:%d' % i): with tf.device('/%s:%d' % ('gpu' if use_gpu else 'cpu', i)):
picolor_out, embedding_cache[i] = picolor_with_scope(x_canvas_gens[i], picolor_out, embedding_cache[i] = picolor_with_scope(x_canvas_gens[i],
x_gray_gens[i], x_gray_gens[i],
return_embedding=True) return_embedding=True)
......
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