| 191 | |
| 192 | |
| 193 | def parse_args(): |
| 194 | parser = argparse.ArgumentParser( |
| 195 | description='Convert MMSegmentation models from ONNX to TensorRT') |
| 196 | parser.add_argument('config', help='Config file of the model') |
| 197 | parser.add_argument('model', help='Path to the input ONNX model') |
| 198 | parser.add_argument( |
| 199 | '--trt-file', type=str, help='Path to the output TensorRT engine') |
| 200 | parser.add_argument( |
| 201 | '--max-shape', |
| 202 | type=int, |
| 203 | nargs=4, |
| 204 | default=[1, 3, 400, 600], |
| 205 | help='Maximum shape of model input.') |
| 206 | parser.add_argument( |
| 207 | '--min-shape', |
| 208 | type=int, |
| 209 | nargs=4, |
| 210 | default=[1, 3, 400, 600], |
| 211 | help='Minimum shape of model input.') |
| 212 | parser.add_argument('--fp16', action='store_true', help='Enable fp16 mode') |
| 213 | parser.add_argument( |
| 214 | '--workspace-size', |
| 215 | type=int, |
| 216 | default=1, |
| 217 | help='Max workspace size in GiB') |
| 218 | parser.add_argument( |
| 219 | '--input-img', type=str, default='', help='Image for test') |
| 220 | parser.add_argument( |
| 221 | '--show', action='store_true', help='Whether to show output results') |
| 222 | parser.add_argument( |
| 223 | '--dataset', |
| 224 | type=str, |
| 225 | default='CityscapesDataset', |
| 226 | help='Dataset name') |
| 227 | parser.add_argument( |
| 228 | '--verify', |
| 229 | action='store_true', |
| 230 | help='Verify the outputs of ONNXRuntime and TensorRT') |
| 231 | parser.add_argument( |
| 232 | '--verbose', |
| 233 | action='store_true', |
| 234 | help='Whether to verbose logging messages while creating \ |
| 235 | TensorRT engine.') |
| 236 | args = parser.parse_args() |
| 237 | return args |
| 238 | |
| 239 | |
| 240 | if __name__ == '__main__': |