MCPcopy Create free account
hub / github.com/NVIDIA/TensorRT / parse_impl

Method parse_impl

tools/Polygraphy/polygraphy/tools/args/model.py:139–211  ·  view source on GitHub ↗

Parses command-line arguments and populates the following attributes: Attributes: input_shapes (TensorMetadata): Input names and their shapes. path (str): Path to the model. model_type (ModelArgs.ModelType): The type of model. extra_m

(self, args)

Source from the content-addressed store, hash-verified

137 )
138
139 def parse_impl(self, args):
140 """
141 Parses command-line arguments and populates the following attributes:
142
143 Attributes:
144 input_shapes (TensorMetadata): Input names and their shapes.
145 path (str): Path to the model.
146 model_type (ModelArgs.ModelType): The type of model.
147 extra_model_info (str):
148 Any extra model information specified after the model argument, separated by a colon.
149 The meaning of this information may be specific to each model type.
150 In most cases, no extra model information is provided.
151 """
152
153 def determine_model_type(model_file):
154 model_type = args_util.get(args, "model_type")
155 if model_type is not None:
156 return model_type.lower()
157
158 if model_file is None:
159 return None
160
161 def use_ext(ext_mapping):
162 file_ext = os.path.splitext(model_file)[-1]
163 if file_ext in ext_mapping:
164 return ext_mapping[file_ext]
165
166 runner_opts = []
167 if self._guess_model_type_from_runners:
168 if not hasattr(self.arg_groups[RunnerSelectArgs], "runners"):
169 G_LOGGER.internal_error(
170 "RunnerSelectArgs must be parsed before ModelArgs when `guess_model_type_from_runners` is enabled!"
171 )
172 runner_opts = list(self.arg_groups[RunnerSelectArgs].runners.keys())
173
174 if args_util.get(args, "ckpt") or os.path.isdir(model_file):
175 return "ckpt"
176 elif "tf" in runner_opts or "trt-legacy" in runner_opts:
177 if args_util.get(args, "caffe_model"):
178 return "caffe"
179 return use_ext(ModelArgs.EXT_MODEL_TYPE_MAPPING) or "frozen"
180 else:
181 model_type = use_ext(ModelArgs.EXT_MODEL_TYPE_MAPPING)
182 if model_type:
183 return model_type
184
185 G_LOGGER.critical(
186 f"Could not automatically determine model type for: {model_file}"
187 f"\nPlease explicitly specify the type with the --model-type option"
188 )
189
190 self.input_shapes = TensorMetadata()
191 if args_util.get(args, "input_shapes"):
192 self.input_shapes = args_util.parse_meta(args_util.get(args, "input_shapes"), includes_dtype=False)
193
194 self.path = None
195 self.extra_model_info = None
196

Callers

nothing calls this directly

Calls 4

TensorMetadataClass · 0.90
getMethod · 0.45
criticalMethod · 0.45
defaultMethod · 0.45

Tested by

no test coverage detected