support appending net and converting them

Summary:
As per rushabhmshah99 request: he wants to append a pre-trained model (without training that) to the model.
So added data_parallel_model.ConvertNetForDevice() to enable that. The unit test shows example how to use this with
AppendNet, and I also added a blurb to the function.

Differential Revision: D5503335

fbshipit-source-id: b2a5db5c1739dc97f46dd0d7606ed555d99255b8

Author: Aapo Kyrola

caffe2/python/data_parallel_model.py

@@ -563,6 +563,44 @@ def Synchronize(model, timeout_sec=30):
     workspace.RunNetOnce(barrier_net)
 
 
+def ConvertNetForDevice(net, device=None):
+    '''
+    Converts all blobs in the net to have namescope gpu_X, and correct
+    device scope. You can use this to enable AppendNet with a
+    forward_pass_builder_fun:
+
+       def builder_fun(model):
+          ...
+          model.net.AppendNet(
+             data_parallel_model.ConvertNetForDevice(othermodel.net))
+          model.param_init_net.AppendNet(
+             data_parallel_model.ConvertNetForDevice(othermodel.param_init_net))
+    '''
+    mnet = copy.deepcopy(net)
+
+    if device is None:
+        device = scope.CurrentDeviceScope()
+
+    device_prefix = "gpu" if device.device_type == caffe2_pb2.CUDA else "cpu"
+
+    namescope = "{}_{}/".format(device_prefix, device.cuda_gpu_id)
+    for op in mnet.Proto().op:
+        if "RecurrentNetwork" in op.type:
+            raise("RecurrentNetwork conversion not yet supported")
+        for i, inputb in enumerate(op.input):
+            op.input[i] = namescope + inputb
+        for i, outputb in enumerate(op.output):
+            op.output[i] = namescope + outputb
+        for i, blob in enumerate(op.control_input):
+            op.control_input[i] = namescope + blob
+        op.device_option.CopyFrom(device)
+    for i, einp in enumerate(mnet.Proto().external_input):
+        mnet.Proto().external_input[i] = namescope + einp
+    for i, eoutp in enumerate(mnet.Proto().external_output):
+        mnet.Proto().external_output[i] = namescope + eoutp
+    return mnet
+
+
 def _ForEachGPU(gpu_ids, f, scoped=False, *args, **kwargs):
     for gpu_id in gpu_ids:
         device_opt = core.DeviceOption(caffe2_pb2.CUDA, gpu_id)

caffe2/python/data_parallel_model_test.py

@@ -9,7 +9,7 @@
 from multiprocessing import Process, Queue
 from caffe2.proto import caffe2_pb2
 from caffe2.python import core, cnn, data_parallel_model, dyndep, optimizer, \
-    rnn_cell, workspace
+    rnn_cell, workspace, model_helper, brew
 from caffe2.python.test_util import TestCase
 from future.utils import viewkeys
 
@@ -208,6 +208,59 @@ def add_optimizer(model):
         self.assertFalse(core.BlobReference("cpu_1/data") in checkpoint_params)
         self.assertTrue(core.BlobReference("optimizer_iteration") in checkpoint_params)
 
+    def test_net_conversion_and_append_net(self):
+        other = model_helper.ModelHelper()
+        fc1 = brew.fc(other, "data", "other_fc1", dim_in=3*227*227, dim_out=10)
+        fc2 = brew.fc(other, fc1, "other_fc2", dim_in=10, dim_out=10)
+        brew.fc(other, fc2, "other_fc3", dim_in=10, dim_out=10)
+
+        def add_input_ops(model):
+            model.net.UniformFill([], ["data"], shape=[4, 227, 227, 3])
+            model.net.UniformFill([], ["label"], shape=[4])
+
+        def add_model_ops(model, loss_scale):
+            model.NHWC2NCHW("data", "data_nchw")
+            model.Conv("data_nchw", 'conv1', 3, 64,
+                       weight_init=("MSRAFill", {}), kernel=7,
+                       stride=2, pad=3, no_bias=0)
+            model.SpatialBN('conv1', 'conv1_spatbn_relu', 64, epsilon=1e-3)
+            model.Relu('conv1_spatbn_relu', 'conv1_spatbn_relu')
+            model.MaxPool('conv1_spatbn_relu', 'pool1', kernel=3, stride=2)
+            model.FC('pool1', 'fc', dim_in=(64 * 56 * 56), dim_out=10)
+
+            # Append the net and param_init_net of the other model
+            appendnet = data_parallel_model.ConvertNetForDevice(other.net)
+            model.net.AppendNet(appendnet)
+
+            model.param_init_net.AppendNet(
+                data_parallel_model.ConvertNetForDevice(other.param_init_net))
+
+            model.Sigmoid('fc', 'fc_sigm')
+            model.Softmax('fc_sigm', 'softmax')
+            loss = model.AveragedLoss('softmax', 'loss')
+            return [loss]
+
+        def add_optimizer(model):
+            optimizer.build_sgd(model, 0.1, policy="fixed", momentum=0.9)
+
+        model = cnn.CNNModelHelper(
+            order="NCHW",
+            name="test",
+        )
+        data_parallel_model.Parallelize_CPU(
+            model,
+            input_builder_fun=add_input_ops,
+            forward_pass_builder_fun=add_model_ops,
+            optimizer_builder_fun=add_optimizer,
+            devices=range(4)
+        )
+
+        # Just create and run net and confirm no exception is thrown
+        workspace.RunNetOnce(model.param_init_net)
+        workspace.CreateNet(model.net)
+        workspace.RunNet(model.net)
+
+
     def test_synchronization_barrier(self):
 
         def run(comm_rank, comm_size, tmpdir):