add softmax and argmax while exporting models

paddlex/cv/models/base.py

@@ -34,6 +34,7 @@ from paddlex.utils import (seconds_to_hms, get_single_card_bs, dict2str,
                            _get_shared_memory_size_in_M, EarlyStop)
 import paddlex.utils.logging as logging
 from .slim.prune import _pruner_eval_fn, _pruner_template_input, sensitive_prune
+from .utils.infer_nets import InferNet
 
 
 class BaseModel:
@@ -578,13 +579,18 @@ class BaseModel:
         pipeline_info["version"] = "1.0.0"
         return pipeline_info
 
+    def _build_inference_net(self):
+        infer_net = InferNet(self.net, self.model_type)
+        infer_net.eval()
+        return infer_net
+
     def _export_inference_model(self, save_dir, image_shape=None):
         save_dir = osp.join(save_dir, 'inference_model')
-        self.net.eval()
         self.test_inputs = self._get_test_inputs(image_shape)
+        infer_net = self._build_inference_net()
 
         if self.status == 'Quantized':
-            self.quantizer.save_quantized_model(self.net,
+            self.quantizer.save_quantized_model(infer_net,
                                                 osp.join(save_dir, 'model'),
                                                 self.test_inputs)
             quant_info = self.get_quant_info()
@@ -594,7 +600,7 @@ class BaseModel:
                 yaml.dump(quant_info, f)
         else:
             static_net = paddle.jit.to_static(
-                self.net, input_spec=self.test_inputs)
+                infer_net, input_spec=self.test_inputs)
             paddle.jit.save(static_net, osp.join(save_dir, 'model'))
 
         if self.status == 'Pruned':

paddlex/cv/models/classifier.py

@@ -440,7 +440,7 @@ class BaseClassifier(BaseModel):
         with paddle.no_grad():
             outputs = self.run(self.net, im, mode='test')
         prediction = outputs['prediction'].numpy()
-        prediction = self._postprocess(prediction, true_topk, self.labels)
+        prediction = self._postprocess(prediction, true_topk)
         if isinstance(img_file, (str, np.ndarray)):
             prediction = prediction[0]
 
@@ -456,16 +456,18 @@ class BaseClassifier(BaseModel):
 
         if to_tensor:
             batch_im = paddle.to_tensor(batch_im)
+        else:
+            batch_im = np.asarray(batch_im)
 
         return batch_im,
 
-    def _postprocess(self, results, true_topk, labels):
+    def _postprocess(self, results, true_topk):
         preds = list()
         for i, pred in enumerate(results):
             pred_label = np.argsort(pred)[::-1][:true_topk]
             preds.append([{
                 'category_id': l,
-                'category': labels[l],
+                'category': self.labels[l],
                 'score': results[i][l]
             } for l in pred_label])
 

paddlex/cv/models/segmenter.py

@@ -107,7 +107,7 @@ class BaseSegmenter(BaseModel):
                 logit, origin_shape, transforms=inputs[2])
             label_map = paddle.argmax(
                 score_map, axis=1, keepdim=True, dtype='int32')
-            score_map = paddle.max(score_map, axis=1, keepdim=True)
+            score_map = paddle.nn.functional.softmax(score_map, axis=1)
             score_map = paddle.squeeze(score_map)
             label_map = paddle.squeeze(label_map)
             outputs = {'label_map': label_map, 'score_map': score_map}
@@ -510,6 +510,8 @@ class BaseSegmenter(BaseModel):
             batch_ori_shape.append(ori_shape)
         if to_tensor:
             batch_im = paddle.to_tensor(batch_im)
+        else:
+            batch_im = np.asarray(batch_im)
 
         return batch_im, batch_ori_shape
 

paddlex/cv/models/utils/infer_nets.py

@@ -0,0 +1,43 @@
+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import paddle
+
+
+class PostProcessor(paddle.nn.Layer):
+    def __init__(self, model_type):
+        self.model_type = model_type
+
+    def forward(self, net_outputs):
+        if self.model_type == 'classifier':
+            outputs = paddle.nn.functional.softmax(net_outputs, axis=1)
+        elif self.model_type == 'segmenter':
+            outputs = paddle.squeeze(paddle.nn.functional.softmax(net_outputs, axis=1)), \
+                      paddle.squeeze(paddle.argmax(net_outputs, axis=1))
+        else:
+            outputs = net_outputs
+        return outputs
+
+
+class InferNet(paddle.nn.Layer):
+    def __init__(self, net, model_type):
+        super(InferNet, self).__init__()
+        self.net = net
+        self.postprocessor = PostProcessor(model_type)
+
+    def forward(self, x):
+        net_outputs = self.net(x)
+        outputs = self.postprocessor(net_outputs)
+
+        return outputs

paddlex/deploy.py

@@ -14,6 +14,7 @@
 
 import os.path as osp
 import numpy as np
+import paddle.nn.functional as F
 from paddle.inference import Config
 from paddle.inference import create_predictor
 from paddle.inference import PrecisionType
@@ -131,19 +132,26 @@ class Predictor(object):
         preprocessed_samples = self._model._preprocess(
             images, transforms, to_tensor=False)
         if self._model.model_type == 'classifier':
-            batch_samples = {'image': preprocessed_samples[0]}
+            preprocessed_samples = {'image': preprocessed_samples[0]}
         elif self._model.model_type == 'segmenter':
-            batch_samples = {
+            preprocessed_samples = {
                 'image': preprocessed_samples[0],
                 'ori_shape': preprocessed_samples[1]
             }
         elif self._model.model_type == 'detector':
-            batch_samples = preprocessed_samples
+            pass
         else:
             logging.error(
                 "Invalid model type {}".format(self._model.model_type),
                 exit=True)
-        return batch_samples
+        return preprocessed_samples
+
+    def postprocess(self, net_outputs, topk=1):
+        if self._model.model_type == 'classifier':
+            true_topk = min(self._model.num_classes, topk)
+            preds = self._model._postprocess(net_outputs, true_topk)
+
+        return preds
 
     def raw_predict(self, inputs):
         """ 接受预处理过后的数据进行预测
@@ -151,6 +159,23 @@ class Predictor(object):
             Args:
                 inputs(dict): 预处理过后的数据
         """
+        input_names = self.predictor.get_input_names()
+        for name in input_names:
+            input_tensor = self.predictor.get_input_handle(name)
+            input_tensor.copy_from_cpu(inputs[name])
+
+        self.timer.inference_time_s.start()
+        self.predictor.run()
+        output_names = self.predictor.get_output_names()
+        if self._model.model_type == 'classifier':
+            net_outputs = F.softmax(self.predictor.get_output_handle(name))
+
+        net_outputs = list()
+        for name in output_names:
+            output_tensor = self.predictor.get_output_handle(name)
+            net_outputs.append(output_tensor.copy_to_cpu())
+
+        return net_outputs
 
     def predict(self, img_file, topk=1, transforms=None):
         """ 图片预测
@@ -161,7 +186,7 @@ class Predictor(object):
                 topk(int): 分类预测时使用，表示预测前topk的结果。
                 transforms (paddlex.transforms): 数据预处理操作。
         """
-        if transforms is None and not hasattr(self, 'test_transforms'):
+        if transforms is None and not hasattr(self._model, 'test_transforms'):
             raise Exception("Transforms need to be defined, now is None.")
         if transforms is None:
             transforms = self._model.test_transforms
@@ -171,14 +196,13 @@ class Predictor(object):
             images = img_file
 
         self.timer.preprocess_time_s.start()
-        batch_samples = self.preprocess(images, transforms)
+        preprocessed_input = self.preprocess(images, transforms)
         self.timer.preprocess_time_s.end()
 
-        input_names = self.predictor.get_input_names()
-        for name in input_names:
-            input_tensor = self.predictor.get_input_handle(name)
-            input_tensor.copy_from_cpu(batch_samples[name])
-
         self.timer.inference_time_s.start()
-        self.predictor.run()
-        output_names = self.predictor.get_output_names()
+        net_outputs = self.raw_predict(preprocessed_input)
+        self.timer.inference_time_s.end()
+
+        self.timer.postprocess_time_s.start()
+        results = self.postprocess(net_outputs, topk)
+        print(results)