支持使用配置文件

.gitignore

@@ -6,6 +6,9 @@ __pycache__/
 *.bin
 *.ipynb
 !template.ipynb
+*.yaml
+!pre-commit-config.yaml
+!config-default.yaml
 
 # C extensions
 *.so

config-default.yaml

@@ -0,0 +1,18 @@
+num_workers: 4
+train_batch_size: 8
+valid_batch_size: 8
+max_node_count: 1024
+train_num_chunk: 10
+limit_valid_set_size: 10000
+
+# 模型相关
+word_embedding_size: 384
+hidden_size: 768
+num_layers: 6
+num_heads: 8
+model_lr: 3.0e-5
+classifier_lr: 1.0e-3
+
+# 验证
+valid_chunk_gap: 4
+evaluate_step_gap: 50

detecter/train/BigCloneBench.py

@@ -9,14 +9,16 @@
 
 
 class Trainer(torch.nn.Module):
-    def __init__(self, model: AstAttention, classifier: Classifier):
+    def __init__(self, model: AstAttention, classifier: Classifier, evaluate_step_gap: int):
         super().__init__()
         self.model: AstAttention = model
         self.classifier: Classifier = classifier
+        self.evaluate_step_gap = evaluate_step_gap
+
         self.loss_fn = torch.nn.CrossEntropyLoss()
 
-        self.evaluator = MulticlassF1Score(num_classes=2)
-        self.evaluator_all = MulticlassF1Score(num_classes=2)
+        self.half_evaluator = MulticlassF1Score(num_classes=2)
+        self.final_evaluator = MulticlassF1Score(num_classes=2)
         self.loss_list = []
 
     def device(self) -> bool:
@@ -30,29 +32,32 @@ def forward(self, batch: Union[torch.Tensor, torch.Tensor, torch.Tensor]):
         label = label.to(device)
         nodes = nodes.to(device)
         mask = mask.to(device)
-        self.evaluator.to(device)
-        self.evaluator_all.to(device)
+        self.half_evaluator.to(device)
+        self.final_evaluator.to(device)
 
         hidden = self.model(nodes, mask)[0]
         score = self.classifier(hidden)
         loss = self.loss_fn(score, label.long())
-        self.evaluator.update(score, label.long())
-        self.evaluator_all.update(score, label.long())
+        self.half_evaluator.update(score, label.long())
+        self.final_evaluator.update(score, label.long())
 
-        logger.debug("f1   {}".format(self.evaluator.compute().item()))
-        self.evaluator.reset()
-        logger.debug("loss {}".format(loss.item()))
         self.loss_list.append(loss.item())
 
+        if len(self.loss_list) % self.evaluate_step_gap == 0:
+            logger.debug("f1   {}".format(self.half_evaluator.compute().item()))
+            self.half_evaluator.reset()
+            logger.debug("loss {}".format(sum(self.loss_list[-self.evaluate_step_gap :]) / self.evaluate_step_gap))
+
         return loss
 
     def evaluate(self) -> float:
-        f1 = self.evaluator_all.compute()
-        self.evaluator_all.reset()
-        logger.info("aggr evaluate {}".format(f1))
+        logger.info("aggr evaluate {}".format(self.final_evaluator.compute().item()))
         loss = sum(self.loss_list) / len(self.loss_list)
         logger.info("aggr loss     {}".format(loss))
+
         self.loss_list = []
+        self.final_evaluator.reset()
+        self.half_evaluator.reset()
         return loss
 
 

trainBCB.py

@@ -2,6 +2,7 @@
 from typing import *
 
 import torch
+import yaml
 from torch.cuda.amp import GradScaler
 from torch.utils import data
 from tqdm import tqdm
@@ -28,50 +29,65 @@ def array_split(arr, n):
     stderr.setLevel(logging.INFO)
     logger.addHandler(stderr)
 
-    batch_size = 8
+    with open("config-default.yaml", "r") as f:
+        cfg = yaml.safe_load(f)
+    try:
+        with open("config.yaml", "r") as f:
+            new_cfg = yaml.safe_load(f)
+        if new_cfg:
+            cfg = {**cfg, **new_cfg}
+    except IOError:
+        pass
+
     ds = BigCloneBench.DataSet(
-        "dataset/BigCloneBench/data.jsonl.txt", "dataset/BigCloneBench/test.txt", max_node_count=1024, fixed_prune=False
+        "dataset/BigCloneBench/data.jsonl.txt",
+        "dataset/BigCloneBench/test.txt",
+        max_node_count=cfg["max_node_count"],
+        fixed_prune=False,
     )
 
-    indices_list = array_split(list(range(len(ds))), 10)
+    indices_list = array_split(list(range(len(ds))), cfg["train_num_chunk"])
     subdataset_list = [data.Subset(ds, indices) for indices in indices_list]
     loaders = [
         data.DataLoader(
             subdataset,
-            batch_size=batch_size,
+            batch_size=cfg["train_batch_size"],
             collate_fn=BigCloneBench.collate_fn,
-            num_workers=4,
+            num_workers=cfg["num_workers"],
             pin_memory=True,
         )
         for subdataset in subdataset_list
     ]
 
     ds = BigCloneBench.DataSet(
-        "dataset/BigCloneBench/data.jsonl.txt", "dataset/BigCloneBench/valid.txt", max_node_count=1024, fixed_prune=True
+        "dataset/BigCloneBench/data.jsonl.txt",
+        "dataset/BigCloneBench/valid.txt",
+        max_node_count=cfg["max_node_count"],
+        fixed_prune=True,
     )
-    ds = data.Subset(ds, list(range(10000)))
+    ds = data.Subset(ds, list(range(cfg["limit_valid_set_size"])))
     v_loader = data.DataLoader(
         ds,
-        batch_size=batch_size,
+        batch_size=cfg["valid_batch_size"],
         collate_fn=BigCloneBench.collate_fn,
-        num_workers=4,
+        num_workers=cfg["num_workers"],
         pin_memory=True,
     )
 
-    model = AstAttention(384, 768, num_layers=6, num_heads=8).cuda()
-    classifier = Classifier(768, 2).cuda()
-    trainer = Trainer(model=model, classifier=classifier).cuda()
+    model = AstAttention(cfg["word_embedding_size"], cfg["hidden_size"], cfg["num_layers"], cfg["num_heads"]).cuda()
+    classifier = Classifier(cfg["hidden_size"], 2).cuda()
+    trainer = Trainer(model=model, classifier=classifier, evaluate_step_gap=cfg["evaluate_step_gap"]).cuda()
 
     optimizer = torch.optim.AdamW(
         [
             {
                 "params": model.parameters(),
-                "lr": 3e-5,
+                "lr": cfg["model_lr"],
                 "weight_decay": 0.1,
             },
             {
                 "params": classifier.parameters(),
-                "lr": 1e-3,
+                "lr": cfg["classifier_lr"],
             },
         ]
     )
@@ -106,7 +122,7 @@ def array_split(arr, n):
     logger.info("epoch {}".format(epoch))
 
     while True:
-        if trained_chunks == len(loaders):
+        if trained_chunks >= len(loaders):
             epoch += 1
             trained_chunks = 0
             logger.info("===========================")
@@ -125,7 +141,7 @@ def array_split(arr, n):
 
         torch.save(check_point(trainer, optimizer, scaler, (epoch, trained_chunks)), TRAINER_CKPT_PATH)
 
-        if trained_chunks % 2 == 0:
+        if (trained_chunks + len(loaders) * epoch) % cfg["valid_chunk_gap"] == 0:
             with torch.inference_mode():
                 trainer.eval()
                 for batch in tqdm(v_loader, desc="VALID"):