def export_to_caffe2(self,
model,
export_path: str,
export_onnx_path: str = None) -> List[str]:
"""
export pytorch model to caffe2 by first using ONNX to convert logic in forward
function to a caffe2 net, and then prepend/append additional operators to
the caffe2 net according to the model
Args:
model (Model): pytorch model to export
export_path (str): path to save the exported caffe2 model
export_onnx_path (str): path to save the exported onnx model
Returns:
final_output_names: list of caffe2 model output names
"""
print(f"Saving caffe2 model to: {export_path}")
# caffe2/onnx doesn't support internal uri(i.e. manifold)
# workaround: save to a temp file and copy to model_path
# this will be deprecated soon after caffe2 fully deprecated
_, temp_path = tempfile.mkstemp(prefix="pytext")
c2_prepared = onnx.pytorch_to_caffe2(
model,
self.dummy_model_input,
self.input_names,
self.output_names,
temp_path,
export_onnx_path,
)
c2_prepared, final_input_names = self.prepend_operators(
c2_prepared, self.input_names)
# Required because of https://github.com/pytorch/pytorch/pull/6456/files
with c2_prepared.workspace._ctx:
predict_net = core.Net(c2_prepared.predict_net)
init_net = core.Net(c2_prepared.init_net)
net_outputs, final_out_names = self.postprocess_output(
init_net, predict_net, c2_prepared.workspace,
self.output_names, model)
for output in net_outputs:
predict_net.AddExternalOutput(output)
c2_prepared.predict_net = predict_net.Proto()
c2_prepared.init_net = init_net.Proto()
# Save predictor net to file
onnx.export_nets_to_predictor_file(
c2_prepared,
final_input_names,
final_out_names,
temp_path,
self.get_extra_params(),
)
PathManager.copy_from_local(temp_path, export_path, overwrite=True)
return final_out_names
def export_to_caffe2(self,
model,
export_path: str,
export_onnx_path: str = None) -> List[str]:
"""
export pytorch model to caffe2 by first using ONNX to convert logic in forward
function to a caffe2 net, and then prepend/append additional operators to
the caffe2 net according to the model
Args:
model (Model): pytorch model to export
export_path (str): path to save the exported caffe2 model
export_onnx_path (str): path to save the exported onnx model
Returns:
final_output_names: list of caffe2 model output names
"""
print(f"Saving caffe2 model to: {export_path}")
c2_prepared = onnx.pytorch_to_caffe2(
model,
self.dummy_model_input,
self.input_names,
self.output_names,
export_path,
export_onnx_path,
)
c2_prepared, final_input_names = self.prepend_operators(
c2_prepared, self.input_names)
# Required because of https://github.com/pytorch/pytorch/pull/6456/files
with c2_prepared.workspace._ctx:
predict_net = core.Net(c2_prepared.predict_net)
init_net = core.Net(c2_prepared.init_net)
net_outputs, final_out_names = self.postprocess_output(
init_net, predict_net, c2_prepared.workspace,
self.output_names, model)
for output in net_outputs:
predict_net.AddExternalOutput(output)
c2_prepared.predict_net = predict_net.Proto()
c2_prepared.init_net = init_net.Proto()
# Save predictor net to file
onnx.export_nets_to_predictor_file(
c2_prepared,
final_input_names,
final_out_names,
export_path,
self.get_extra_params(),
)
return final_out_names