Source code for pytext.exporters.exporter

#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved

from typing import 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_utils


[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
[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)
[docs] @classmethod def get_feature_metadata( cls, feature_config: FeatureConfig, feature_meta: Dict[str, FieldMeta] ): # 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 isinstance(feat_config, ConfigBase): 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) dummy_model_input.append( torch.tensor([1, 1], dtype=torch.long) ) # token lengths input_names.append("tokens_lens") return input_names, tuple(dummy_model_input), feature_itos_map
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 # 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" )
[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_utils.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 get_extra_params(self) -> List[str]: """ Returns: list of blobs to be added as extra params to the caffe2 model """ return []
[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 """ c2_prepared = onnx_utils.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_utils.export_nets_to_predictor_file( c2_prepared, final_input_names, final_out_names, export_path, self.get_extra_params(), ) return final_out_names