Merge pull request #3210 from anzhella-pankratova/python_demos/use_output_tensor_as_key

Wovchena · web-flow · commit bb680b042f60 · 2022-02-14T18:02:18.000+03:00
Python demos: modify output result accessing in case of single output layer
diff --git a/demos/3d_segmentation_demo/python/3d_segmentation_demo.py b/demos/3d_segmentation_demo/python/3d_segmentation_demo.py
@@ -278,6 +278,7 @@ def main():
     n, c, d, h, w = model.inputs[0].shape
 
     compiled_model = core.compile_model(model, args.target_device)
+    output_tensor = compiled_model.outputs[0]
     infer_request = compiled_model.create_infer_request()
     log.info('The model {} is loaded to {}'.format(args.path_to_model, args.target_device))
 
@@ -308,8 +309,8 @@ def main():
         original_data = data_crop
         original_size = original_data.shape[-3:]
 
-    result = infer_request.infer({input_tensor_name: data_crop})
-    result = next(iter(result.values()))
+    input_data = {input_tensor_name: data_crop}
+    result = infer_request.infer(input_data)[output_tensor]
     batch, channels, out_d, out_h, out_w = result.shape
 
     list_img = []
diff --git a/demos/action_recognition_demo/python/action_recognition_demo/models.py b/demos/action_recognition_demo/python/action_recognition_demo/models.py
@@ -93,12 +93,13 @@ def __init__(self, model_path, core, target_device, num_requests, model_type):
             log.error("Demo supports only models with 1 output")
             sys.exit(1)
 
-        self.compiled_model = core.compile_model(self.model, target_device)
+        compiled_model = core.compile_model(self.model, target_device)
+        self.output_tensor = compiled_model.outputs[0]
         self.input_name = self.model.inputs[0].get_any_name()
         self.input_shape = self.model.inputs[0].shape
 
         self.num_requests = num_requests
-        self.infer_requests = [self.compiled_model.create_infer_request() for _ in range(self.num_requests)]
+        self.infer_requests = [compiled_model.create_infer_request() for _ in range(self.num_requests)]
         log.info('The {} model {} is loaded to {}'.format(model_type, model_path, target_device))
 
     def async_infer(self, frame, req_id):
@@ -107,7 +108,7 @@ def async_infer(self, frame, req_id):
 
     def wait_request(self, req_id):
         self.infer_requests[req_id].wait()
-        return next(iter(self.infer_requests[req_id].results.values()))
+        return self.infer_requests[req_id].results[self.output_tensor]
 
 
 class DummyDecoder:
diff --git a/demos/colorization_demo/python/colorization_demo.py b/demos/colorization_demo/python/colorization_demo.py
@@ -80,6 +80,7 @@ def main(args):
     assert len(model.outputs) == 1, "Expected number of outputs is equal 1"
 
     compiled_model = core.compile_model(model, device_name=args.device)
+    output_tensor = compiled_model.outputs[0]
     infer_request = compiled_model.create_infer_request()
     log.info('The model {} is loaded to {}'.format(args.model, args.device))
 
@@ -115,7 +116,7 @@ def main(args):
 
         inputs[input_tensor_name] = np.expand_dims(img_l_rs, axis=[0, 1])
 
-        res = next(iter(infer_request.infer(inputs).values()))
+        res = infer_request.infer(inputs)[output_tensor]
 
         update_res = np.squeeze(res)
 
diff --git a/demos/face_recognition_demo/python/ie_module.py b/demos/face_recognition_demo/python/ie_module.py
@@ -31,12 +31,13 @@ def __init__(self, core, model_path, model_type):
     def deploy(self, device, plugin_config, max_requests=1):
         self.max_requests = max_requests
         compiled_model = self.core.compile_model(self.model, device, config=plugin_config)
+        self.output_tensor = compiled_model.outputs[0]
         self.infer_queue = AsyncInferQueue(compiled_model, self.max_requests)
         self.infer_queue.set_callback(self.completion_callback)
         log.info('The {} model {} is loaded to {}'.format(self.model_type, self.model_path, device))
 
     def completion_callback(self, infer_request, id):
-        self.outputs[id] = next(iter(infer_request.results.values()))
+        self.outputs[id] = infer_request.results[self.output_tensor]
 
     def enqueue(self, input):
         self.clear()
diff --git a/demos/gpt2_text_prediction_demo/python/gpt2_text_prediction_demo.py b/demos/gpt2_text_prediction_demo/python/gpt2_text_prediction_demo.py
@@ -99,6 +99,7 @@ def main():
 
     # load model to the device
     compiled_model = ie.compile_model(model, args.device)
+    output_tensor = compiled_model.outputs[0]
     infer_request = compiled_model.create_infer_request()
     log.info('The model {} is loaded to {}'.format(args.model, args.device))
 
@@ -148,13 +149,12 @@ def prompts():
 
             # infer by IE
             t_start = time.perf_counter()
-            res = infer_request.infer(inputs)
+            outputs = infer_request.infer(inputs)[output_tensor]
             t_end = time.perf_counter()
             t_count += 1
             log.info("Sequence of length {} is processed with {:0.2f} requests/sec ({:0.2} sec per request)".format(
                 model_input.shape[1], 1 / (t_end - t_start), t_end - t_start))
 
-            outputs = next(iter(res.values()))
             next_token_logits = outputs[:, cur_input_len-1, :]
 
             # pre-process distribution
diff --git a/demos/image_inpainting_demo/python/inpainting.py b/demos/image_inpainting_demo/python/inpainting.py
@@ -26,6 +26,7 @@ def __init__(self, core, model_path, device='CPU'):
         self.image_input_layer, self.mask_input_layer = sorted([node.get_any_name() for node in model.inputs])
 
         compiled_model = core.compile_model(model, device)
+        self.output_tensor = compiled_model.outputs[0]
         self.infer_request = compiled_model.create_infer_request()
 
         self.nchw_layout = model.input(self.image_input_layer).shape[1] == 3
@@ -44,11 +45,9 @@ def __init__(self, core, model_path, device='CPU'):
         self.input_height = input_height
         self.input_width = input_width
 
-
     def infer(self, image, mask):
-        output = self.infer_request.infer(inputs={self.image_input_layer: image, self.mask_input_layer: mask})
-        return next(iter(output.values()))
-
+        input_data = {self.image_input_layer: image, self.mask_input_layer: mask}
+        return self.infer_request.infer(input_data)[self.output_tensor]
 
     def process(self, image, mask):
         if self.nchw_layout:
diff --git a/demos/image_retrieval_demo/python/image_retrieval_demo/image_retrieval.py b/demos/image_retrieval_demo/python/image_retrieval_demo/image_retrieval.py
@@ -40,16 +40,16 @@ def __init__(self, model_path, device, cpu_extension):
         self.model = core.read_model(model_path)
         self.input_tensor_name = "Placeholder"
         compiled_model = core.compile_model(self.model, device)
+        self.output_tensor = compiled_model.outputs[0]
         self.infer_request = compiled_model.create_infer_request()
         log.info('The model {} is loaded to {}'.format(model_path, device))
 
     def predict(self, image):
         ''' Takes input image and returns L2-normalized embedding vector. '''
 
-        assert len(image.shape) == 4
         image = np.transpose(image, (0, 3, 1, 2))
-        out = next(iter(self.infer_request.infer({self.input_tensor_name: image}).values()))
-        return out
+        input_data = {self.input_tensor_name: image}
+        return self.infer_request.infer(input_data)[self.output_tensor]
 
 
 class ImageRetrieval:
diff --git a/demos/image_translation_demo/python/image_translation_demo/models.py b/demos/image_translation_demo/python/image_translation_demo/models.py
@@ -25,6 +25,7 @@ def __init__(self, core, model_path, device='CPU'):
         self.input_image_size = model.input(self.reference_image).shape
 
         compiled_model = core.compile_model(model, device)
+        self.output_tensor = compiled_model.outputs[0]
         self.infer_request = compiled_model.create_infer_request()
 
     def infer(self, input_semantics, reference_image, reference_semantics):
@@ -33,8 +34,7 @@ def infer(self, input_semantics, reference_image, reference_semantics):
             self.reference_image: reference_image,
             self.reference_semantics: reference_semantics
         }
-        result = self.infer_request.infer(input_data)
-        return next(iter(result.values()))
+        return self.infer_request.infer(input_data)[self.output_tensor]
 
 
 class SegmentationModel:
@@ -49,9 +49,9 @@ def __init__(self, core, model_path, device='CPU'):
         self.input_size = model.inputs[0].shape
 
         compiled_model = core.compile_model(model, device)
+        self.output_tensor = compiled_model.outputs[0]
         self.infer_request = compiled_model.create_infer_request()
 
     def infer(self, input):
         input_data = {self.input_tensor_name: input}
-        result = self.infer_request.infer(input_data)
-        return next(iter(result.values()))
+        return self.infer_request.infer(input_data)[self.output_tensor]
diff --git a/demos/place_recognition_demo/python/place_recognition_demo/place_recognition.py b/demos/place_recognition_demo/python/place_recognition_demo/place_recognition.py
@@ -41,17 +41,17 @@ def __init__(self, model_path, device, cpu_extension):
         self.input_size = self.model.input(self.input_tensor_name).shape
         self.nchw_layout = self.input_size[1] == 3
         compiled_model = core.compile_model(self.model, device)
+        self.output_tensor = compiled_model.outputs[0]
         self.infer_request = compiled_model.create_infer_request()
         log.info('The model {} is loaded to {}'.format(model_path, device))
 
     def predict(self, image):
         ''' Takes input image and returns L2-normalized embedding vector. '''
 
-        assert len(image.shape) == 4
         if self.nchw_layout:
             image = np.transpose(image, (0, 3, 1, 2))
-        out = next(iter(self.infer_request.infer({self.input_tensor_name: image}).values()))
-        return out
+        input_data = {self.input_tensor_name: image}
+        return self.infer_request.infer(input_data)[self.output_tensor]
 
 
 class PlaceRecognition:
diff --git a/demos/single_human_pose_estimation_demo/python/detector.py b/demos/single_human_pose_estimation_demo/python/detector.py
@@ -21,6 +21,7 @@ def __init__(self, core, model_path, label_class, thr=0.3, device='CPU'):
             raise RuntimeError("Expected model output shape with {} outputs".format(OUTPUT_SIZE))
 
         compiled_model = core.compile_model(self.model, device)
+        self.output_tensor = compiled_model.outputs[0]
         self.infer_request = compiled_model.create_infer_request()
         self.input_tensor_name = self.model.inputs[0].get_any_name()
         if self.nchw_layout:
@@ -38,7 +39,8 @@ def _preprocess(self, img):
         return img[None, ]
 
     def _infer(self, prep_img):
-        output = next(iter(self.infer_request.infer({self.input_tensor_name: prep_img}).values()))
+        input_data = {self.input_tensor_name: prep_img}
+        output = self.infer_request.infer(input_data)[self.output_tensor]
         return output[0][0]
 
     def _postprocess(self, bboxes):
diff --git a/demos/single_human_pose_estimation_demo/python/estimator.py b/demos/single_human_pose_estimation_demo/python/estimator.py
@@ -99,6 +99,7 @@ def __init__(self, core, model_path, device='CPU'):
             raise RuntimeError("Expected model output shape [1, {}, H, W]".format(OUTPUT_CHANNELS_SIZE))
 
         compiled_model = core.compile_model(self.model, device)
+        self.output_tensor = compiled_model.outputs[0]
         self.infer_request = compiled_model.create_infer_request()
         self.input_tensor_name = self.model.inputs[0].get_any_name()
         self._transform = TransformedCrop(input_shape[2], input_shape[3], output_shape[2], output_shape[3])
@@ -107,7 +108,8 @@ def _preprocess(self, img, bbox):
         return self._transform(img, bbox)
 
     def _infer(self, prep_img):
-        output = next(iter(self.infer_request.infer({self.input_tensor_name: prep_img}).values()))
+        input_data = {self.input_tensor_name: prep_img}
+        output = self.infer_request.infer(input_data)[self.output_tensor]
         return output[0]
 
     @staticmethod
diff --git a/demos/sound_classification_demo/python/sound_classification_demo.py b/demos/sound_classification_demo/python/sound_classification_demo.py
@@ -163,6 +163,7 @@ def main():
         sys.exit(1)
 
     compiled_model = ie.compile_model(model, args.device)
+    output_tensor = compiled_model.outputs[0]
     infer_request = compiled_model.create_infer_request()
     log.info('The model {} is loaded to {}'.format(args.model, args.device))
 
@@ -178,7 +179,7 @@ def main():
     clips = 0
     for idx, chunk in enumerate(audio.chunks(length, hop, num_chunks=batch_size)):
         chunk = np.reshape(chunk, model.inputs[0].shape)
-        output = next(iter(infer_request.infer({input_tensor_name: chunk}).values()))
+        output = infer_request.infer({input_tensor_name: chunk})[output_tensor]
         clips += batch_size
         for batch, data in enumerate(output):
             chunk_start_time = (idx*batch_size + batch)*hop / audio.samplerate
diff --git a/demos/speech_recognition_quartznet_demo/python/speech_recognition_quartznet_demo.py b/demos/speech_recognition_quartznet_demo/python/speech_recognition_quartznet_demo.py
@@ -64,12 +64,13 @@ def __init__(self, core, model_path, input_shape, device):
             raise RuntimeError(f'QuartzNet output third dimension size must be {len(self.alphabet) + 1}')
         model.reshape({self.input_tensor_name: PartialShape(input_shape)})
         compiled_model = core.compile_model(model, device)
+        self.output_tensor = compiled_model.outputs[0]
         self.infer_request = compiled_model.create_infer_request()
         log.info('The model {} is loaded to {}'.format(model_path, device))
 
     def infer(self, melspectrogram):
         input_data = {self.input_tensor_name: melspectrogram}
-        return next(iter(self.infer_request.infer(input_data).values()))
+        return self.infer_request.infer(input_data)[self.output_tensor]
 
     @classmethod
     def audio_to_melspectrum(cls, audio, sampling_rate):
diff --git a/demos/speech_recognition_wav2vec_demo/python/speech_recognition_wav2vec_demo.py b/demos/speech_recognition_wav2vec_demo/python/speech_recognition_wav2vec_demo.py
@@ -69,6 +69,7 @@ def __init__(self, core, model_path, input_shape, device, vocab_file):
             raise RuntimeError(f'Wav2Vec output third dimension size must be {len(self.alphabet)}')
         model.reshape({self.input_tensor_name: PartialShape(input_shape)})
         compiled_model = core.compile_model(model, device)
+        self.output_tensor = compiled_model.outputs[0]
         self.infer_request = compiled_model.create_infer_request()
         log.info('The model {} is loaded to {}'.format(model_path, device))
         self._init_vocab(vocab_file)
@@ -93,7 +94,7 @@ def preprocess(sound):
 
     def infer(self, audio):
         input_data = {self.input_tensor_name: audio}
-        return next(iter(self.infer_request.infer(input_data).values()))
+        return self.infer_request.infer(input_data)[self.output_tensor]
 
     def decode(self, logits):
         token_ids = np.squeeze(np.argmax(logits, -1))
diff --git a/demos/text_spotting_demo/python/text_spotting_demo.py b/demos/text_spotting_demo/python/text_spotting_demo.py
@@ -202,11 +202,12 @@ def main():
     mask_rcnn_infer_request = mask_rcnn_compiled_model.create_infer_request()
     log.info('The Mask-RCNN model {} is loaded to {}'.format(args.mask_rcnn_model, args.device))
 
-    text_enc_compiled_model =  core.compile_model(text_enc_model, args.device)
+    text_enc_compiled_model = core.compile_model(text_enc_model, args.device)
+    text_enc_output_tensor = text_enc_compiled_model.outputs[0]
     text_enc_infer_request = text_enc_compiled_model.create_infer_request()
     log.info('The Text Recognition Encoder model {} is loaded to {}'.format(args.text_enc_model, args.device))
 
-    text_dec_compiled_model =  core.compile_model(text_dec_model, args.device)
+    text_dec_compiled_model = core.compile_model(text_dec_model, args.device)
     text_dec_infer_request = text_dec_compiled_model.create_infer_request()
     log.info('The Text Recognition Decoder model {} is loaded to {}'.format(args.text_dec_model, args.device))
 
@@ -288,7 +289,8 @@ def main():
 
         texts = []
         for feature in text_features:
-            feature = next(iter(text_enc_infer_request.infer({'input': np.expand_dims(feature, axis=0)}).values()))
+            input_data = {'input': np.expand_dims(feature, axis=0)}
+            feature = text_enc_infer_request.infer(input_data)[text_enc_output_tensor]
             feature = np.reshape(feature, (feature.shape[0], feature.shape[1], -1))
             feature = np.transpose(feature, (0, 2, 1))
 
diff --git a/demos/whiteboard_inpainting_demo/python/utils/ie_tools.py b/demos/whiteboard_inpainting_demo/python/utils/ie_tools.py
@@ -37,8 +37,8 @@ def _preprocess(self, img):
 
     def forward(self, img):
         """Performs forward pass of the wrapped IE model"""
-        res = self.infer_request.infer(inputs={self.input_tensor_name: self._preprocess(img)})
-        return next(iter(res.values()))
+        input_data = {self.input_tensor_name: self._preprocess(img)}
+        return self.infer_request.infer(input_data)[self.output_tensor]
 
     def get_detections(self, input):
         raise NotImplementedError
@@ -71,4 +71,5 @@ def load_model(self, core, model_path, device, cpu_extension=''):
         self.input_tensor_name = self.model.inputs[0].get_any_name()
         # Loading model to the plugin
         compiled_model = core.compile_model(self.model, device)
+        self.output_tensor = compiled_model.outputs[0]
         self.infer_request = compiled_model.create_infer_request()
