support uniform output

Author: zhangyubo0722

ppcls/engine/engine.py

@@ -16,6 +16,7 @@
 
 import os
 import shutil
+import copy
 import platform
 import paddle
 import paddle.distributed as dist
@@ -38,6 +39,7 @@
 from ppcls.utils.ema import ExponentialMovingAverage
 from ppcls.utils.save_load import load_dygraph_pretrain
 from ppcls.utils.save_load import init_model
+from ppcls.utils.save_result import update_train_results
 from ppcls.utils import save_load, save_predict_result
 
 from ppcls.data.utils.get_image_list import get_image_list
@@ -169,8 +171,8 @@ def __init__(self, config, mode="train"):
                         self.config["DataLoader"]["Eval"], "Gallery",
                         self.device, self.use_dali)
                     self.query_dataloader = build_dataloader(
-                        self.config["DataLoader"]["Eval"], "Query",
-                        self.device, self.use_dali)
+                        self.config["DataLoader"]["Eval"], "Query", self.device,
+                        self.use_dali)
 
         # build loss
         if self.mode == "train":
@@ -210,8 +212,8 @@ def __init__(self, config, mode="train"):
                                    self.config["Global"]["eval_during_train"]):
             if self.eval_mode == "classification":
                 if "Metric" in self.config and "Eval" in self.config["Metric"]:
-                    self.eval_metric_func = build_metrics(self.config["Metric"]
-                                                          ["Eval"])
+                    self.eval_metric_func = build_metrics(self.config["Metric"][
+                        "Eval"])
                 else:
                     self.eval_metric_func = None
             elif self.eval_mode == "retrieval":
@@ -266,8 +268,7 @@ def __init__(self, config, mode="train"):
             self.model = paddle.DataParallel(self.model)
             if self.mode == 'train' and len(self.train_loss_func.parameters(
             )) > 0:
-                self.train_loss_func = paddle.DataParallel(
-                    self.train_loss_func)
+                self.train_loss_func = paddle.DataParallel(self.train_loss_func)
 
             # set different seed in different GPU manually in distributed environment
             if seed is None:
@@ -313,6 +314,8 @@ def train(self):
         }
         # global iter counter
         self.global_step = 0
+        uniform_output_enabled = self.config['Global'].get(
+            "uniform_output_enabled", False)
 
         if self.config.Global.checkpoints is not None:
             metric_info = init_model(self.config.Global, self.model,
@@ -384,41 +387,89 @@ def train(self):
                 # save best model from best_acc or best_ema_acc
                 if max(acc, acc_ema) >= max(best_metric["metric"],
                                             best_metric_ema):
+                    metric_info = {
+                        "metric": max(acc, acc_ema),
+                        "epoch": epoch_id
+                    }
+                    prefix = "best_model"
                     save_load.save_model(
                         self.model,
                         self.optimizer,
-                        {"metric": max(acc, acc_ema),
-                         "epoch": epoch_id},
-                        self.output_dir,
+                        metric_info,
+                        os.path.join(self.output_dir, prefix)
+                        if uniform_output_enabled else self.output_dir,
                         ema=ema_module,
                         model_name=self.config["Arch"]["name"],
-                        prefix="best_model",
+                        prefix=prefix,
                         loss=self.train_loss_func,
                         save_student_model=True)
+                    if uniform_output_enabled:
+                        save_path = os.path.join(self.output_dir, prefix,
+                                                 "inference")
+                        self.export(save_path, uniform_output_enabled)
+                        if self.ema:
+                            ema_save_path = os.path.join(
+                                self.output_dir, prefix, "inference_ema")
+                            self.export(ema_save_path, uniform_output_enabled)
+                        update_train_results(
+                            self.config, prefix, metric_info, ema=self.ema)
+                        save_load.save_model_info(metric_info, self.output_dir,
+                                                  prefix)
 
                 self.model.train()
 
             # save model
             if save_interval > 0 and epoch_id % save_interval == 0:
+                metric_info = {"metric": acc, "epoch": epoch_id}
+                prefix = "epoch_{}".format(epoch_id)
                 save_load.save_model(
                     self.model,
-                    self.optimizer, {"metric": acc,
-                                     "epoch": epoch_id},
-                    self.output_dir,
+                    self.optimizer,
+                    metric_info,
+                    os.path.join(self.output_dir, prefix)
+                    if uniform_output_enabled else self.output_dir,
                     ema=ema_module,
                     model_name=self.config["Arch"]["name"],
-                    prefix="epoch_{}".format(epoch_id),
+                    prefix=prefix,
                     loss=self.train_loss_func)
+                if uniform_output_enabled:
+                    save_path = os.path.join(self.output_dir, prefix,
+                                             "inference")
+                    self.export(save_path, uniform_output_enabled)
+                    if self.ema:
+                        ema_save_path = os.path.join(self.output_dir, prefix,
+                                                     "inference_ema")
+                        self.export(ema_save_path, uniform_output_enabled)
+                    update_train_results(
+                        self.config,
+                        prefix,
+                        metric_info,
+                        done_flag=epoch_id == self.config["Global"]["epochs"],
+                        ema=self.ema)
+                    save_load.save_model_info(metric_info, self.output_dir,
+                                              prefix)
             # save the latest model
+            metric_info = {"metric": acc, "epoch": epoch_id}
+            prefix = "latest"
             save_load.save_model(
                 self.model,
-                self.optimizer, {"metric": acc,
-                                 "epoch": epoch_id},
-                self.output_dir,
+                self.optimizer,
+                metric_info,
+                os.path.join(self.output_dir, prefix)
+                if uniform_output_enabled else self.output_dir,
                 ema=ema_module,
                 model_name=self.config["Arch"]["name"],
-                prefix="latest",
+                prefix=prefix,
                 loss=self.train_loss_func)
+            if uniform_output_enabled:
+                save_path = os.path.join(self.output_dir, prefix, "inference")
+                self.export(save_path, uniform_output_enabled)
+                if self.ema:
+                    ema_save_path = os.path.join(self.output_dir, prefix,
+                                                 "inference_ema")
+                    self.export(ema_save_path, uniform_output_enabled)
+                save_load.save_model_info(metric_info, self.output_dir, prefix)
+                self.model.train()
 
         if self.vdl_writer is not None:
             self.vdl_writer.close()
@@ -479,33 +530,45 @@ def infer(self):
                     image_file_list.clear()
             except Exception as ex:
                 logger.error(
-                    "Exception occured when parse line: {} with msg: {}".
-                    format(image_file, ex))
+                    "Exception occured when parse line: {} with msg: {}".format(
+                        image_file, ex))
                 continue
         if save_path:
             save_predict_result(save_path, results)
         return results
 
-    def export(self):
-        assert self.mode == "export"
+    def export(self,
+               save_path=None,
+               uniform_output_enabled=False,
+               ema_module=None):
+        assert self.mode == "export" or uniform_output_enabled
+        if paddle.distributed.get_rank() != 0:
+            return
         use_multilabel = self.config["Global"].get(
             "use_multilabel",
             False) or "ATTRMetric" in self.config["Metric"]["Eval"][0]
-        model = ExportModel(self.config["Arch"], self.model, use_multilabel)
-        if self.config["Global"]["pretrained_model"] is not None:
+        model = self.model_ema.module if self.ema else self.model
+        if isinstance(self.model, paddle.DataParallel):
+            model = copy.deepcopy(model._layers)
+        else:
+            model = copy.deepcopy(model)
+        model = ExportModel(self.config["Arch"], model
+                            if not ema_module else ema_module, use_multilabel)
+        if self.config["Global"][
+                "pretrained_model"] is not None and not uniform_output_enabled:
             load_dygraph_pretrain(model.base_model,
                                   self.config["Global"]["pretrained_model"])
-
         model.eval()
-
         # for re-parameterization nets
-        for layer in self.model.sublayers():
+        for layer in model.sublayers():
             if hasattr(layer, "re_parameterize") and not getattr(layer,
                                                                  "is_repped"):
                 layer.re_parameterize()
-
-        save_path = os.path.join(self.config["Global"]["save_inference_dir"],
-                                 "inference")
+        if not save_path:
+            save_path = os.path.join(
+                self.config["Global"]["save_inference_dir"], "inference")
+        else:
+            save_path = os.path.join(save_path, "inference")
 
         model = paddle.jit.to_static(
             model,
@@ -520,12 +583,12 @@ def export(self):
                                                           save_path + "_int8")
         else:
             paddle.jit.save(model, save_path)
-        if self.config["Global"].get("export_for_fd", False):
-            dst_path = os.path.join(
-                self.config["Global"]["save_inference_dir"], 'inference.yml')
+        if self.config["Global"].get("export_for_fd",
+                                     False) or uniform_output_enabled:
+            dst_path = os.path.join(os.path.dirname(save_path), 'inference.yml')
             dump_infer_config(self.config, dst_path)
         logger.info(
-            f"Export succeeded! The inference model exported has been saved in \"{self.config['Global']['save_inference_dir']}\"."
+            f"Export succeeded! The inference model exported has been saved in \"{save_path}\"."
         )
 
     def _init_amp(self):

ppcls/utils/__init__.py

@@ -17,7 +17,7 @@
 from . import misc
 from . import model_zoo
 
-from .config import get_config
+from .config import get_config, convert_to_dict
 from .dist_utils import all_gather
 from .metrics import accuracy_score
 from .metrics import hamming_distance

ppcls/utils/config.py

@@ -20,7 +20,16 @@
 from . import check
 from collections import OrderedDict
 
-__all__ = ['get_config']
+__all__ = ['get_config', 'convert_to_dict']
+
+
+def convert_to_dict(obj):
+    if isinstance(obj, dict):
+        return {k: convert_to_dict(v) for k, v in obj.items()}
+    elif isinstance(obj, list):
+        return [convert_to_dict(i) for i in obj]
+    else:
+        return obj
 
 
 class AttrDict(dict):
@@ -223,16 +232,49 @@ def setup_orderdict():
     yaml.add_representer(OrderedDict, represent_dictionary_order)
 
 
-def dump_infer_config(config, path):
+def dump_infer_config(inference_config, path):
     setup_orderdict()
     infer_cfg = OrderedDict()
+    config = copy.deepcopy(inference_config)
     if config.get("Infer"):
         transforms = config["Infer"]["transforms"]
     elif config["DataLoader"]["Eval"].get("Query"):
-        transforms = config["DataLoader"]["Eval"]["Query"]["dataset"]["transform_ops"]
+        transforms = config["DataLoader"]["Eval"]["Query"]["dataset"][
+            "transform_ops"]
         transforms.append({"ToCHWImage": None})
     else:
         logger.error("This config does not support dump transform config!")
+    transform = next((item for item in transforms if 'CropImage' in item), None)
+    if transform:
+        dynamic_shapes = transform["CropImage"]["size"]
+    else:
+        transform = next((item for item in transforms
+                          if 'ResizeImage' in item), None)
+        if transform:
+            dynamic_shapes = transform["ResizeImage"]["size"][0]
+        else:
+            dynamic_shapes = 224
+    # Configuration required config for high-performance inference.
+    if config["Global"].get("hpi_config_path", None):
+        hpi_config = convert_to_dict(
+            parse_config(config["Global"]["hpi_config_path"]))
+        if hpi_config["Hpi"]["backend_config"].get("paddle_tensorrt", None):
+            hpi_config["Hpi"]["backend_config"]["paddle_tensorrt"][
+                "dynamic_shapes"]["x"] = [[
+                    1, 3, dynamic_shapes, dynamic_shapes
+                ] for i in range(3)]
+            hpi_config["Hpi"]["backend_config"]["paddle_tensorrt"][
+                "max_batch_size"] = 1
+        if hpi_config["Hpi"]["backend_config"].get("tensorrt", None):
+            hpi_config["Hpi"]["backend_config"]["tensorrt"]["dynamic_shapes"][
+                "x"] = [[1, 3, dynamic_shapes, dynamic_shapes]
+                        for i in range(3)]
+            hpi_config["Hpi"]["backend_config"]["tensorrt"][
+                "max_batch_size"] = 1
+        infer_cfg["Hpi"] = hpi_config["Hpi"]
+    if config["Global"].get("pdx_model_name", None):
+        infer_cfg["Global"] = {}
+        infer_cfg["Global"]["model_name"] = config["Global"]["pdx_model_name"]
     for transform in transforms:
         if "NormalizeImage" in transform:
             transform["NormalizeImage"]["channel_num"] = 3
@@ -262,7 +304,7 @@ def dump_infer_config(config, path):
         postprocess_dict.pop("name")
         dic = OrderedDict()
         for item in postprocess_dict.items():
-           dic[item[0]] = item[1]
+            dic[item[0]] = item[1]
         dic['label_list'] = label_names
 
         if postprocess_name:

ppcls/utils/save_load.py

@@ -18,6 +18,7 @@
 
 import errno
 import os
+import json
 
 import paddle
 from . import logger
@@ -108,8 +109,7 @@ def load_distillation_model(model, pretrained_model):
     student = model.student if hasattr(model,
                                        "student") else model._layers.student
     load_dygraph_pretrain(teacher, path=pretrained_model[0])
-    logger.info("Finish initing teacher model from {}".format(
-        pretrained_model))
+    logger.info("Finish initing teacher model from {}".format(pretrained_model))
     # load student model
     if len(pretrained_model) >= 2:
         load_dygraph_pretrain(student, path=pretrained_model[1])
@@ -188,8 +188,7 @@ def save_model(net,
     params_state_dict = net.state_dict()
     if loss is not None:
         loss_state_dict = loss.state_dict()
-        keys_inter = set(params_state_dict.keys()) & set(loss_state_dict.keys(
-        ))
+        keys_inter = set(params_state_dict.keys()) & set(loss_state_dict.keys())
         assert len(keys_inter) == 0, \
             f"keys in model and loss state_dict must be unique, but got intersection {keys_inter}"
         params_state_dict.update(loss_state_dict)
@@ -210,3 +209,15 @@ def save_model(net,
     paddle.save([opt.state_dict() for opt in optimizer], model_path + ".pdopt")
     paddle.save(metric_info, model_path + ".pdstates")
     logger.info("Already save model in {}".format(model_path))
+
+
+def save_model_info(model_info, save_path, prefix):
+    """
+    save model info to the target path
+    """
+    save_path = os.path.join(save_path, prefix)
+    if not os.path.exists(save_path):
+        os.makedirs(save_path)
+    with open(os.path.join(save_path, f'{prefix}.info.json'), 'w') as f:
+        json.dump(model_info, f)
+    logger.info("Already save model info in {}".format(save_path))