Convert a Pytorch Module to PPQ graph, this conversation exports pytorch model to onnx first, then PPQ will read from onnx file to build graph. Args: model (torch.nn.Module): Converting Model. onnx_file_name (str, optional): Exporting Onnx File name, b
(
model: torch.nn.Module, onnx_file_name: str = 'MyModel.onnx',
sample: Union[dict, list, torch.Tensor] = None, device = 'cuda',
verbose: bool = False, opset: int = 11, do_constant_folding: bool = None,
dynamic_axes: dict = None, training: bool = False)
| 76 | return load_graph(import_file, from_framework=NetworkFramework.NATIVE) |
| 77 | |
| 78 | def load_torch_model( |
| 79 | model: torch.nn.Module, onnx_file_name: str = 'MyModel.onnx', |
| 80 | sample: Union[dict, list, torch.Tensor] = None, device = 'cuda', |
| 81 | verbose: bool = False, opset: int = 11, do_constant_folding: bool = None, |
| 82 | dynamic_axes: dict = None, training: bool = False) -> BaseGraph: |
| 83 | """ |
| 84 | Convert a Pytorch Module to PPQ graph, |
| 85 | this conversation exports pytorch model to onnx first, |
| 86 | then PPQ will read from onnx file to build graph. |
| 87 | |
| 88 | Args: |
| 89 | model (torch.nn.Module): Converting Model. |
| 90 | |
| 91 | onnx_file_name (str, optional): Exporting Onnx File name, by default it is 'MyModel.onnx'. |
| 92 | |
| 93 | sample (Union[dict, list, torch.Tensor], optional): |
| 94 | A sample input for tracing your pytorch model. |
| 95 | |
| 96 | Pytorch will feed this sample input to your model, |
| 97 | recording record every layer and functions during executing. |
| 98 | |
| 99 | Each function recorded will be translated into onnx. |
| 100 | The input and output shape will also depends on this sample input. |
| 101 | |
| 102 | device (str, optional): _description_. Defaults to 'cuda'. |
| 103 | |
| 104 | verbose (bool, optional): _description_. Defaults to True. |
| 105 | |
| 106 | opset (int, optional): _description_. Defaults to 11. |
| 107 | |
| 108 | do_constant_folding (bool, optional): |
| 109 | Apply the constant-folding optimization. |
| 110 | Constant-folding will replace some of the ops that have all constant inputs |
| 111 | with pre-computed constant nodes. |
| 112 | |
| 113 | dynamic_axes (dict, optional): |
| 114 | (dict<string, dict<int, string>> or dict<string, list(int)>, default empty dict): |
| 115 | By default the exported model will have the shapes of all input and output tensors |
| 116 | set to exactly match those given in ``args``. To specify axes of tensors as |
| 117 | dynamic (i.e. known only at run-time), set ``dynamic_axes`` to a dict with schema: |
| 118 | |
| 119 | Returns: |
| 120 | BaseGraph: Converted PPQ Graph Instance. |
| 121 | """ |
| 122 | |
| 123 | if opset not in {11, 13}: |
| 124 | ppq_warning(f'The Opset of Exporting Model should be 11 or 13, however {opset} was given.') |
| 125 | |
| 126 | model = model.eval() |
| 127 | model = model.to(device) |
| 128 | |
| 129 | _training = torch.onnx.TrainingMode.TRAINING if training else torch.onnx.TrainingMode.EVAL |
| 130 | torch.onnx.export( |
| 131 | model=model, args=sample, |
| 132 | verbose=verbose, f=onnx_file_name, opset_version=opset, |
| 133 | do_constant_folding=do_constant_folding, |
| 134 | dynamic_axes=dynamic_axes, |
| 135 | training=_training) |
no test coverage detected