|
|
|
@ -75,7 +75,7 @@ class AutoBackend(nn.Module): |
|
|
|
|
fp16 &= pt or jit or onnx or engine or nn_module # FP16 |
|
|
|
|
nhwc = coreml or saved_model or pb or tflite or edgetpu # BHWC formats (vs torch BCWH) |
|
|
|
|
stride = 32 # default stride |
|
|
|
|
model = None # TODO: resolves ONNX inference, verify effect on other backends |
|
|
|
|
model, metadata = None, None |
|
|
|
|
cuda = torch.cuda.is_available() and device.type != 'cpu' # use CUDA |
|
|
|
|
if not (pt or triton or nn_module): |
|
|
|
|
w = attempt_download_asset(w) # download if not local |
|
|
|
@ -105,10 +105,7 @@ class AutoBackend(nn.Module): |
|
|
|
|
model = torch.jit.load(w, _extra_files=extra_files, map_location=device) |
|
|
|
|
model.half() if fp16 else model.float() |
|
|
|
|
if extra_files['config.txt']: # load metadata dict |
|
|
|
|
d = json.loads(extra_files['config.txt'], |
|
|
|
|
object_hook=lambda d: {int(k) if k.isdigit() else k: v |
|
|
|
|
for k, v in d.items()}) |
|
|
|
|
stride, names = int(d['stride']), d['names'] |
|
|
|
|
metadata = json.loads(extra_files['config.txt'], object_hook=lambda x: dict(x.items())) |
|
|
|
|
elif dnn: # ONNX OpenCV DNN |
|
|
|
|
LOGGER.info(f'Loading {w} for ONNX OpenCV DNN inference...') |
|
|
|
|
check_requirements('opencv-python>=4.5.4') |
|
|
|
@ -120,23 +117,23 @@ class AutoBackend(nn.Module): |
|
|
|
|
providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if cuda else ['CPUExecutionProvider'] |
|
|
|
|
session = onnxruntime.InferenceSession(w, providers=providers) |
|
|
|
|
output_names = [x.name for x in session.get_outputs()] |
|
|
|
|
meta = session.get_modelmeta().custom_metadata_map # metadata |
|
|
|
|
if 'stride' in meta: |
|
|
|
|
stride, names = int(meta['stride']), eval(meta['names']) |
|
|
|
|
metadata = session.get_modelmeta().custom_metadata_map # metadata |
|
|
|
|
elif xml: # OpenVINO |
|
|
|
|
LOGGER.info(f'Loading {w} for OpenVINO inference...') |
|
|
|
|
check_requirements('openvino') # requires openvino-dev: https://pypi.org/project/openvino-dev/ |
|
|
|
|
from openvino.runtime import Core, Layout, get_batch # noqa |
|
|
|
|
ie = Core() |
|
|
|
|
if not Path(w).is_file(): # if not *.xml |
|
|
|
|
w = next(Path(w).glob('*.xml')) # get *.xml file from *_openvino_model dir |
|
|
|
|
network = ie.read_model(model=w, weights=Path(w).with_suffix('.bin')) |
|
|
|
|
w = Path(w) |
|
|
|
|
if not w.is_file(): # if not *.xml |
|
|
|
|
w = next(w.glob('*.xml')) # get *.xml file from *_openvino_model dir |
|
|
|
|
network = ie.read_model(model=str(w), weights=w.with_suffix('.bin')) |
|
|
|
|
if network.get_parameters()[0].get_layout().empty: |
|
|
|
|
network.get_parameters()[0].set_layout(Layout('NCHW')) |
|
|
|
|
batch_dim = get_batch(network) |
|
|
|
|
if batch_dim.is_static: |
|
|
|
|
batch_size = batch_dim.get_length() |
|
|
|
|
executable_network = ie.compile_model(network, device_name='CPU') # device_name="MYRIAD" for NCS2 |
|
|
|
|
metadata = w.parent / 'metadata.yaml' |
|
|
|
|
elif engine: # TensorRT |
|
|
|
|
LOGGER.info(f'Loading {w} for TensorRT inference...') |
|
|
|
|
import tensorrt as trt # https://developer.nvidia.com/nvidia-tensorrt-download |
|
|
|
@ -148,7 +145,7 @@ class AutoBackend(nn.Module): |
|
|
|
|
# Read file |
|
|
|
|
with open(w, 'rb') as f, trt.Runtime(logger) as runtime: |
|
|
|
|
meta_len = int.from_bytes(f.read(4), byteorder='little') # read metadata length |
|
|
|
|
meta = json.loads(f.read(meta_len).decode('utf-8')) # read metadata |
|
|
|
|
metadata = json.loads(f.read(meta_len).decode('utf-8')) # read metadata |
|
|
|
|
model = runtime.deserialize_cuda_engine(f.read()) # read engine |
|
|
|
|
context = model.create_execution_context() |
|
|
|
|
bindings = OrderedDict() |
|
|
|
@ -171,18 +168,17 @@ class AutoBackend(nn.Module): |
|
|
|
|
bindings[name] = Binding(name, dtype, shape, im, int(im.data_ptr())) |
|
|
|
|
binding_addrs = OrderedDict((n, d.ptr) for n, d in bindings.items()) |
|
|
|
|
batch_size = bindings['images'].shape[0] # if dynamic, this is instead max batch size |
|
|
|
|
stride, names = int(meta['stride']), meta['names'] |
|
|
|
|
elif coreml: # CoreML |
|
|
|
|
LOGGER.info(f'Loading {w} for CoreML inference...') |
|
|
|
|
import coremltools as ct |
|
|
|
|
model = ct.models.MLModel(w) |
|
|
|
|
names, stride, task = (model.user_defined_metadata.get(k) for k in ('names', 'stride', 'task')) |
|
|
|
|
names, stride = eval(names), int(stride) |
|
|
|
|
metadata = model.user_defined_metadata |
|
|
|
|
elif saved_model: # TF SavedModel |
|
|
|
|
LOGGER.info(f'Loading {w} for TensorFlow SavedModel inference...') |
|
|
|
|
import tensorflow as tf |
|
|
|
|
keras = False # assume TF1 saved_model |
|
|
|
|
model = tf.keras.models.load_model(w) if keras else tf.saved_model.load(w) |
|
|
|
|
metadata = Path(w) / 'metadata.yaml' |
|
|
|
|
elif pb: # GraphDef https://www.tensorflow.org/guide/migrate#a_graphpb_or_graphpbtxt |
|
|
|
|
LOGGER.info(f'Loading {w} for TensorFlow GraphDef inference...') |
|
|
|
|
import tensorflow as tf |
|
|
|
@ -221,23 +217,23 @@ class AutoBackend(nn.Module): |
|
|
|
|
with contextlib.suppress(zipfile.BadZipFile): |
|
|
|
|
with zipfile.ZipFile(w, 'r') as model: |
|
|
|
|
meta_file = model.namelist()[0] |
|
|
|
|
meta = ast.literal_eval(model.read(meta_file).decode('utf-8')) |
|
|
|
|
stride, names = int(meta['stride']), meta['names'] |
|
|
|
|
metadata = ast.literal_eval(model.read(meta_file).decode('utf-8')) |
|
|
|
|
elif tfjs: # TF.js |
|
|
|
|
raise NotImplementedError('YOLOv8 TF.js inference is not supported') |
|
|
|
|
elif paddle: # PaddlePaddle |
|
|
|
|
LOGGER.info(f'Loading {w} for PaddlePaddle inference...') |
|
|
|
|
check_requirements('paddlepaddle-gpu' if cuda else 'paddlepaddle') |
|
|
|
|
import paddle.inference as pdi |
|
|
|
|
if not Path(w).is_file(): # if not *.pdmodel |
|
|
|
|
w = next(Path(w).rglob('*.pdmodel')) # get *.pdmodel file from *_paddle_model dir |
|
|
|
|
weights = Path(w).with_suffix('.pdiparams') |
|
|
|
|
config = pdi.Config(str(w), str(weights)) |
|
|
|
|
w = Path(w) |
|
|
|
|
if not w.is_file(): # if not *.pdmodel |
|
|
|
|
w = next(w.rglob('*.pdmodel')) # get *.pdmodel file from *_paddle_model dir |
|
|
|
|
config = pdi.Config(str(w), str(w.with_suffix('.pdiparams'))) |
|
|
|
|
if cuda: |
|
|
|
|
config.enable_use_gpu(memory_pool_init_size_mb=2048, device_id=0) |
|
|
|
|
predictor = pdi.create_predictor(config) |
|
|
|
|
input_handle = predictor.get_input_handle(predictor.get_input_names()[0]) |
|
|
|
|
output_names = predictor.get_output_names() |
|
|
|
|
metadata = w.parents[1] / 'metadata.yaml' |
|
|
|
|
elif triton: # NVIDIA Triton Inference Server |
|
|
|
|
LOGGER.info('Triton Inference Server not supported...') |
|
|
|
|
''' |
|
|
|
@ -254,14 +250,16 @@ class AutoBackend(nn.Module): |
|
|
|
|
f'\n\n{EXPORT_FORMATS_TABLE}') |
|
|
|
|
|
|
|
|
|
# Load external metadata YAML |
|
|
|
|
w = Path(w) |
|
|
|
|
if xml or saved_model or paddle: |
|
|
|
|
metadata = (w if saved_model else w.parents[1] if paddle else w.parent) / 'metadata.yaml' |
|
|
|
|
if metadata.exists(): |
|
|
|
|
metadata = yaml_load(metadata) |
|
|
|
|
stride, names = int(metadata['stride']), metadata['names'] # load metadata |
|
|
|
|
else: |
|
|
|
|
LOGGER.warning(f"WARNING ⚠️ Metadata not found at '{metadata}'") |
|
|
|
|
if isinstance(metadata, (str, Path)) and Path(metadata).exists(): |
|
|
|
|
metadata = yaml_load(metadata) |
|
|
|
|
if metadata: |
|
|
|
|
stride = int(metadata['stride']) |
|
|
|
|
task = metadata['task'] |
|
|
|
|
batch = int(metadata['batch']) |
|
|
|
|
imgsz = eval(metadata['imgsz']) if isinstance(metadata['imgsz'], str) else metadata['imgsz'] |
|
|
|
|
names = eval(metadata['names']) if isinstance(metadata['names'], str) else metadata['names'] |
|
|
|
|
elif not (pt or triton or nn_module): |
|
|
|
|
LOGGER.warning(f"WARNING ⚠️ Metadata not found for 'model={weights}'") |
|
|
|
|
|
|
|
|
|
# Check names |
|
|
|
|
if 'names' not in locals(): # names missing |
|
|
|
|