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)