Add -layout option (#10272)

samples/python/speech_sample/speech_sample.py

@@ -17,8 +17,8 @@ from arg_parser import parse_args
 from file_options import read_utterance_file, write_utterance_file
 from utils import (GNA_ATOM_FREQUENCY, GNA_CORE_FREQUENCY,
                    compare_with_reference, get_scale_factor, log,
-                   parse_outputs_from_args, parse_scale_factors,
-                   set_scale_factors)
+                   parse_input_layouts, parse_outputs_from_args,
+                   parse_scale_factors, set_scale_factors)
 
 
 def do_inference(data: Dict[str, np.ndarray], infer_request: InferRequest, cw_l: int = 0, cw_r: int = 0) -> np.ndarray:
@@ -26,7 +26,7 @@ def do_inference(data: Dict[str, np.ndarray], infer_request: InferRequest, cw_l:
     frames_to_infer = {}
     result = {}
 
-    batch_size = infer_request.get_input_tensor(0).shape[0]
+    batch_size = infer_request.model_inputs[0].shape[0]
     num_of_frames = next(iter(data.values())).shape[0]
 
     for output in infer_request.model_outputs:
@@ -83,24 +83,32 @@ def main():
             output_layer_names, output_layer_ports = parse_outputs_from_args(args)
             model.add_outputs(list(zip(output_layer_names, output_layer_ports)))
 
+        if args.layout:
+            layouts = parse_input_layouts(args, model.inputs)
+
         ppp = PrePostProcessor(model)
 
         for i in range(len(model.inputs)):
-            ppp.input(i).tensor() \
-                .set_element_type(Type.f32) \
-                .set_layout(Layout('NC'))  # noqa: N400
+            ppp.input(i).tensor().set_element_type(Type.f32)
 
-            ppp.input(i).model().set_layout(Layout('NC'))
+            input_name = model.input(i).get_any_name()
+
+            if args.layout and input_name in layouts.keys():
+                ppp.input(i).tensor().set_layout(Layout(layouts[input_name]))
+                ppp.input(i).model().set_layout(Layout(layouts[input_name]))
 
         for i in range(len(model.outputs)):
             ppp.output(i).tensor().set_element_type(Type.f32)
 
         model = ppp.build()
 
-        if args.context_window_left == args.context_window_right == 0:
-            set_batch(model, args.batch_size)
-        else:
-            set_batch(model, 1)
+        if args.batch_size:
+            batch_size = args.batch_size if args.context_window_left == args.context_window_right == 0 else 1
+
+            if any([not _input.node.layout.empty for _input in model.inputs]):
+                set_batch(model, batch_size)
+            else:
+                log.warning('Layout is not set for any input, so custom batch size is not set')
 
 # ---------------------------Step 4. Configure plugin ---------------------------------------------------------
     devices = args.device.replace('HETERO:', '').split(',')