feat: add convnext model and config (#201)

Co-authored-by: GuoxiaWang <mingzilaochongtu@gmail.com>

Author: qizhaoaoe

README.md

@@ -23,6 +23,7 @@
 * [Swin](./task/classification/swin/)
 * [DeiT](./task/classification/deit/)
 * [CaiT](./task/classification/cait/)
+* [ConvNeXt](./task/classification/convnext)
 * [MoCo v3](./task/ssl/mocov3/)
 * [MAE](./task/ssl/mae/)
 * [ConvMAE](./task/ssl/mae/)

plsc/engine/classification/evaluation.py

@@ -41,6 +41,7 @@ def default_eval(engine, epoch_id=0):
         dataset) if not engine.use_dali else engine.eval_dataloader.size
     max_iter = len(engine.eval_dataloader) - 1 if platform.system(
     ) == "Windows" else len(engine.eval_dataloader)
+
     for iter_id, batch in enumerate(engine.eval_dataloader):
         if iter_id >= max_iter:
             break
@@ -63,7 +64,6 @@ def default_eval(engine, epoch_id=0):
                 custom_black_list=engine.fp16_custom_black_list,
                 level=engine.fp16_level):
             out = engine.model(batch[0])
-
             # calc loss
             if engine.eval_loss_func is not None:
                 loss_dict = engine.eval_loss_func(out, batch[1])
@@ -132,6 +132,7 @@ def default_eval(engine, epoch_id=0):
                 len(engine.eval_dataloader), metric_msg, time_msg, ips_msg))
 
         tic = time.time()
+
     if engine.use_dali:
         engine.eval_dataloader.reset()
 

plsc/engine/classification/train.py

@@ -84,7 +84,9 @@ def default_train_one_epoch(engine, epoch_id):
         if iter_id % engine.print_batch_step == 0:
             log_info(engine, batch_size, epoch_id, iter_id)
         tic = time.time()
-
+        # ema update
+        if engine.enabled_ema:
+            engine.ema.update()
         # eval model and save model if possible
         eval_metric_info = {
             "epoch": epoch_id,

plsc/engine/engine.py

@@ -38,7 +38,7 @@
 from plsc.optimizer import build_optimizer
 from plsc.utils import io
 from plsc.core import recompute_warp, GradScaler, param_sync
-
+from plsc.models.utils import EMA
 from . import classification
 from . import recognition
 
@@ -284,13 +284,36 @@ def worker_init_fn(worker_id):
                 self.data_parallel_recompute = self.config[
                     "DistributedStrategy"].get("recompute", None) is not None
 
+        self.enabled_ema = True if "EMA" in self.config else False
+        if self.enabled_ema and self.mode == 'train':
+            ema_cfg = self.config.get("EMA", {})
+            self.ema_eval = ema_cfg.pop('ema_eval', False)
+            self.ema_eval_start_epoch = ema_cfg.pop('eval_start_epoch', 0)
+            if self.ema_eval:
+                logger.info(
+                    f'You have enable ema evaluation and start from {self.ema_eval_start_epoch} epoch, and it will save the best ema state.'
+                )
+            else:
+                logger.info(
+                    f'You have enable ema, and also can set ema_eval=True and eval_start_epoch to enable ema evaluation.'
+                )
+            self.ema = EMA(self.optimizer._param_groups, **ema_cfg)
+            self.ema.register()
+
     def train(self):
         assert self.mode == "train"
         self.best_metric = {
             "metric": 0.0,
             "epoch": 0,
             "global_step": 0,
         }
+
+        if self.enabled_ema and self.ema_eval:
+            self.ema_best_metric = {
+                "metric": 0.0,
+                "epoch": 0,
+                "global_step": 0,
+            }
         # key:
         # val: metrics list word
         self.output_info = dict()
@@ -301,6 +324,12 @@ def train(self):
 
         # load checkpoint and resume
         if self.config["Global"]["checkpoint"] is not None:
+            if self.enabled_ema:
+                ema_metric_info = io.load_ema_checkpoint(
+                    self.config["Global"]["checkpoint"] + '_ema', self.ema)
+                if ema_metric_info is not None and self.ema_eval:
+                    self.ema_best_metric.update(ema_metric_info)
+
             metric_info = io.load_checkpoint(
                 self.config["Global"]["checkpoint"], self.model,
                 self.optimizer, self.scaler)
@@ -362,19 +391,36 @@ def train(self):
                     step=epoch_id,
                     writer=self.vdl_writer)
 
+                if self.enabled_ema and self.ema_eval and epoch_id > self.ema_eval_start_epoch:
+                    self.ema.apply_shadow()
+                    ema_eval_metric_info = self.eval(epoch_id)
+
+                    if ema_eval_metric_info["metric"] > self.ema_best_metric[
+                            "metric"]:
+                        self.ema_best_metric = ema_eval_metric_info.copy()
+                        io.save_ema_checkpoint(
+                            self.model,
+                            self.ema,
+                            self.output_dir,
+                            self.ema_best_metric,
+                            model_name=self.config["Model"]["name"],
+                            prefix="best_model_ema",
+                            max_num_checkpoint=self.config["Global"][
+                                "max_num_latest_checkpoint"], )
+
+                    logger.info("[Eval][Epoch {}][ema best metric: {}]".format(
+                        epoch_id, self.ema_best_metric["metric"]))
+                    logger.scaler(
+                        name="ema_eval_metric",
+                        value=eval_metric_info["metric"],
+                        step=epoch_id,
+                        writer=self.vdl_writer)
+
+                    self.ema.restore()
+
             # save model
-            if epoch_id % self.save_interval == 0:
-                if self.config["Global"]["max_num_latest_checkpoint"] != 0:
-                    io.save_checkpoint(
-                        self.model,
-                        self.optimizer,
-                        self.scaler,
-                        eval_metric_info,
-                        self.output_dir,
-                        model_name=self.config["Model"]["name"],
-                        prefix="epoch_{}".format(epoch_id),
-                        max_num_checkpoint=self.config["Global"][
-                            "max_num_latest_checkpoint"], )
+            if epoch_id % self.save_interval == 0 or epoch_id == self.config[
+                    "Global"]["epochs"]:
                 # save the latest model
                 io.save_checkpoint(
                     self.model,
@@ -387,6 +433,46 @@ def train(self):
                     max_num_checkpoint=self.config["Global"][
                         "max_num_latest_checkpoint"], )
 
+                if self.config["Global"]["max_num_latest_checkpoint"] != 0:
+                    io.save_checkpoint(
+                        self.model,
+                        self.optimizer,
+                        self.scaler,
+                        eval_metric_info,
+                        self.output_dir,
+                        model_name=self.config["Model"]["name"],
+                        prefix="epoch_{}".format(epoch_id),
+                        max_num_checkpoint=self.config["Global"][
+                            "max_num_latest_checkpoint"], )
+
+                if self.enabled_ema:
+                    if epoch_id == self.config["Global"]["epochs"]:
+                        self.ema.apply_shadow()
+
+                    io.save_ema_checkpoint(
+                        self.model,
+                        self.ema,
+                        self.output_dir,
+                        None,
+                        model_name=self.config["Model"]["name"],
+                        prefix="latest_ema",
+                        max_num_checkpoint=self.config["Global"][
+                            "max_num_latest_checkpoint"], )
+
+                    if self.config["Global"]["max_num_latest_checkpoint"] != 0:
+                        io.save_ema_checkpoint(
+                            self.model,
+                            self.ema,
+                            self.output_dir,
+                            None,
+                            model_name=self.config["Model"]["name"],
+                            prefix="epoch_{}_ema".format(epoch_id),
+                            max_num_checkpoint=self.config["Global"][
+                                "max_num_latest_checkpoint"], )
+
+                    if epoch_id == self.config["Global"]["epochs"]:
+                        self.ema.restore()
+
         if self.vdl_writer is not None:
             self.vdl_writer.close()
 

plsc/models/__init__.py

@@ -25,6 +25,7 @@
 from .mae import *
 from .convmae import *
 from .swin_transformer import *
+from .convnext import *
 
 __all__ = ["build_model"]