Add -layout option (#10272)

samples/python/speech_sample/arg_parser.py

@@ -4,8 +4,8 @@
 import argparse
 
 
-def parse_args() -> argparse.Namespace:
-    """Parse and return command line arguments"""
+def build_arg_parser() -> argparse.ArgumentParser:
+    """Create and return argument parser"""
     parser = argparse.ArgumentParser(add_help=False)
     args = parser.add_argument_group('Options')
     model = parser.add_mutually_exclusive_group(required=True)
@@ -25,8 +25,10 @@ def parse_args() -> argparse.Namespace:
                       'CPU, GPU, MYRIAD, GNA_AUTO, GNA_HW, GNA_SW_FP32, GNA_SW_EXACT and HETERO with combination of GNA'
                       ' as the primary device and CPU as a secondary (e.g. HETERO:GNA,CPU) are supported. '
                       'The sample will look for a suitable plugin for device specified. Default value is CPU.')
-    args.add_argument('-bs', '--batch_size', default=1, type=int, choices=range(1, 9), metavar='[1-8]',
-                      help='Optional. Batch size 1-8 (default 1).')
+    args.add_argument('-bs', '--batch_size', type=int, choices=range(1, 9), metavar='[1-8]',
+                      help='Optional. Batch size 1-8.')
+    args.add_argument('-layout', type=str,
+                      help='Optional. Custom layout in format: "input0[value0],input1[value1]" or "[value]" (applied to all inputs)')
     args.add_argument('-qb', '--quantization_bits', default=16, type=int, choices=(8, 16), metavar='[8, 16]',
                       help='Optional. Weight bits for quantization: 8 or 16 (default 16).')
     args.add_argument('-sf', '--scale_factor', type=str,
@@ -67,6 +69,12 @@ def parse_args() -> argparse.Namespace:
                       help='Optional. The maximum percent of error for PWL function. '
                       'The value must be in <0, 100> range. The default value is 1.0.')
 
+    return parser
+
+
+def parse_args() -> argparse.Namespace:
+    """Parse and validate command-line arguments"""
+    parser = build_arg_parser()
     args = parser.parse_args()
 
     if args.context_window_left < 0:

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(',')

samples/python/speech_sample/utils.py

@@ -6,9 +6,10 @@ import argparse
 import logging as log
 import re
 import sys
-from typing import List, Tuple
+from typing import Dict, List, Tuple
 
 import numpy as np
+from openvino.runtime import Output
 
 # Operating Frequency for GNA HW devices for Core and Atom architecture
 GNA_CORE_FREQUENCY = 400
@@ -78,3 +79,10 @@ def parse_outputs_from_args(args: argparse.Namespace) -> Tuple[List[str], List[i
     except ValueError:
         log.error('Incorrect value for -oname/--output_layers option, please specify a port for each output layer.')
         sys.exit(-4)
+
+
+def parse_input_layouts(args: argparse.Namespace, inputs: List[Output]) -> Dict[str, str]:
+    if ',' in args.layout:
+        return dict([_input.split('[') for _input in args.layout[:-1].split('],')])
+    else:
+        return {_input.get_any_name(): args.layout[1:-1] for _input in inputs}