def _sum_deterministic_policy(self, model_names, path):
net = core.Net('DeterministicPolicy')
C2.set_net(net)
output = 'ActionProbabilities'
workspace.FeedBlob(output, np.array([1.0]))
model_outputs = []
for model in model_names:
model_output = '{}_Output'.format(model)
workspace.FeedBlob(model_output, np.array([1.0], dtype=np.float32))
model_outputs.append(model_output)
max_action = C2.FlattenToVec(
C2.ArgMax(C2.Transpose(C2.Sum(*model_outputs)))
)
one_blob = C2.NextBlob('one')
workspace.FeedBlob(one_blob, np.array([1.0], dtype=np.float32))
C2.net().SparseToDense(
[
max_action,
one_blob,
model_outputs[0],
],
[output],
)
meta = PredictorExportMeta(
net,
[one_blob],
model_outputs,
[output],
)
save_to_db('minidb', path, meta)
def get_predictor_export_meta(self):
"""
Returns a PredictorExportMeta object
"""
return PredictorExportMeta(
self._net,
self._parameters,
self._input_blobs,
self._output_blobs,
extra_init_net=self._init_net,
)
def get_predictor_export_meta(self):
"""
Returns a PredictorExportMeta object
"""
return PredictorExportMeta(
self._net,
self._parameters,
self._input_blobs,
self._output_blobs,
extra_init_net=self._init_net,
net_type="async_scheduling",
num_workers=8,
)
def _dummy_model_copy(self, model_name, path):
net = core.Net(model_name)
C2.set_net(net)
inp = 'Input'
output = 'Output'
workspace.FeedBlob(inp, np.array([1.0]))
workspace.FeedBlob(output, np.array([1.0]))
net.Copy([inp], [output])
meta = PredictorExportMeta(
net,
[],
[inp],
[output],
)
save_to_db('minidb', path, meta)
def save_sum_deterministic_policy(model_names, path, db_type):
net = core.Net("DeterministicPolicy")
C2.set_net(net)
output = "ActionProbabilities"
workspace.FeedBlob(output, np.array([1.0]))
model_outputs = []
for model in model_names:
model_output = "{}_Output".format(model)
workspace.FeedBlob(model_output, np.array([[1.0]], dtype=np.float32))
model_outputs.append(model_output)
max_action = C2.FlattenToVec(C2.ArgMax(C2.Transpose(C2.Sum(*model_outputs))))
one_blob = C2.NextBlob("one")
workspace.FeedBlob(one_blob, np.array([1.0], dtype=np.float32))
C2.net().SparseToDense([max_action, one_blob, model_outputs[0]], [output])
meta = PredictorExportMeta(net, [one_blob], model_outputs, [output])
save_to_db(db_type, path, meta)
def get_predictor_export_meta(self):
return PredictorExportMeta(self._net, self._parameters,
self._input_blobs, self._output_blobs)
def get_predictor_export_meta_and_workspace(self,
feature_extractor=None,
output_transformer=None):
"""
ONNX would load blobs into a private workspace. We returns the workspace
here instead of copying the blobs to the global workspace in order to
save memory in the export state. Returning private workspace, we only
need memory for PyTorch model, ONNX buffer, and Caffe2 model. Including
optimizer parameters, this means we can train and save a model
a quarter of the size of machine memory.
We should revisit this once PyTorch 1.0 is ready.
Args:
feature_extractor: An instance of FeatureExtractorBase
"""
# 1. Get Caffe2 model
c2_model, input_blobs, output_blobs = self.get_caffe2_model()
ws = c2_model.workspace
# Initializing constants in the model
init_net = core.Net(c2_model.init_net)
ws.CreateNet(init_net)
ws.RunNet(init_net)
# Per ONNX code comment, input blobs are not initilized
model_inputs = c2_model.uninitialized
assert len(model_inputs) > 0, "Model is expected to have some input"
parameters = [b for b in ws.Blobs() if b not in model_inputs]
# Input blobs in order
model_input_blobs = [b for b in input_blobs if b in model_inputs]
predict_net = core.Net("predict_net")
output_blob_names = self.output_blob_names()
assert len(output_blobs) == len(output_blob_names), (
"output_blobs and output_blob_names must have the same lengths. "
"Check that your model don't reuse output tensors. "
"output_blobs: {}; output_blob_names: {}".format(
output_blobs, output_blob_names))
blob_remap = {
onnx_name: explicit_name
for onnx_name, explicit_name in zip(output_blobs,
output_blob_names)
}
shapes = {}
# 2. Create feature extractor net
if feature_extractor:
feature_extractor_nets = feature_extractor.create_net()
# Initializing feature extractor parameters
ws.CreateNet(feature_extractor_nets.init_net)
ws.RunNet(feature_extractor_nets.init_net)
feature_extractor_params = set(
feature_extractor_nets.init_net.Proto().external_output)
assert (len(set(parameters) & feature_extractor_params) == 0
), "Blob names collide! Please open a bug report"
parameters += feature_extractor_params
extracted_blobs = [
str(b) for b in
feature_extractor_nets.net.output_record().field_blobs()
]
assert len(model_input_blobs) == len(extracted_blobs), (
"The lengths of model_input_blobs and extracted_blobs must match. "
"model_input_blobs: {}; extracted_blobs: {}".format(
model_input_blobs, extracted_blobs))
blob_remap.update({
onnx_name: extracted_name
for onnx_name, extracted_name in zip(model_input_blobs,
extracted_blobs)
})
predict_net.AppendNet(feature_extractor_nets.net)
del predict_net.Proto().external_output[:]
input_blobs = [
b for b in predict_net.Proto().external_input
if b not in feature_extractor_params
]
shapes.update({b: [] for b in input_blobs})
else:
input_blobs = model_input_blobs
# 3. Rename the input blobs of model to match the output of feature
# extractor net
model_net = core.Net(c2_model.predict_net).Clone("remapped_model_net",
blob_remap=blob_remap)
# 5. Join feature extractor net & model net
predict_net.AppendNet(model_net)
if output_transformer is not None:
output_field_names = self.output_field_names()
original_output = schema.from_column_list(
col_names=output_field_names,
col_blobs=[core.BlobReference(b) for b in output_blob_names],
)
output_transformer_nets = output_transformer.create_net(
original_output)
# Initializing output transformer parameters
ws.CreateNet(output_transformer_nets.init_net)
ws.RunNet(output_transformer_nets.init_net)
output_transformer_params = set(
output_transformer_nets.init_net.Proto().external_output)
assert (len(set(parameters) & output_transformer_params) == 0
), "Blob names collide! Please open a bug report"
parameters += output_transformer_params
del predict_net.Proto().external_output[:]
predict_net.AppendNet(output_transformer_nets.net)
# These shapes are not really used but required, so just pass fake ones
shapes.update({b: [] for b in predict_net.Proto().external_output})
return (
PredictorExportMeta(
predict_net,
parameters,
input_blobs,
predict_net.Proto().external_output,
shapes=shapes,
net_type="async_scheduling",
num_workers=8,
),
ws,
)
