def test_parse_shape_with_dim_param_quoted(self, name, quote): meta_args = [ "{name}:[{quote}batch{quote},3,224,224]".format(name=name, quote=quote) ] meta = args_util.parse_meta(meta_args, includes_dtype=False) assert meta[name].shape == ("batch", 3, 224, 224)
def parse(self, args): def determine_model_type(): if args_util.get(args, "model_type") is not None: return args.model_type.lower() if args_util.get(args, "model_file") is None: return None def use_ext(ext_mapping): file_ext = os.path.splitext(args.model_file)[-1] if file_ext in ext_mapping: return ext_mapping[file_ext] runners = util.default(args_util.get(args, "runners"), []) if args_util.get(args, "ckpt") or os.path.isdir(args.model_file): return "ckpt" elif "tf" in runners or "trt_legacy" in runners: if args.caffe_model: return "caffe" return use_ext(ModelArgs.EXT_MODEL_TYPE_MAPPING) or "frozen" else: model_type = use_ext(ModelArgs.EXT_MODEL_TYPE_MAPPING) if model_type: return model_type G_LOGGER.exit( "Could not automatically determine model type for: {:}\n" "Please explicitly specify the type with the --model-type option" .format(args.model_file)) if args_util.get(args, "input_shapes"): self.input_shapes = args_util.parse_meta( args_util.get(args, "input_shapes"), includes_dtype=False) # TensorMetadata else: self.input_shapes = TensorMetadata() self.model_file = args_util.get(args, "model_file") if self.model_file: G_LOGGER.verbose("Model: {:}".format(self.model_file)) if not os.path.exists(self.model_file): G_LOGGER.warning("Model path does not exist: {:}".format( self.model_file)) self.model_file = os.path.abspath(self.model_file) model_type_str = util.default(self._model_type, determine_model_type()) self.model_type = ModelArgs.ModelType( model_type_str) if model_type_str else None if self.model_type == "trt-network-script" and ( not self.model_file or not self.model_file.endswith(".py")): G_LOGGER.exit( "TensorRT network scripts must exist and have '.py' extensions. " "Note: Provided network script path was: {:}".format( self.model_file))
def get_shapes(lst, idx): nonlocal default_shapes default_shapes = copy.copy(default_shapes) if idx < len(lst): default_shapes.update( args_util.parse_meta(lst[idx], includes_dtype=False)) # Don't care about dtype, and need to override dynamic dimensions shapes = { name: util.override_dynamic_shape(shape) for name, (_, shape) in default_shapes.items() } for name, shape in shapes.items(): if tuple(default_shapes[name].shape) != tuple(shape): G_LOGGER.warning( "Input tensor: {:} | For TensorRT profile, overriding dynamic shape: {:} to: {:}" .format(name, default_shapes[name].shape, shape), mode=LogMode.ONCE, ) return shapes
def test_parse_shape_dtype_auto(self, name): meta_args = ["{name}:auto:auto".format(name=name)] meta = args_util.parse_meta(meta_args) assert meta[name].shape is None assert meta[name].dtype is None
def test_parse_shape_dtype(self, name): meta_args = ["{name}:[1,3,224,224]:float32".format(name=name)] meta = args_util.parse_meta(meta_args) assert meta[name].shape == (1, 3, 224, 224) assert meta[name].dtype == np.float32
def test_parse_dtype_only(self, name): meta_args = ["{name}:float32".format(name=name)] meta = args_util.parse_meta(meta_args, includes_shape=False) assert meta[name].shape is None assert meta[name].dtype == np.float32
def test_parse_shape_single_dim(self, name): meta_args = ["{name}:[1]".format(name=name)] meta = args_util.parse_meta(meta_args, includes_dtype=False) assert meta[name].shape == (1, )
def test_parse_shape_scalar(self, name): meta_args = ["{name}:[]".format(name=name)] meta = args_util.parse_meta(meta_args, includes_dtype=False) assert meta[name].shape == tuple()
def test_parse_empty_shape(self, name): meta_args = ["{name}:[0,3,0,224]".format(name=name)] meta = args_util.parse_meta(meta_args, includes_dtype=False) assert meta[name].shape == (0, 3, 0, 224) assert meta[name].dtype is None
def test_parse_shape_only(self, name): meta_args = ["{name}:[1,3,224,224]".format(name=name)] meta = args_util.parse_meta(meta_args, includes_dtype=False) assert meta[name].shape == (1, 3, 224, 224) assert meta[name].dtype is None
def test_parse_legacy(self, name): # Legacy argument format used comma. meta_args = ["{:},1x3x224x224".format(name)] meta = args_util.parse_meta(meta_args, includes_dtype=False) assert meta[name].shape == (1, 3, 224, 224) assert meta[name].dtype is None
def run_impl(self, args): def missing_meta_tensors(input_metadata, output_metadata): missing = TensorMetadata() for name, (dtype, shape) in input_metadata.items(): if dtype is None or shape is None: missing.add(name, dtype, shape) for name, (dtype, shape) in output_metadata.items(): if dtype is None: missing.add(name, dtype, shape) return missing model = super().load_model() user_input_metadata = args_util.parse_meta(args.input_meta) user_output_metadata = args_util.parse_meta(args.output_meta, includes_shape=False) # Loads an ONNX-GS graph and create new I/O metadata w/ info missing in user_input/output_metadata. def load_graph_and_io_meta(model): graph = gs.import_onnx(model) TENSOR_MAP = graph.tensors() def get_tensor(name): if name not in TENSOR_MAP: G_LOGGER.exit( "Tensor: {:} does not exist in the model.".format( name)) return TENSOR_MAP[name] # Makes a TensorMetadata for inputs/outputs using either the user provided information # or details derived from tensors. def make_io_meta(user_meta, tensors): if not user_meta: return tools_util.meta_from_gs_tensors(tensors) new_meta = copy.copy(user_meta) for name, (dtype, shape) in new_meta.items(): tensor = get_tensor(name) new_meta.add(name, dtype or tensor.dtype, shape or tensor.shape) return new_meta input_metadata = make_io_meta(user_input_metadata, graph.inputs) output_metadata = make_io_meta(user_output_metadata, graph.outputs) return graph, input_metadata, output_metadata graph, input_metadata, output_metadata = load_graph_and_io_meta(model) # If we've already done ONNX shape inference, we should not do it again here. skip_shape_inference = self.arg_groups[ OnnxShapeInferenceArgs].force_fallback or self.arg_groups[ OnnxShapeInferenceArgs].do_shape_inference if missing_meta_tensors(input_metadata, output_metadata) and not skip_shape_inference: G_LOGGER.info( "Running ONNX shape inference to derive shapes and/or data types for `auto` arguments.\n" "To avoid this, you can specify the shapes and data types explicitly." ) model = onnx_backend.infer_shapes(model) graph, input_metadata, output_metadata = load_graph_and_io_meta( model) missing_tensors = missing_meta_tensors(input_metadata, output_metadata) if missing_tensors or self.arg_groups[ OnnxShapeInferenceArgs].force_fallback: # Use ONNX runtime with static shapes to infer shapes when all else fails # Returns a TensorMetadata for all tensors in the graph. if not self.arg_groups[OnnxShapeInferenceArgs].force_fallback: G_LOGGER.warning( "Some tensor shapes or dtypes are missing in the model. Note: Tensors with missing information:\n{:}\n" "Will run inference to determine shapes. This may cause some dynamic " "dimensions to become static.\n" "To avoid this, please provide metadata on the command-line. " .format(missing_tensors)) else: G_LOGGER.info( "Forcing fallback shape inference. This will cause dynamic dimensions to become static." ) _, layerwise_meta = self.arg_groups[ OnnxShapeInferenceArgs].fallback_inference(model) def update_meta_from_layerwise(meta, user_meta, set_shapes=True): for name in meta: user_dtype, user_shape = None, None if name in user_meta: user_dtype, user_shape = user_meta[ name].dtype, user_meta[name].shape # Choose between what the user set, what's in the model, and what # fallback shape inference said. def choose_meta(user, model, fallback): if self.arg_groups[ OnnxShapeInferenceArgs].force_fallback: return user or fallback return user or model or fallback if name in layerwise_meta: meta[name].dtype = choose_meta( user_dtype, meta[name].dtype, layerwise_meta[name].dtype) if set_shapes: meta[name].shape = choose_meta( user_shape, meta[name].shape, layerwise_meta[name].shape) G_LOGGER.verbose( "Updated tensor: {:} metadata to: {:}".format( name, meta[name])) return meta input_metadata = update_meta_from_layerwise( input_metadata, user_input_metadata) output_metadata = update_meta_from_layerwise( output_metadata, user_output_metadata, set_shapes=self.arg_groups[OnnxShapeInferenceArgs]. force_fallback) graph = onnx_backend.extract_subgraph(graph, input_metadata, output_metadata) super().save_model(super().export_graph(graph))