pocketpose.converters.torch2onnx
#
Module Contents#
Classes#
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