Revert "refactor: mv all dataloaders to engine.dataloader_dict"

This reverts commit 284e2a6.

Author: TingquanGao

ppcls/data/__init__.py

@@ -187,37 +187,10 @@ def mix_collate_fn(batch):
             collate_fn=batch_collate_fn,
             worker_init_fn=init_fn)
 
-    total_samples = len(
-        data_loader.dataset) if not use_dali else data_loader.size
-    max_iter = len(data_loader) - 1 if platform.system() == "Windows" else len(
-        data_loader)
-    data_loader.max_iter = max_iter
-    data_loader.total_samples = total_samples
-
     logger.debug("build data_loader({}) success...".format(data_loader))
     return data_loader
 
 
-# TODO(gaotingquan): perf
-class DataIterator(object):
-    def __init__(self, dataloader, use_dali=False):
-        self.dataloader = dataloader
-        self.use_dali = use_dali
-        self.iterator = iter(dataloader)
-
-    def get_batch(self):
-        # fetch data batch from dataloader
-        try:
-            batch = next(self.iterator)
-        except Exception:
-            # NOTE: reset DALI dataloader manually
-            if self.use_dali:
-                self.dataloader.reset()
-            self.iterator = iter(self.dataloader)
-            batch = next(self.iterator)
-        return batch
-
-
 def build_dataloader(engine):
     if "class_num" in engine.config["Global"]:
         global_class_num = engine.config["Global"]["class_num"]
@@ -249,15 +222,12 @@ def build_dataloader(engine):
         iter_per_epoch = len(train_dataloader) - 1 if platform.system(
         ) == "Windows" else len(train_dataloader)
         if engine.config["Global"].get("iter_per_epoch", None):
-            # TODO(gaotingquan): iter_per_epoch should be set in Dataloader.Train, not Global
             # set max iteration per epoch mannualy, when training by iteration(s), such as XBM, FixMatch.
             iter_per_epoch = engine.config["Global"].get("iter_per_epoch")
         iter_per_epoch = iter_per_epoch // engine.update_freq * engine.update_freq
-        # engine.iter_per_epoch = iter_per_epoch
+        engine.iter_per_epoch = iter_per_epoch
         train_dataloader.iter_per_epoch = iter_per_epoch
         dataloader_dict["Train"] = train_dataloader
-        # TODO(gaotingquan): set the iterator field in config, such as Dataloader.Train.convert_iterator=True
-        dataloader_dict["TrainIter"] = DataIterator(train_dataloader, use_dali)
 
     if engine.config["DataLoader"].get('UnLabelTrain', None) is not None:
         dataloader_dict["UnLabelTrain"] = build(
@@ -279,4 +249,5 @@ def build_dataloader(engine):
                     engine.config["DataLoader"]["Eval"], "Gallery", use_dali)
                 dataloader_dict["Query"] = build(
                     engine.config["DataLoader"]["Eval"], "Query", use_dali)
+
     return dataloader_dict

ppcls/engine/engine.py

@@ -63,7 +63,7 @@ def __init__(self, config, mode="train"):
 
         # init train_func and eval_func
         self.train_epoch_func = build_train_epoch_func(self.config)
-        self.eval_func = build_eval_func(self.config)
+        self.eval_epoch_func = build_eval_func(self.config)
 
         # set device
         self._init_device()
@@ -73,6 +73,12 @@ def __init__(self, config, mode="train"):
 
         # build dataloader
         self.dataloader_dict = build_dataloader(self)
+        self.train_dataloader, self.unlabel_train_dataloader, self.eval_dataloader = self.dataloader_dict[
+            "Train"], self.dataloader_dict[
+                "UnLabelTrain"], self.dataloader_dict["Eval"]
+        self.gallery_query_dataloader, self.gallery_dataloader, self.query_dataloader = self.dataloader_dict[
+            "GalleryQuery"], self.dataloader_dict[
+                "Gallery"], self.dataloader_dict["Query"]
 
         # build loss
         self.train_loss_func, self.unlabel_train_loss_func, self.eval_loss_func = build_loss(
@@ -88,7 +94,9 @@ def __init__(self, config, mode="train"):
         self._init_pretrained()
 
         # build optimizer
-        self.optimizer, self.lr_sch = build_optimizer(self)
+        self.optimizer, self.lr_sch = build_optimizer(
+            self.config, self.train_dataloader,
+            [self.model, self.train_loss_func])
 
         # AMP training and evaluating
         self._init_amp()

ppcls/engine/evaluation/classification.py

@@ -35,10 +35,13 @@ def classification_eval(engine, epoch_id=0):
     print_batch_step = engine.config["Global"]["print_batch_step"]
 
     tic = time.time()
-    total_samples = engine.dataloader_dict["Eval"].total_samples
     accum_samples = 0
-    max_iter = engine.dataloader_dict["Eval"].max_iter
-    for iter_id, batch in enumerate(engine.dataloader_dict["Eval"]):
+    total_samples = len(
+        engine.eval_dataloader.
+        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
         if iter_id == 5:
@@ -58,9 +61,9 @@ def classification_eval(engine, epoch_id=0):
                         "flatten_contiguous_range", "greater_than"
                     },
                     level=engine.amp_level):
-                out = engine.model(batch)
+                out = engine.model(batch[0])
         else:
-            out = engine.model(batch)
+            out = engine.model(batch[0])
 
         # just for DistributedBatchSampler issue: repeat sampling
         current_samples = batch_size * paddle.distributed.get_world_size()
@@ -92,8 +95,7 @@ def classification_eval(engine, epoch_id=0):
                 paddle.distributed.all_gather(pred_list, out)
                 preds = paddle.concat(pred_list, 0)
 
-            if accum_samples > total_samples and not engine.config[
-                    "Global"].get("use_dali", False):
+            if accum_samples > total_samples and not engine.use_dali:
                 if isinstance(preds, list):
                     preds = [
                         pred[:total_samples + current_samples - accum_samples]
@@ -149,11 +151,12 @@ def classification_eval(engine, epoch_id=0):
                 ])
                 metric_msg += ", {}".format(engine.eval_metric_func.avg_info)
             logger.info("[Eval][Epoch {}][Iter: {}/{}]{}, {}, {}".format(
-                epoch_id, iter_id, max_iter, metric_msg, time_msg, ips_msg))
+                epoch_id, iter_id,
+                len(engine.eval_dataloader), metric_msg, time_msg, ips_msg))
 
         tic = time.time()
-    if engine.config["Global"].get("use_dali", False):
-        engine.dataloader_dict["Eval"].reset()
+    if engine.use_dali:
+        engine.eval_dataloader.reset()
 
     if "ATTRMetric" in engine.config["Metric"]["Eval"][0]:
         metric_msg = ", ".join([

ppcls/engine/train/regular_train_epoch.py

@@ -22,8 +22,19 @@
 def regular_train_epoch(engine, epoch_id, print_batch_step):
     tic = time.time()
 
-    for iter_id in range(engine.dataloader_dict["Train"].iter_per_epoch):
-        batch = engine.dataloader_dict["TrainIter"].get_batch()
+    if not hasattr(engine, "train_dataloader_iter"):
+        engine.train_dataloader_iter = iter(engine.train_dataloader)
+
+    for iter_id in range(engine.iter_per_epoch):
+        # fetch data batch from dataloader
+        try:
+            batch = next(engine.train_dataloader_iter)
+        except Exception:
+            # NOTE: reset DALI dataloader manually
+            if engine.use_dali:
+                engine.train_dataloader.reset()
+            engine.train_dataloader_iter = iter(engine.train_dataloader)
+            batch = next(engine.train_dataloader_iter)
 
         profiler.add_profiler_step(engine.config["profiler_options"])
         if iter_id == 5:

ppcls/engine/train/utils.py

@@ -54,14 +54,13 @@ def log_info(trainer, batch_size, epoch_id, iter_id):
     ips_msg = "ips: {:.5f} samples/s".format(
         batch_size / trainer.time_info["batch_cost"].avg)
 
-    eta_sec = ((trainer.config["Global"]["epochs"] - epoch_id + 1
-                ) * trainer.dataloader_dict["Train"].iter_per_epoch - iter_id
-               ) * trainer.time_info["batch_cost"].avg
+    eta_sec = (
+        (trainer.config["Global"]["epochs"] - epoch_id + 1) *
+        trainer.iter_per_epoch - iter_id) * trainer.time_info["batch_cost"].avg
     eta_msg = "eta: {:s}".format(str(datetime.timedelta(seconds=int(eta_sec))))
     logger.info("[Train][Epoch {}/{}][Iter: {}/{}]{}, {}, {}, {}, {}".format(
-        epoch_id, trainer.config["Global"][
-            "epochs"], iter_id, trainer.dataloader_dict["Train"]
-        .iter_per_epoch, lr_msg, metric_msg, time_msg, ips_msg, eta_msg))
+        epoch_id, trainer.config["Global"]["epochs"], iter_id, trainer.
+        iter_per_epoch, lr_msg, metric_msg, time_msg, ips_msg, eta_msg))
 
     for i, lr in enumerate(trainer.lr_sch):
         logger.scaler(

ppcls/metric/__init__.py

@@ -70,9 +70,8 @@ def build_metrics(engine):
     if mode == 'train' and "Metric" in config and "Train" in config[
             "Metric"] and config["Metric"]["Train"]:
         metric_config = config["Metric"]["Train"]
-        if hasattr(engine.dataloader_dict["Train"],
-                   "collate_fn") and engine.dataloader_dict[
-                       "Train"].collate_fn is not None:
+        if hasattr(engine.train_dataloader, "collate_fn"
+                   ) and engine.train_dataloader.collate_fn is not None:
             for m_idx, m in enumerate(metric_config):
                 if "TopkAcc" in m:
                     msg = f"Unable to calculate accuracy when using \"batch_transform_ops\". The metric \"{m}\" has been removed."

ppcls/optimizer/__init__.py

@@ -45,15 +45,11 @@ def build_lr_scheduler(lr_config, epochs, step_each_epoch):
 
 
 # model_list is None in static graph
-def build_optimizer(engine):
-    if engine.mode != "train":
-        return None, None
-    config, iter_per_epoch, model_list = engine.config, engine.dataloader_dict[
-        "Train"].iter_per_epoch, [engine.mode, engine.train_loss_func]
+def build_optimizer(config, dataloader, model_list=None):
     optim_config = copy.deepcopy(config["Optimizer"])
     epochs = config["Global"]["epochs"]
     update_freq = config["Global"].get("update_freq", 1)
-    step_each_epoch = iter_per_epoch // update_freq
+    step_each_epoch = dataloader.iter_per_epoch // update_freq
     if isinstance(optim_config, dict):
         # convert {'name': xxx, **optim_cfg} to [{name: {scope: xxx, **optim_cfg}}]
         optim_name = optim_config.pop("name")