pocketpose.converters.torch2onnx#

Module Contents#

Classes#

Torch2ONNXConverter

Converts PyTorch models to ONNX format.

class pocketpose.converters.torch2onnx.Torch2ONNXConverter(overwrite=True, log_level=logging.INFO, ops=17)#

Bases: pocketpose.converters.base_converter.BaseConverter

Converts PyTorch models to ONNX format.

_sanitize_names()#
_convert(model, save_path, *args, **kwargs)#

Converts a PyTorch model to ONNX format (.onnx) and saves it to disk.

Args:

model (torch.nn.Module): PyTorch model to convert (e.g. torchvision.models.resnet50()) save_path (str): Path to save the converted model to. This should be a .onnx file

Returns:

onnx.ModelProto: The converted ONNX model