#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
import tempfile
from typing import Callable, Dict, List, Tuple, Union
import torch
from caffe2.python import core
from caffe2.python.onnx.backend_rep import Caffe2Rep
from pytext.config import ConfigBase
from pytext.config.component import Component, ComponentType
from pytext.config.field_config import FeatureConfig
from pytext.data import CommonMetadata
from pytext.fields import FieldMeta
from pytext.utils import onnx
from pytext.utils.file_io import PathManager
from pytext.utils.usage import log_class_usage
[docs]class ModelExporter(Component):
"""
Model exporter exports a PyTorch model to Caffe2 model using ONNX
Attributes:
input_names (List[Str]): names of the input variables to model forward
function, in a flattened way.
e.g: forward(tokens, dict) where tokens is List[Tensor] and dict is
a tuple of value and length: (List[Tensor], List[Tensor]) the
input names should looks like ['token', 'dict_value', 'dict_length']
dummy_model_input (Tuple[torch.Tensor]): dummy values to define the
shape of input tensors, should exactly match the shape of the model
forward function
vocab_map (Dict[str, List[str]]): dict of input feature names
to corresponding index_to_string array, e.g:
::
{
"text": ["<UNK>", "W1", "W2", "W3", "W4", "W5", "W6", "W7", "W8"],
"dict": ["<UNK>", "D1", "D2", "D3", "D4", "D5", "D6", "D7", "D8"]
}
output_names (List[Str]): names of output variables
"""
__COMPONENT_TYPE__ = ComponentType.EXPORTER
[docs] class Config(ConfigBase):
export_logits: bool = False
export_raw_to_metrics: bool = False
[docs] @classmethod
def from_config(
cls,
config,
feature_config: FeatureConfig,
target_config: Union[ConfigBase, List[ConfigBase]],
meta: CommonMetadata,
*args,
**kwargs,
):
"""
Gather all the necessary metadata from configs and global metadata to be
used in exporter
"""
input_names, dummy_model_input, vocab_map = cls.get_feature_metadata(
feature_config, meta.features
)
if not isinstance(target_config, list):
target_config = [target_config]
output_names = [
name for target in target_config for name in target.export_output_names
]
return cls(config, input_names, dummy_model_input, vocab_map, output_names)
def __init__(self, config, input_names, dummy_model_input, vocab_map, output_names):
super().__init__(config)
self.input_names = input_names
self.output_names = output_names
self.dummy_model_input = dummy_model_input
self.vocab_map = vocab_map or {}
# validate feature vocab
for name in self.vocab_map:
if name not in self.input_names:
raise ValueError(
f"{name} is not found in input names {self.input_names}, \
there's a mismatch"
)
log_class_usage(__class__)
[docs] def prepend_operators(
self, c2_prepared: Caffe2Rep, input_names: List[str]
) -> Tuple[Caffe2Rep, List[str]]:
"""
Prepend operators to the converted caffe2 net, do nothing by default
Args:
c2_prepared (Caffe2Rep): caffe2 net rep
input_names (List[str]): current input names to the caffe2 net
Returns:
c2_prepared (Caffe2Rep): caffe2 net with prepended operators
input_names (List[str]): list of input names for the new net
"""
return onnx.add_feats_numericalize_ops(c2_prepared, self.vocab_map, input_names)
[docs] def postprocess_output(
self,
init_net: core.Net,
predict_net: core.Net,
workspace: core.workspace,
output_names: List[str],
py_model,
):
"""
Postprocess the model output, generate additional blobs for human readable
prediction. By default it use export function of output layer from pytorch
model to append additional operators to caffe2 net
Args:
init_net (caffe2.python.Net): caffe2 init net created by the current graph
predict_net (caffe2.python.Net): caffe2 net created by the current graph
workspace (caffe2.python.workspace): caffe2 current workspace
output_names (List[str]): current output names of the caffe2 net
py_model (Model): original pytorch model object
Returns:
result: list of blobs that will be added to the caffe2 model
final_output_names: list of output names of the blobs to add
"""
model_out = py_model(*self.dummy_model_input)
res = py_model.output_layer.export_to_caffe2(
workspace, init_net, predict_net, model_out, *output_names
)
# optionally include the last decoder layer of pytorch model
final_output_names = [str(output) for output in res] + (
output_names if self.config.export_logits else []
)
return res, final_output_names
[docs] 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
[docs] def export_to_metrics(self, model, metric_channels):
"""
Exports the pytorch model to tensorboard as a graph.
Args:
model (Model): pytorch model to export
metric_channels (List[Channel]): outputs of model's execution graph
"""
for mc in metric_channels or []:
if self.config.export_raw_to_metrics:
mc.export(model, self.dummy_model_input, operator_export_type="RAW")
else:
mc.export(model, self.dummy_model_input)
@classmethod
def _get_exportable_metadata(
cls,
exportable_filter: Callable,
feature_config: FeatureConfig,
feature_meta: Dict[str, FieldMeta],
) -> Tuple[List[str], List, Dict]:
# The number of names in input_names *must* be equal to the number of
# tensors passed in dummy_input
input_names: List[str] = []
dummy_model_input: List = []
feature_itos_map = {}
for name, feat_config in feature_config._asdict().items():
if exportable_filter(feat_config):
input_names.extend(feat_config.export_input_names)
if getattr(feature_meta[name], "vocab", None):
feature_itos_map[feat_config.export_input_names[0]] = feature_meta[
name
].vocab.itos
dummy_model_input.append(feature_meta[name].dummy_model_input)
return input_names, dummy_model_input, feature_itos_map
@classmethod
def _add_feature_lengths(cls, input_names: List[str], dummy_model_input: List):
"""If any of the input_names have tokens or seq_tokens, add the length
of those tokens to dummy_input
"""
if "tokens_vals" in input_names:
dummy_model_input.append(
torch.tensor([1, 1], dtype=torch.long)
) # token lengths
input_names.append("tokens_lens")
if "seq_tokens_vals" in input_names:
dummy_model_input.append(
torch.tensor([1, 1], dtype=torch.long)
) # seq lengths
input_names.append("seq_tokens_lens")