help="set the size of the workspace for trt export")
arguments.add_argument("--trt-fp16",
action='store_true',
help="trt flag ---- export model in mixed precision mode")
arguments.add_argument("--capture-cuda-graph",
type=int,
default=1,
help="capture cuda graph for obtaining speedup. possible values: 0, 1. default: 1. ")
arguments.add_argument('--quantize',
action='store_true',
help='apply quantization for supported nodes')
arguments.add_argument('--calibrate',
action='store_true',
help='apply calibration for supported nodes')
# remainder args
arguments.add_argument('model_arguments',nargs=argparse.REMAINDER,help='arguments that will be ignored by deployer lib and will be forwarded to your deployer script')
#
args=parser.parse_args(argv)
deployer=Deployer(args)
#
returndeployer,args.model_arguments[1:]
classDeployerLibrary:
def__init__(self,args):
self.args=args
self.platform=None
defset_platform(self,platform):
''' sets the platform
:: platform :: "pytorch_libtorch" or "onnxruntime_onnx" or "tensorrt_plan"
'''
self.platform=platform
defbuild_trt_engine(self,model_file,shapes):
''' takes a path to an onnx file, and shape information, returns a trt engine
:: model_file :: path to an onnx model
:: shapes :: dictionary containing min shape, max shape, opt shape for the trt engine
#parser.add_argument('--dataset_size', type=int, default=0, help="Number of images or tensors to load. Default is 0 which means all samples")
#parser.add_argument('--data_preprocess', type=str, required=True, choices=['preprocess_method1', 'preprocess_method2', 'None'], help="Refer to Readme.md for guidance on choosing this option.")