pytext.exporters package

Submodules

pytext.exporters.exporter module

class pytext.exporters.exporter.ModelExporter(config, input_names, dummy_model_input, vocab_map, output_names)[source]

Bases: pytext.config.component.Component

Model exporter exports a PyTorch model to Caffe2 model using ONNX

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

Config[source]

alias of ModelExporter.Config

export_to_caffe2(model, export_path: str, export_onnx_path: str = None) → List[str][source]

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

Parameters:
  • 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:

list of caffe2 model output names

Return type:

final_output_names

classmethod from_config(config, feature_config: pytext.config.field_config.FeatureConfig, target_config: Union[pytext.config.pytext_config.ConfigBase, List[pytext.config.pytext_config.ConfigBase]], meta: pytext.data.data_handler.CommonMetadata, *args, **kwargs)[source]

Gather all the necessary metadata from configs and global metadata to be used in exporter

get_extra_params() → List[str][source]
Returns:list of blobs to be added as extra params to the caffe2 model
classmethod get_feature_metadata(feature_config: pytext.config.field_config.FeatureConfig, feature_meta: Dict[str, pytext.fields.field.FieldMeta])[source]
postprocess_output(init_net: caffe2.python.core.Net, predict_net: caffe2.python.core.Net, workspace: <module 'caffe2.python.workspace' from '/home/docs/checkouts/readthedocs.org/user_builds/pytext-pytext/envs/latest/lib/python3.6/site-packages/caffe2/python/workspace.py'>, output_names: List[str], py_model)[source]

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

Parameters:
  • 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:

list of blobs that will be added to the caffe2 model final_output_names: list of output names of the blobs to add

Return type:

result

prepend_operators(c2_prepared: caffe2.python.onnx.backend_rep.Caffe2Rep, input_names: List[str]) → Tuple[caffe2.python.onnx.backend_rep.Caffe2Rep, List[str]][source]

Prepend operators to the converted caffe2 net, do nothing by default

Parameters:
  • c2_prepared (Caffe2Rep) – caffe2 net rep
  • input_names (List[str]) – current input names to the caffe2 net
Returns:

caffe2 net with prepended operators input_names (List[str]): list of input names for the new net

Return type:

c2_prepared (Caffe2Rep)

Module contents

class pytext.exporters.ModelExporter(config, input_names, dummy_model_input, vocab_map, output_names)[source]

Bases: pytext.config.component.Component

Model exporter exports a PyTorch model to Caffe2 model using ONNX

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

Config[source]

alias of ModelExporter.Config

export_to_caffe2(model, export_path: str, export_onnx_path: str = None) → List[str][source]

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

Parameters:
  • 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:

list of caffe2 model output names

Return type:

final_output_names

classmethod from_config(config, feature_config: pytext.config.field_config.FeatureConfig, target_config: Union[pytext.config.pytext_config.ConfigBase, List[pytext.config.pytext_config.ConfigBase]], meta: pytext.data.data_handler.CommonMetadata, *args, **kwargs)[source]

Gather all the necessary metadata from configs and global metadata to be used in exporter

get_extra_params() → List[str][source]
Returns:list of blobs to be added as extra params to the caffe2 model
classmethod get_feature_metadata(feature_config: pytext.config.field_config.FeatureConfig, feature_meta: Dict[str, pytext.fields.field.FieldMeta])[source]
postprocess_output(init_net: caffe2.python.core.Net, predict_net: caffe2.python.core.Net, workspace: <module 'caffe2.python.workspace' from '/home/docs/checkouts/readthedocs.org/user_builds/pytext-pytext/envs/latest/lib/python3.6/site-packages/caffe2/python/workspace.py'>, output_names: List[str], py_model)[source]

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

Parameters:
  • 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:

list of blobs that will be added to the caffe2 model final_output_names: list of output names of the blobs to add

Return type:

result

prepend_operators(c2_prepared: caffe2.python.onnx.backend_rep.Caffe2Rep, input_names: List[str]) → Tuple[caffe2.python.onnx.backend_rep.Caffe2Rep, List[str]][source]

Prepend operators to the converted caffe2 net, do nothing by default

Parameters:
  • c2_prepared (Caffe2Rep) – caffe2 net rep
  • input_names (List[str]) – current input names to the caffe2 net
Returns:

caffe2 net with prepended operators input_names (List[str]): list of input names for the new net

Return type:

c2_prepared (Caffe2Rep)