update logger

Author: Ze-Yang

train.py

@@ -87,30 +87,31 @@ def __init__(self, model, model_old, device, opts, trainer_state=None, classes=N
         self.lkd_flag = self.lkd > 0. and model_old is not None
         self.kd_need_labels = False
         self.lgkd_flag = opts.lgkd
-        if opts.unkd:
-            self.lkd_loss = UnbiasedKnowledgeDistillationLoss(reduction="none", alpha=opts.alpha)
-        elif opts.lgkd:
-            self.lkd_loss = LabelGuidedKnowledgeDistillationLoss(alpha=opts.alpha,
-                                                                 prev_kd=opts.prev_kd,
-                                                                 novel_kd=opts.novel_kd)
-        elif opts.kd_bce_sig:
-            self.lkd_loss = BCESigmoid(reduction="none", alpha=opts.alpha, shape=opts.kd_bce_sig_shape)
-        elif opts.exkd_gt and self.old_classes > 0 and self.step > 0:
-            self.lkd_loss = ExcludedKnowledgeDistillationLoss(
-                reduction='none', index_new=self.old_classes, new_reduction="gt",
-                initial_nb_classes=opts.inital_nb_classes,
-                temperature_semiold=opts.temperature_semiold
-            )
-            self.kd_need_labels = True
-        elif opts.exkd_sum and self.old_classes > 0 and self.step > 0:
-            self.lkd_loss = ExcludedKnowledgeDistillationLoss(
-                reduction='none', index_new=self.old_classes, new_reduction="sum",
-                initial_nb_classes=opts.inital_nb_classes,
-                temperature_semiold=opts.temperature_semiold
-            )
-            self.kd_need_labels = True
-        else:
-            self.lkd_loss = KnowledgeDistillationLoss(alpha=opts.alpha)
+        if self.step > 0:
+            if opts.unkd:
+                self.lkd_loss = UnbiasedKnowledgeDistillationLoss(reduction="none", alpha=opts.alpha)
+            elif opts.lgkd:
+                self.lkd_loss = LabelGuidedKnowledgeDistillationLoss(alpha=opts.alpha,
+                                                                     prev_kd=opts.prev_kd,
+                                                                     novel_kd=opts.novel_kd)
+            elif opts.kd_bce_sig:
+                self.lkd_loss = BCESigmoid(reduction="none", alpha=opts.alpha, shape=opts.kd_bce_sig_shape)
+            elif opts.exkd_gt and self.old_classes > 0:
+                self.lkd_loss = ExcludedKnowledgeDistillationLoss(
+                    reduction='none', index_new=self.old_classes, new_reduction="gt",
+                    initial_nb_classes=opts.inital_nb_classes,
+                    temperature_semiold=opts.temperature_semiold
+                )
+                self.kd_need_labels = True
+            elif opts.exkd_sum and self.old_classes > 0:
+                self.lkd_loss = ExcludedKnowledgeDistillationLoss(
+                    reduction='none', index_new=self.old_classes, new_reduction="sum",
+                    initial_nb_classes=opts.inital_nb_classes,
+                    temperature_semiold=opts.temperature_semiold
+                )
+                self.kd_need_labels = True
+            else:
+                self.lkd_loss = KnowledgeDistillationLoss(alpha=opts.alpha)
 
         # ICARL
         self.icarl_combined = False

utils/logger.py

@@ -1,3 +1,4 @@
+import functools
 import logging
 import os
 import sys
@@ -92,24 +93,21 @@ def add_results(self, results, epoch=None):
             text += "</table>"
             self.writer.add_text(tag, text)
 
-    def setup_logger(self, output=None, distributed_rank=0, color=True, name="PLOP", abbrev_name=None):
+    @functools.lru_cache()  # so that calling setup_logger multiple times won't add many handlers
+    def setup_logger(self, output=None, distributed_rank=0, color=True, name="LGKD", abbrev_name=None):
         """
         Args:
             output (str): a file name or a directory to save log. If None, will not save log file.
                 If ends with ".txt" or ".log", assumed to be a file name.
                 Otherwise, logs will be saved to `output/log.txt`.
             name (str): the root module name of this logger
-            abbrev_name (str): an abbreviation of the module, to avoid long names in logs.
-                Set to "" to not log the root module in logs.
-                By default, will abbreviate "detectron2" to "d2" and leave other
-                modules unchanged.
         """
         logger = logging.getLogger(name)
         logger.setLevel(logging.DEBUG)
         logger.propagate = False
 
         if abbrev_name is None:
-            abbrev_name = "plop" if name == "PLOP" else name
+            abbrev_name = name
 
         plain_formatter = logging.Formatter(
             "[%(asctime)s] %(name)s %(levelname)s: %(message)s", datefmt="%m/%d %H:%M:%S"
@@ -170,5 +168,6 @@ def formatMessage(self, record):
 
 # cache the opened file object, so that different calls to `setup_logger`
 # with the same file name can safely write to the same file.
+@functools.lru_cache(maxsize=None)
 def _cached_log_stream(filename):
     return open(filename, "a")

utils/loss.py

@@ -1,6 +1,9 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+import logging
+
+logger = logging.getLogger('LGKD.' + __name__)
 
 
 def get_loss(loss_type):
@@ -454,7 +457,7 @@ def __init__(self, reduction='mean', alpha=1., prev_kd=10, novel_kd=1):
         self.alpha = alpha
         self.prev_kd = prev_kd
         self.novel_kd = novel_kd
-        print("prev kd: {}\t novel kd: {}".format(self.prev_kd, self.novel_kd))
+        logger.info("prev kd: {}\t novel kd: {}".format(self.prev_kd, self.novel_kd))
 
     def forward(self, new_logits, old_logits, targets):
         targets = targets.clone()