def call_impl(self): """ save_uff (bool): Whether to write the generated UFF and corresponding PBTXT files. """ import uff from polygraphy.backend.tf import util as tf_util G_LOGGER.module_info(uff) graph, output_names = self.tf_loader() output_names = [name.split(":")[0] for name in output_names] # GraphDefs don't have names, so we have to name it something generic. output_filename = None if not self.uff_path else "out.uff" # Generate the UFF model and get information about the input_buffers/output_buffers. uff_model, input_nodes, _ = uff.from_tensorflow( graph.as_graph_def(), return_graph_info=True, quiet=(G_LOGGER.severity > G_LOGGER.VERBOSE), debug_mode=(G_LOGGER.severity == G_LOGGER.EXTRA_VERBOSE), text=self.uff_path, save_preprocessed=self.uff_path, output_filename=output_filename, preprocessor=self.preprocessor, ) input_names = [node.name for node in input_nodes] input_shapes = [ tuple(int(dim.size) for dim in node.attr["shape"].shape.dim) for node in input_nodes ] return uff_model, input_names, input_shapes, output_names
def import_mod(): from polygraphy import config from polygraphy.logger import G_LOGGER, LogMode def install_mod(raise_error=True): modname = name.split(".")[0] pkg = _MODULE_TO_PKG_NAME.get(modname, modname) extra_flags = _MODULE_EXTRA_FLAGS.get(modname, []) if version == LATEST_VERSION: extra_flags.append("--upgrade") elif version is not None: pkg += version cmd = config.INSTALL_CMD + [pkg] + extra_flags G_LOGGER.info( "{:} is required, but not installed. Attempting to install now.\n" "Running: {:}".format(pkg, " ".join(cmd))) status = sp.run(cmd) if status.returncode != 0: G_LOGGER.log( "Could not automatically install required package: {:}. Please install it manually." .format(pkg), severity=G_LOGGER.CRITICAL if raise_error else G_LOGGER.WARNING, ) mod = importlib.import_module(name) return mod mod = None try: mod = importlib.import_module(name) except: if config.AUTOINSTALL_DEPS: mod = install_mod() else: G_LOGGER.error( "Module: {:} is required but could not be imported.\n" "You can try setting POLYGRAPHY_AUTOINSTALL_DEPS=1 in your environment variables " "to allow Polygraphy to automatically install missing packages.\n" "Note that this may cause existing packages to be overwritten - hence, it may be " "desirable to use a Python virtual environment or container. " .format(name)) raise # Auto-upgrade if necessary if version is not None and hasattr( mod, "__version__") and not _version_ok(mod.__version__, version): if config.AUTOINSTALL_DEPS: G_LOGGER.info( "Note: Package: '{name}' version {cur_ver} is installed, but version {rec_ver} is recommended.\n" "Upgrading...".format(name=name, cur_ver=mod.__version__, rec_ver=version)) mod = install_mod( raise_error=False ) # We can try to use the other version if install fails. elif version != LATEST_VERSION: G_LOGGER.warning( "Package: '{name}' version {cur_ver} is installed, but version {rec_ver} is recommended.\n" "Consider installing the recommended version or setting POLYGRAPHY_AUTOINSTALL_DEPS=1 in your " "environment variables to do so automatically. ".format( name=name, cur_ver=mod.__version__, rec_ver=version), mode=LogMode.ONCE, ) if log: G_LOGGER.module_info(mod) return mod