def test_set_batch_int():
model = create_test_model()
model_param1 = model.get_parameters()[0]
model_param2 = model.get_parameters()[1]
# batch == 2
model_param1.set_layout(Layout("NC"))
assert get_batch(model) == 2
# set batch to 1
set_batch(model, 1)
assert get_batch(model) == 1
# check if shape of param 1 has changed
assert model_param1.get_output_shape(0) == PartialShape([1, 1])
# check if shape of param 2 has not changed
assert model_param2.get_output_shape(0) == PartialShape([2, 1])
def test_set_batch_int():
param1 = ops.parameter(Shape([2, 1]), dtype=np.float32, name="data1")
param2 = ops.parameter(Shape([2, 1]), dtype=np.float32, name="data2")
add = ops.add(param1, param2)
func = Model(add, [param1, param2], "TestFunction")
func_param1 = func.get_parameters()[0]
func_param2 = func.get_parameters()[1]
# batch == 2
func_param1.set_layout(Layout("NC"))
assert get_batch(func) == 2
# set batch to 1
set_batch(func, 1)
assert get_batch(func) == 1
# check if shape of param 1 has changed
assert str(func_param1.get_output_shape(0) == {1, 1})
# check if shape of param 2 has not changed
assert str(func_param2.get_output_shape(0) == {2, 1})
def test_get_batch():
param1 = ops.parameter(Shape([2, 1]), dtype=np.float32, name="data1")
param2 = ops.parameter(Shape([2, 1]), dtype=np.float32, name="data2")
add = ops.add(param1, param2)
func = Model(add, [param1, param2], "TestFunction")
param = func.get_parameters()[0]
param.set_layout(Layout("NC"))
assert get_batch(func) == 2
def test_get_batch_CHWN():
param1 = ops.parameter(Shape([3, 1, 3, 4]), dtype=np.float32, name="data1")
param2 = ops.parameter(Shape([3, 1, 3, 4]), dtype=np.float32, name="data2")
param3 = ops.parameter(Shape([3, 1, 3, 4]), dtype=np.float32, name="data3")
add = ops.add(param1, param2)
add2 = ops.add(add, param3)
func = Model(add2, [param1, param2, param3], "TestFunction")
param = func.get_parameters()[0]
param.set_layout(Layout("CHWN"))
assert get_batch(func) == 4
def __init__(self,
weights='yolov5s.pt',
device=torch.device('cpu'),
dnn=False,
data=None,
fp16=False):
# Usage:
# PyTorch: weights = *.pt
# TorchScript: *.torchscript
# ONNX Runtime: *.onnx
# ONNX OpenCV DNN: *.onnx with --dnn
# OpenVINO: *.xml
# CoreML: *.mlmodel
# TensorRT: *.engine
# TensorFlow SavedModel: *_saved_model
# TensorFlow GraphDef: *.pb
# TensorFlow Lite: *.tflite
# TensorFlow Edge TPU: *_edgetpu.tflite
from models.experimental import attempt_download, attempt_load # scoped to avoid circular import
super().__init__()
w = str(weights[0] if isinstance(weights, list) else weights)
pt, jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs = self.model_type(
w) # get backend
w = attempt_download(w) # download if not local
fp16 &= (pt or jit or onnx or engine) and device.type != 'cpu' # FP16
stride, names = 32, [f'class{i}'
for i in range(1000)] # assign defaults
if data: # assign class names (optional)
with open(data, errors='ignore') as f:
names = yaml.safe_load(f)['names']
if pt: # PyTorch
model = attempt_load(weights if isinstance(weights, list) else w,
device=device)
stride = max(int(model.stride.max()), 32) # model stride
names = model.module.names if hasattr(
model, 'module') else model.names # get class names
model.half() if fp16 else model.float()
self.model = model # explicitly assign for to(), cpu(), cuda(), half()
elif jit: # TorchScript
LOGGER.info(f'Loading {w} for TorchScript inference...')
extra_files = {'config.txt': ''} # model metadata
model = torch.jit.load(w, _extra_files=extra_files)
model.half() if fp16 else model.float()
if extra_files['config.txt']:
d = json.loads(extra_files['config.txt']) # extra_files dict
stride, names = int(d['stride']), d['names']
elif dnn: # ONNX OpenCV DNN
LOGGER.info(f'Loading {w} for ONNX OpenCV DNN inference...')
check_requirements(('opencv-python>=4.5.4', ))
net = cv2.dnn.readNetFromONNX(w)
elif onnx: # ONNX Runtime
LOGGER.info(f'Loading {w} for ONNX Runtime inference...')
cuda = torch.cuda.is_available()
check_requirements(
('onnx', 'onnxruntime-gpu' if cuda else 'onnxruntime'))
import onnxruntime
providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'
] if cuda else ['CPUExecutionProvider']
session = onnxruntime.InferenceSession(w, providers=providers)
meta = session.get_modelmeta().custom_metadata_map # metadata
if 'stride' in meta:
stride, names = int(meta['stride']), eval(meta['names'])
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
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'))
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 Intel NCS2
output_layer = next(iter(executable_network.outputs))
meta = Path(w).with_suffix('.yaml')
if meta.exists():
stride, names = self._load_metadata(meta) # load metadata
elif engine: # TensorRT
LOGGER.info(f'Loading {w} for TensorRT inference...')
import tensorrt as trt # https://developer.nvidia.com/nvidia-tensorrt-download
check_version(trt.__version__, '7.0.0',
hard=True) # require tensorrt>=7.0.0
Binding = namedtuple('Binding',
('name', 'dtype', 'shape', 'data', 'ptr'))
logger = trt.Logger(trt.Logger.INFO)
with open(w, 'rb') as f, trt.Runtime(logger) as runtime:
model = runtime.deserialize_cuda_engine(f.read())
bindings = OrderedDict()
fp16 = False # default updated below
for index in range(model.num_bindings):
name = model.get_binding_name(index)
dtype = trt.nptype(model.get_binding_dtype(index))
shape = tuple(model.get_binding_shape(index))
data = torch.from_numpy(np.empty(
shape, dtype=np.dtype(dtype))).to(device)
bindings[name] = Binding(name, dtype, shape, data,
int(data.data_ptr()))
if model.binding_is_input(index) and dtype == np.float16:
fp16 = True
binding_addrs = OrderedDict(
(n, d.ptr) for n, d in bindings.items())
context = model.create_execution_context()
batch_size = bindings['images'].shape[0]
elif coreml: # CoreML
LOGGER.info(f'Loading {w} for CoreML inference...')
import coremltools as ct
model = ct.models.MLModel(w)
else: # TensorFlow (SavedModel, GraphDef, Lite, Edge TPU)
if saved_model: # 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)
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
def wrap_frozen_graph(gd, inputs, outputs):
x = tf.compat.v1.wrap_function(
lambda: tf.compat.v1.import_graph_def(gd, name=""),
[]) # wrapped
ge = x.graph.as_graph_element
return x.prune(tf.nest.map_structure(ge, inputs),
tf.nest.map_structure(ge, outputs))
gd = tf.Graph().as_graph_def() # graph_def
with open(w, 'rb') as f:
gd.ParseFromString(f.read())
frozen_func = wrap_frozen_graph(gd,
inputs="x:0",
outputs="Identity:0")
elif tflite or edgetpu: # https://www.tensorflow.org/lite/guide/python#install_tensorflow_lite_for_python
try: # https://coral.ai/docs/edgetpu/tflite-python/#update-existing-tf-lite-code-for-the-edge-tpu
from tflite_runtime.interpreter import Interpreter, load_delegate
except ImportError:
import tensorflow as tf
Interpreter, load_delegate = tf.lite.Interpreter, tf.lite.experimental.load_delegate,
if edgetpu: # Edge TPU https://coral.ai/software/#edgetpu-runtime
LOGGER.info(
f'Loading {w} for TensorFlow Lite Edge TPU inference...'
)
delegate = {
'Linux': 'libedgetpu.so.1',
'Darwin': 'libedgetpu.1.dylib',
'Windows': 'edgetpu.dll'
}[platform.system()]
interpreter = Interpreter(
model_path=w,
experimental_delegates=[load_delegate(delegate)])
else: # Lite
LOGGER.info(
f'Loading {w} for TensorFlow Lite inference...')
interpreter = Interpreter(
model_path=w) # load TFLite model
interpreter.allocate_tensors() # allocate
input_details = interpreter.get_input_details() # inputs
output_details = interpreter.get_output_details() # outputs
elif tfjs:
raise Exception(
'ERROR: YOLOv5 TF.js inference is not supported')
else:
raise Exception(f'ERROR: {w} is not a supported format')
self.__dict__.update(locals()) # assign all variables to self
def test_get_batch():
model = create_test_model()
param = model.get_parameters()[0]
param.set_layout(Layout("NC"))
assert get_batch(model) == 2
