torchexpo.modules

class torchexpo.modules.ImageClassificationModule(model, model_name, model_example)

Image Classification module for all vision models under vision/image_classification

Parameters

• model – Pretrained model which is used for conversion

• model_name – Name of the model

• model_example – Example tensor used as example input while extraction

extract_onnx(opset_version=None)

Extract onnx from image classification model

Parameters

opset_version – Opset version used while ONNX conversion

extract_torchscript()

Extract torchscript module from image classification model

class torchexpo.modules.ImageSegmentationModule(model, model_name, model_example)

Image Segmentation module for all vision models under vision/image_segmentation

Parameters

• model – Pretrained model which is used for conversion

• model_name – Name of the model

• model_example – Example tensor used as example input while extraction

extract_onnx(opset_version=11)

Extract onnx from image segmentation model

Parameters

opset_version – Opset version used while ONNX conversion

extract_torchscript()

Extract torchscript module from image segmentation model

class torchexpo.modules.SentimentAnalysisModule(model, model_name, model_example)

Sentiment Analysis module for all nlp models under nlp/sentiment_analysis

Parameters

• model – Pretrained model which is used for conversion

• model_name – Name of the model

• model_example – Example tensor used as example input while extraction

extract_onnx(opset_version=None)

Extract onnx from sentiment analysis model

Parameters

opset_version – Opset version used while ONNX conversion

extract_torchscript()

Extract torchscript module from sentiment analysis model

class torchexpo.modules.TorchExpoModule(model, model_name, model_example)

Base class for all PyTorch models

Parameters

• model – Pretrained model which is used for conversion

• model_name – Name of the model

• model_example – Example tensor used as example input while extraction

extract_caffe2_mobile()

Extracts model in Caffe2 Mobile format

extract_onnx(opset_version=None)

Extracts model in ONNX format

Parameters

opset_version – Opset version used while ONNX conversion

extract_torchscript()

Extracts model in TorchScript format

get_extracted_file_name()

Returns file name for output

print_message(output_type)

Prints message

Parameters

output_type – Name of the output format used for printing