Code modification for testcases of various network models in directory example

example/bayesian-methods/bdk_demo.py

@@ -156,42 +156,42 @@ def get_toy_sym(teacher=True, teacher_noise_precision=None):
     return net
 
 
-def dev():
-    return mx.gpu()
+def dev(gpu_id=None):
+    return mx.gpu(gpu_id) if gpu_id else mx.cpu()
 
 
-def run_mnist_SGD(training_num=50000):
+def run_mnist_SGD(training_num=50000, gpu_id=None):
     X, Y, X_test, Y_test = load_mnist(training_num)
     minibatch_size = 100
     net = get_mnist_sym()
     data_shape = (minibatch_size,) + X.shape[1::]
-    data_inputs = {'data': nd.zeros(data_shape, ctx=dev()),
-                   'softmax_label': nd.zeros((minibatch_size,), ctx=dev())}
+    data_inputs = {'data': nd.zeros(data_shape, ctx=dev(gpu_id)),
+                   'softmax_label': nd.zeros((minibatch_size,), ctx=dev(gpu_id))}
     initializer = mx.init.Xavier(factor_type="in", magnitude=2.34)
-    exe, exe_params, _ = SGD(sym=net, dev=dev(), data_inputs=data_inputs, X=X, Y=Y,
+    exe, exe_params, _ = SGD(sym=net, dev=dev(gpu_id), data_inputs=data_inputs, X=X, Y=Y,
                              X_test=X_test, Y_test=Y_test,
                              total_iter_num=1000000,
                              initializer=initializer,
                              lr=5E-6, prior_precision=1.0, minibatch_size=100)
 
 
-def run_mnist_SGLD(training_num=50000):
+def run_mnist_SGLD(training_num=50000, gpu_id=None):
     X, Y, X_test, Y_test = load_mnist(training_num)
     minibatch_size = 100
     net = get_mnist_sym()
     data_shape = (minibatch_size,) + X.shape[1::]
-    data_inputs = {'data': nd.zeros(data_shape, ctx=dev()),
-                   'softmax_label': nd.zeros((minibatch_size,), ctx=dev())}
+    data_inputs = {'data': nd.zeros(data_shape, ctx=dev(gpu_id)),
+                   'softmax_label': nd.zeros((minibatch_size,), ctx=dev(gpu_id))}
     initializer = mx.init.Xavier(factor_type="in", magnitude=2.34)
-    exe, sample_pool = SGLD(sym=net, dev=dev(), data_inputs=data_inputs, X=X, Y=Y,
+    exe, sample_pool = SGLD(sym=net, dev=dev(gpu_id), data_inputs=data_inputs, X=X, Y=Y,
                             X_test=X_test, Y_test=Y_test,
                             total_iter_num=1000000,
                             initializer=initializer,
                             learning_rate=4E-6, prior_precision=1.0, minibatch_size=100,
                             thin_interval=100, burn_in_iter_num=1000)
 
 
-def run_mnist_DistilledSGLD(training_num=50000):
+def run_mnist_DistilledSGLD(training_num=50000, gpu_id=None):
     X, Y, X_test, Y_test = load_mnist(training_num)
     minibatch_size = 100
     if training_num >= 10000:
@@ -214,10 +214,10 @@ def run_mnist_DistilledSGLD(training_num=50000):
     logsoftmax = LogSoftmax()
     student_net = get_mnist_sym(output_op=logsoftmax, num_hidden=num_hidden)
     data_shape = (minibatch_size,) + X.shape[1::]
-    teacher_data_inputs = {'data': nd.zeros(data_shape, ctx=dev()),
-                           'softmax_label': nd.zeros((minibatch_size,), ctx=dev())}
-    student_data_inputs = {'data': nd.zeros(data_shape, ctx=dev()),
-                           'softmax_label': nd.zeros((minibatch_size, 10), ctx=dev())}
+    teacher_data_inputs = {'data': nd.zeros(data_shape, ctx=dev(gpu_id)),
+                           'softmax_label': nd.zeros((minibatch_size,), ctx=dev(gpu_id))}
+    student_data_inputs = {'data': nd.zeros(data_shape, ctx=dev(gpu_id)),
+                           'softmax_label': nd.zeros((minibatch_size, 10), ctx=dev(gpu_id))}
     teacher_initializer = BiasXavier(factor_type="in", magnitude=1)
     student_initializer = BiasXavier(factor_type="in", magnitude=1)
     student_exe, student_params, _ = \
@@ -231,17 +231,17 @@ def run_mnist_DistilledSGLD(training_num=50000):
                       teacher_learning_rate=teacher_learning_rate,
                       student_learning_rate=student_learning_rate,
                       teacher_prior_precision=teacher_prior, student_prior_precision=student_prior,
-                      perturb_deviation=perturb_deviation, minibatch_size=100, dev=dev())
+                      perturb_deviation=perturb_deviation, minibatch_size=100, dev=dev(xpu))
 
 
-def run_toy_SGLD():
+def run_toy_SGLD(gpu_id=None):
     X, Y, X_test, Y_test = load_toy()
     minibatch_size = 1
     teacher_noise_precision = 1.0 / 9.0
     net = get_toy_sym(True, teacher_noise_precision)
     data_shape = (minibatch_size,) + X.shape[1::]
-    data_inputs = {'data': nd.zeros(data_shape, ctx=dev()),
-                   'teacher_output_label': nd.zeros((minibatch_size, 1), ctx=dev())}
+    data_inputs = {'data': nd.zeros(data_shape, ctx=dev(gpu_id)),
+                   'teacher_output_label': nd.zeros((minibatch_size, 1), ctx=dev(gpu_id))}
     initializer = mx.init.Uniform(0.07)
     exe, params, _ = \
         SGLD(sym=net, data_inputs=data_inputs,
@@ -253,20 +253,20 @@ def run_toy_SGLD():
              burn_in_iter_num=1000,
              thin_interval=10,
              task='regression',
-             minibatch_size=minibatch_size, dev=dev())
+             minibatch_size=minibatch_size, dev=dev(gpu_id))
 
 
-def run_toy_DistilledSGLD():
+def run_toy_DistilledSGLD(gpu_id=None):
     X, Y, X_test, Y_test = load_toy()
     minibatch_size = 1
     teacher_noise_precision = 1.0
     teacher_net = get_toy_sym(True, teacher_noise_precision)
     student_net = get_toy_sym(False)
     data_shape = (minibatch_size,) + X.shape[1::]
-    teacher_data_inputs = {'data': nd.zeros(data_shape, ctx=dev()),
-                           'teacher_output_label': nd.zeros((minibatch_size, 1), ctx=dev())}
-    student_data_inputs = {'data': nd.zeros(data_shape, ctx=dev())}
-    #                   'softmax_label': nd.zeros((minibatch_size, 10), ctx=dev())}
+    teacher_data_inputs = {'data': nd.zeros(data_shape, ctx=dev(gpu_id)),
+                           'teacher_output_label': nd.zeros((minibatch_size, 1), ctx=dev(gpu_id))}
+    student_data_inputs = {'data': nd.zeros(data_shape, ctx=dev(gpu_id))}
+    #                   'softmax_label': nd.zeros((minibatch_size, 10), ctx=dev(xpu))}
     teacher_initializer = mx.init.Uniform(0.07)
     student_initializer = mx.init.Uniform(0.07)
     student_grad_f = lambda student_outputs, teacher_pred: \
@@ -284,21 +284,21 @@ def run_toy_DistilledSGLD():
                       student_grad_f=student_grad_f,
                       teacher_prior_precision=0.1, student_prior_precision=0.001,
                       perturb_deviation=0.1, minibatch_size=minibatch_size, task='regression',
-                      dev=dev())
+                      dev=dev(gpu_id))
 
 
-def run_toy_HMC():
+def run_toy_HMC(gpu_id=None):
     X, Y, X_test, Y_test = load_toy()
     minibatch_size = Y.shape[0]
     noise_precision = 1 / 9.0
     net = get_toy_sym(True, noise_precision)
     data_shape = (minibatch_size,) + X.shape[1::]
-    data_inputs = {'data': nd.zeros(data_shape, ctx=dev()),
-                   'teacher_output_label': nd.zeros((minibatch_size, 1), ctx=dev())}
+    data_inputs = {'data': nd.zeros(data_shape, ctx=dev(gpu_id)),
+                   'teacher_output_label': nd.zeros((minibatch_size, 1), ctx=dev(gpu_id))}
     initializer = mx.init.Uniform(0.07)
     sample_pool = HMC(net, data_inputs=data_inputs, X=X, Y=Y, X_test=X_test, Y_test=Y_test,
                       sample_num=300000, initializer=initializer, prior_precision=1.0,
-                      learning_rate=1E-3, L=10, dev=dev())
+                      learning_rate=1E-3, L=10, dev=dev(gpu_id))
 
 
 def run_synthetic_SGLD():
@@ -350,21 +350,22 @@ def run_synthetic_SGLD():
                         help="Type of algorithm to use. 0 --> SGD, 1 --> SGLD, other-->DistilledSGLD")
     parser.add_argument("-t", "--training", type=int, default=50000,
                         help="Number of training samples")
+    parser.add_argument("--gpu", type=int, help="if None then use cpu else use gpu")
     args = parser.parse_args()
     training_num = args.training
     if args.dataset == 1:
         if 0 == args.algorithm:
-            run_mnist_SGD(training_num)
+            run_mnist_SGD(training_num, gpu_id=args.gpu)
         elif 1 == args.algorithm:
-            run_mnist_SGLD(training_num)
+            run_mnist_SGLD(training_num, gpu_id=args.gpu)
         else:
-            run_mnist_DistilledSGLD(training_num)
+            run_mnist_DistilledSGLD(training_num, gpu_id=args.gpu)
     elif args.dataset == 0:
         if 1 == args.algorithm:
-            run_toy_SGLD()
+            run_toy_SGLD(gpu_id=args.gpu)
         elif 2 == args.algorithm:
-            run_toy_DistilledSGLD()
+            run_toy_DistilledSGLD(gpu_id=args.gpu)
         elif 3 == args.algorithm:
-            run_toy_HMC()
+            run_toy_HMC(gpu_id=args.gpu)
     else:
         run_synthetic_SGLD()

example/fcn-xs/fcn_xs.py

@@ -28,9 +28,10 @@
 
 logger = logging.getLogger()
 logger.setLevel(logging.INFO)
-ctx = mx.gpu(0)
+
 
 def main():
+    ctx = mx.cpu() if not args.gpu else mx.gpu(args.gpu)
     fcnxs = symbol_fcnxs.get_fcn32s_symbol(numclass=21, workspace_default=1536)
     fcnxs_model_prefix = "model_pascal/FCN32s_VGG16"
     if args.model == "fcn16s":
@@ -85,6 +86,7 @@ def main():
         help='the init type of fcn-xs model, e.g. vgg16, fcnxs')
     parser.add_argument('--retrain', action='store_true', default=False,
         help='true means continue training.')
+    parser.add_argument("--gpu", type=int, help="gpu device_id. if not provide, then use cpu")
     args = parser.parse_args()
     logging.info(args)
     main()

example/rcnn/test.py

@@ -35,7 +35,7 @@ def test_net(sym, imdb, args):
     logger.info('called with args\n{}'.format(pprint.pformat(vars(args))))
 
     # setup context
-    ctx = mx.gpu(args.gpu)
+    ctx = mx.cpu() if not args.gpu else mx.gpu(args.gpu)
 
     # load testing data
     test_data = TestLoader(imdb.roidb, batch_size=1, short=args.img_short_side, max_size=args.img_long_side,
@@ -94,7 +94,7 @@ def parse_args():
     parser.add_argument('--params', type=str, default='', help='path to trained model')
     parser.add_argument('--dataset', type=str, default='voc', help='training dataset')
     parser.add_argument('--imageset', type=str, default='', help='imageset splits')
-    parser.add_argument('--gpu', type=int, default=0, help='gpu device eg. 0')
+    parser.add_argument('--gpu', type=int, default=0, help='gpu device_id eg: 0, if not provide then use cpu.')
     # faster rcnn params
     parser.add_argument('--img-short-side', type=int, default=600)
     parser.add_argument('--img-long-side', type=int, default=1000)

example/rcnn/train.py

@@ -33,7 +33,7 @@ def train_net(sym, roidb, args):
     logger.info('called with args\n{}'.format(pprint.pformat(vars(args))))
 
     # setup multi-gpu
-    ctx = [mx.gpu(int(i)) for i in args.gpus.split(',')]
+    ctx = [mx.cpu()] if not args.gpus else [mx.gpu(int(i)) for i in args.gpus.split(',')]
     batch_size = args.rcnn_batch_size * len(ctx)
 
     # load training data
@@ -127,7 +127,7 @@ def parse_args():
     parser.add_argument('--pretrained', type=str, default='', help='path to pretrained model')
     parser.add_argument('--dataset', type=str, default='voc', help='training dataset')
     parser.add_argument('--imageset', type=str, default='', help='imageset splits')
-    parser.add_argument('--gpus', type=str, default='0', help='gpu devices eg. 0,1')
+    parser.add_argument('--gpus', type=str, help='gpu devices eg: 0,1 , if not provide, then use cpu')
     parser.add_argument('--epochs', type=int, default=10, help='training epochs')
     parser.add_argument('--lr', type=float, default=0.001, help='base learning rate')
     parser.add_argument('--lr-decay-epoch', type=str, default='7', help='epoch to decay lr')