def load_neuropod(neuropod_path, _always_use_native=True, **kwargs): """ Load a neuropod package. Returns a NeuropodExecutor :param neuropod_path The path to a neuropod package :param visible_gpu: The index of the GPU that this Neuropod should run on (if any). This is either `None` or a nonnegative integer. Setting this to `None` will attempt to run this model on CPU. :param load_custom_ops: Whether or not to load custom ops included in the model. """ if _always_use_native: return NativeNeuropodExecutor(neuropod_path, **kwargs) # If we were given a zipfile, extract it to a temp dir and use it neuropod_path = zip_loader.extract_neuropod_if_necessary(neuropod_path) # Figure out what type of neuropod this is neuropod_config = config_utils.read_neuropod_config(neuropod_path) platform = neuropod_config["platform"] if platform == "python": from neuropod.backends.python.executor import PythonNeuropodExecutor return PythonNeuropodExecutor(neuropod_path, **kwargs) elif platform == "torchscript": from neuropod.backends.torchscript.executor import TorchScriptNeuropodExecutor return TorchScriptNeuropodExecutor(neuropod_path, **kwargs) elif platform == "tensorflow": from neuropod.backends.tensorflow.executor import TensorflowNeuropodExecutor return TensorflowNeuropodExecutor(neuropod_path, **kwargs) else: raise ValueError( "Invalid platform found in neuropod config: {}".format(platform))
def __init__(self, neuropod_path): # Read the neuropod config self.neuropod_config = config_utils.read_neuropod_config(neuropod_path) # Generate the tensor to device mapping self.input_device_mapping = { tensor["name"]: self.neuropod_config["input_tensor_device"][tensor["name"]] for tensor in self.inputs }