def SaveModel(args, test_model): prefix = "gpu_0/" # print({prefix + str(b): str(b) for b in test_model.params}) # predictor_export_meta = pred_exp.PredictorExportMeta( # print(workspace.Blobs()) predictor_export_meta = ExtractPredictorNet( test_model.net.Proto(), # parameters=data_parallel_model.GetCheckpointParams(test_model), # parameters=[str(b) for b in test_model.params], input_blobs=[prefix + "data"], output_blobs=[prefix + "softmax"], renames={str(b): str(b)[6:] for b in test_model.params}) # save the test_model for the current epoch model_path = "%s/%s_pred.mdl" % ( args.file_store_path, args.save_model_name, ) # set db_type to be "minidb" instead of "log_file_db", which breaks # the serialization in save_to_db. Need to switch back to log_file_db # after migration print((predictor_export_meta)) pred_exp.save_to_db( db_type="minidb", db_destination=model_path, predictor_export_meta=predictor_export_meta, )
def test_extract_simple(self): from caffe2.python import brew from caffe2.python.model_helper import ModelHelper, ExtractPredictorNet model = ModelHelper(name="test", arg_scope={'order': 'NCHW'}) [data, label] = brew.image_input( model, "reader", ["xx/data", "label"], is_test=1, ) cnv = brew.conv(model, data, 'cnv', 32, 32, 4) a = brew.fc(model, cnv, 'a', 100, 200) pred = brew.fc(model, a, 'pred', 200, 5) brew.softmax(model, [pred, label], "softmax") (predict_net, export_blobs) = ExtractPredictorNet( net_proto=model.net.Proto(), input_blobs=["xx/data"], output_blobs=["pred"], renames={"xx/data": "image"}, ) export_blobs = set(export_blobs) ops = list(predict_net.Proto().op) for op in ops: self.assertFalse(op.type == "Softmax") self.assertFalse("xx/data" in op.input) # Note: image input should not be included self.assertEquals(ops[0].type, "Conv") self.assertEquals(ops[1].type, "FC") self.assertEquals(ops[2].type, "FC") self.assertEquals(len(ops), 3) # test rename happened self.assertEquals(ops[0].input[0], "image") # Check export blobs self.assertTrue("image" not in export_blobs) self.assertTrue("xx/data" not in export_blobs) self.assertEqual(set([str(p) for p in model.params]), export_blobs) # Check external inputs/outputs self.assertTrue("image" in predict_net.Proto().external_input) self.assertEquals(set(["pred"]), set(predict_net.Proto().external_output)) self.assertEqual( set(predict_net.Proto().external_input) - set([str(p) for p in model.params]), set(["image"]))
def test_lstm_extract_predictor_net(self): model = ModelHelper(name="lstm_extract_test") with core.DeviceScope(core.DeviceOption(caffe2_pb2.CPU, 0)): output, _, _, _ = rnn_cell.LSTM( model=model, input_blob="input", seq_lengths="seqlengths", initial_states=("hidden_init", "cell_init"), dim_in=20, dim_out=40, scope="test", drop_states=True, return_last_layer_only=True, ) # Run param init net to get the shapes for all inputs shapes = {} workspace.RunNetOnce(model.param_init_net) for b in workspace.Blobs(): shapes[b] = workspace.FetchBlob(b).shape # But export in CPU (predict_net, export_blobs) = ExtractPredictorNet( net_proto=model.net.Proto(), input_blobs=["input"], output_blobs=[output], device=core.DeviceOption(caffe2_pb2.CPU, 1), ) # Create the net and run once to see it is valid # Populate external inputs with correctly shaped random input # and also ensure that the export_blobs was constructed correctly. workspace.ResetWorkspace() shapes['input'] = [10, 4, 20] shapes['cell_init'] = [1, 4, 40] shapes['hidden_init'] = [1, 4, 40] print(predict_net.Proto().external_input) self.assertTrue('seqlengths' in predict_net.Proto().external_input) for einp in predict_net.Proto().external_input: if einp == 'seqlengths': workspace.FeedBlob( "seqlengths", np.array([10] * 4, dtype=np.int32) ) else: workspace.FeedBlob( einp, np.zeros(shapes[einp]).astype(np.float32), ) if einp != 'input': self.assertTrue(einp in export_blobs) print(str(predict_net.Proto())) self.assertTrue(workspace.CreateNet(predict_net.Proto())) self.assertTrue(workspace.RunNet(predict_net.Proto().name)) # Validate device options set correctly for the RNNs import google.protobuf.text_format as protobuftx for op in predict_net.Proto().op: if op.type == 'RecurrentNetwork': for arg in op.arg: if arg.name == "step_net": step_proto = caffe2_pb2.NetDef() protobuftx.Merge(arg.s.decode("ascii"), step_proto) for step_op in step_proto.op: self.assertEqual(0, step_op.device_option.device_type) self.assertEqual(1, step_op.device_option.cuda_gpu_id) elif arg.name == 'backward_step_net': self.assertEqual(b"", arg.s)