重构训练格式，yaml->py

Author: lipengbo

Config/config.py

@@ -0,0 +1,24 @@
+class Config(object):
+    # ======数据集============
+    size= [224,224] # 图像尺寸
+    sampler="batch_balance" # 采样策略 
+    txt="./Config/dataset.txt" # 数据集路径
+
+    # ======模型============
+    optimizer="sgd" # 优化器
+
+    # 常规分类
+    backbone="mynet" # 主干网络 
+    loss="cross_entropy"  # 损失函数
+
+    # 度量学习
+    # backbone="mynet_metric" # 主干网络  
+    # loss="arcface"  # 损失函数 
+
+    # ======训练============
+    lr=0.01
+    batch=64
+    epochs=80 
+    scheduler="cosine"  # 学习率调度器
+
+cfg=Config()

Config/train.yaml

@@ -4,8 +4,8 @@ numpy
 pyyaml
 seaborn
 matplotlib
-torch
-torchvision
+torch>=1.11.0
+torchvision>=0.12.0
 tensorboard>1.15.0
 torchinfo #模型统计
 timm      #模型库

Package/requirements.txt

@@ -26,6 +26,7 @@
 
 cur_path = os.path.abspath(os.path.dirname(__file__))
 
+
 def analysis_dataset(txt):
     """
     解析dataset.txt 
@@ -37,23 +38,34 @@ def analysis_dataset(txt):
         "val": {"imgs": [], "labels": []},
         "test": {"imgs": [], "labels": []},
     }
-    labels=set()
+    labels = set()
     for path, label, mode in imgs_list:
         assert mode in ["train", "val", "test"]
         labels.add(label)
         dataset[mode]["imgs"].append(path)
         dataset[mode]["labels"].append(label)
-    labels=list(labels)
+    labels = list(labels)
     labels.sort()
 
     index = list(range(0, len(labels)))
-    labels_dict= dict(zip(index,labels))
-    
-    dataset["labels"]=labels
-    dataset["labels_dict"]=labels_dict
+    labels_dict = dict(zip(index, labels))
+
+    dataset["labels"] = labels
+    dataset["labels_dict"] = labels_dict
 
     return dataset
 
+def object2dict(object):
+    '''
+    类对象->字典
+    '''
+    dict={}
+    
+    for key in dir(object):
+        if not key.startswith('__'):
+            dict[key]=getattr(object, key)
+    return dict
+
 def init_env(cfg):
     """
     初始化训练环境
@@ -83,10 +95,10 @@ def init_env(cfg):
 
     # 初始化TensorBoard
     tb_writer = SummaryWriter(tb_path)
-    tb_writer.add_text("Config", str(cfg))
+    tb_writer.add_text("Config", str(object2dict(cfg)))
     print("*" * 28)
     print("TensorBoard | Checkpoint save to ", exp_path, "\n")
-    return tb_writer, checkpoint_path + cfg["Models"]["backbone"]
+    return tb_writer, checkpoint_path + cfg.backbone
 
 
 @torch.no_grad()

Utils/tools.py

@@ -2,7 +2,6 @@
 import os
 import torch
 import argparse
-import yaml
 import copy
 from DataSets.preprocess import PreProcess
 from DataSets import create_datasets, create_dataloader
@@ -19,71 +18,54 @@
 cur_path = os.path.abspath(os.path.dirname(__file__))
 
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Train")
-    parser.add_argument("--yaml", help="训练配置", default=cur_path + "/Config/train.yaml")
-    args = parser.parse_args()
+
+    # !导入训练配置!
+    from Config.config import cfg
 
     # 初始化环境
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    cfg = yaml.load(open(args.yaml, "r"), Loader=yaml.FullLoader)
-    labels_list = analysis_dataset(cfg["DataSet"]["txt"])["labels"]
+    labels_list = analysis_dataset(cfg.txt)["labels"]
 
     tb_writer, checkpoint_path = init_env(cfg)
 
     # 模型
-    model = create_backbone(cfg["Models"]["backbone"], num_classes=len(labels_list))
+    model = create_backbone(cfg.backbone, num_classes=len(labels_list))
     vis_model = copy.deepcopy(model)
     TASK = "metric" if hasattr(model, "embedding_size") else "class"
     # 区分任务
     if TASK == "metric":
         # 数据集
-        train_set = create_datasets(
-            txt=cfg["DataSet"]["txt"],
-            mode="train",
-            size=cfg["DataSet"]["size"],
-            use_augment=True,
-        )
-        train_dataloader = create_dataloader(
-            batch_size=cfg["DataSet"]["batch"],
-            dataset=train_set,
-            sampler_name=cfg["DataSet"]["sampler"],
-        )
-        val_set = create_datasets(
-            txt=cfg["DataSet"]["txt"], mode="val", size=cfg["DataSet"]["size"]
-        )
+        train_set = create_datasets(txt=cfg.txt, mode="train", size=cfg.size, use_augment=True)
+        train_dataloader = create_dataloader(batch_size=cfg.batch, dataset=train_set, sampler_name=cfg.sampler)
+        val_set = create_datasets(txt=cfg.txt, mode="val", size=cfg.size)
 
         # 难样例挖掘
         mining_func = miners.MultiSimilarityMiner()
 
         # 损失函数(分类器)
         loss_func = create_metric_loss(
-            name=cfg["Models"]["loss"],
+            name=cfg.loss,
             num_classes=len(labels_list),
             embedding_size=model.embedding_size,
         ).to(device)
-        params = [{"params": loss_func.parameters(), "lr": cfg["Train"]["lr"]}]
+        params = [{"params": loss_func.parameters(), "lr": cfg.lr}]
 
     else:
         # 数据集
-        train_set = create_datasets(
-            txt=cfg["DataSet"]["txt"],
-            mode="train",
-            size=cfg["DataSet"]["size"],
-            use_augment=True,
-        )
-        val_set = create_datasets(txt=cfg["DataSet"]["txt"], mode="val", size=cfg["DataSet"]["size"])
+        train_set = create_datasets(txt=cfg.txt, mode="train", size=cfg.size, use_augment=True,)
+        val_set = create_datasets(txt=cfg.txt, mode="val", size=cfg.size)
 
         # 数据集加载器
         train_dataloader = create_dataloader(
-            batch_size=cfg["DataSet"]["batch"],
+            batch_size=cfg.batch,
             dataset=train_set,
-            sampler_name=cfg["DataSet"]["sampler"],
+            sampler_name=cfg.sampler,
         )
 
-        val_dataloader = create_dataloader(batch_size=cfg["DataSet"]["batch"], dataset=val_set)
+        val_dataloader = create_dataloader(batch_size=cfg.batch, dataset=val_set)
 
         # 损失函数
-        loss_func = create_class_loss(cfg["Models"]["loss"]).to(device)
+        loss_func = create_class_loss(cfg.loss).to(device)
         params = []
 
     # 模型转为GPU
@@ -95,18 +77,18 @@
     # 优化器
     params.append({"params": model.parameters()})
     optimizer = create_optimizer(
-        params, cfg["Models"]["optimizer"], lr=cfg["Train"]["lr"]
+        params, cfg.optimizer, lr=cfg.lr
     )
 
     # 学习率调度器
     lr_scheduler = create_scheduler(
-        sched_name=cfg["Train"]["scheduler"],
-        epochs=cfg["Train"]["epochs"],
+        sched_name=cfg.scheduler,
+        epochs=cfg.epochs,
         optimizer=optimizer,
     )
     best_score = 0.0
-    for epoch in range(cfg["Train"]["epochs"]):
-        print("start epoch {}/{}...".format(epoch, cfg["Train"]["epochs"]))
+    for epoch in range(cfg.epochs):
+        print("start epoch {}/{}...".format(epoch, cfg.epochs))
         tb_writer.add_scalar("Train/lr", optimizer.param_groups[-1]["lr"], epoch)
         optimizer.zero_grad()
 
@@ -140,9 +122,7 @@
 
             ema_model.update(model)
 
-            lr_scheduler.step_update(
-                num_updates=epoch * len(train_dataloader) + batch_idx
-            )
+            lr_scheduler.step_update(num_updates=epoch * len(train_dataloader) + batch_idx)
 
             if batch_idx % 100 == 0:
                 iter_num = int(batch_idx + epoch * len(train_dataloader))
@@ -159,9 +139,7 @@
         elif TASK == "metric":  # 度量学习
             score = eval_metric_model(model, train_set, val_set)
             ema_score = eval_metric_model(ema_model.module, train_set, val_set)
-            tb_writer.add_scalars(
-                "Eval", {"precision": score, "ema_precision": ema_score}, epoch
-            )
+            tb_writer.add_scalars("Eval", {"precision": score, "ema_precision": ema_score}, epoch)
         model.train()
 
         # 保存最优模型