Merge pull request #41 from kaylode/v1.0.1

Update to V1.0.1

Author: kaylode

requirements.txt

@@ -12,6 +12,8 @@ tabulate==0.8.10
 loguru==0.6.0
 seaborn==0.12.0
 wandb==0.13.3
+plotly==5.10.0
+matplotlib==3.4.3
 
 ## Classification
 timm

theseus/base/callbacks/checkpoint_callbacks.py

@@ -54,6 +54,7 @@ def save_checkpoint(self, trainer, iters, outname='last'):
         weights = {
             'model': trainer.model.model.state_dict(),
             'optimizer': trainer.optimizer.state_dict(),
+            'scheduler': trainer.scheduler.state_dict(),
             'iters': iters,
             'best_value': self.best_value,
         }

theseus/base/callbacks/wandb_callbacks.py

@@ -108,6 +108,7 @@ def on_start(self, logs: Dict=None):
         # Save all config files
         self.wandb_logger.log_file(
             tag='configs', 
+            base_folder=self.save_dir,
             value = osp.join(self.save_dir, '*.yaml'))
 
         # Init logging model for debug

theseus/base/datasets/dataset.py

@@ -18,6 +18,14 @@ def __init__(self, datasets: Iterable[data.Dataset], **kwargs) -> None:
         self.classnames = datasets[0].classnames
         self.collate_fn = datasets[0].collate_fn
 
+    def __getattr__(self, attr):
+        if hasattr(self, attr):
+            return getattr(self, attr)
+            
+        if hasattr(self.datasets[0], attr):
+            return getattr(self.datasets[0], attr)
+
+        raise AttributeError
 
 class ChainDataset(data.ConcatDataset):
     """

theseus/base/losses/multi_loss.py

@@ -11,15 +11,15 @@ def __init__(self, losses: Iterable[nn.Module], weights=None, **kwargs):
         self.losses = losses
         self.weights = [1.0 for _ in range(len(losses))] if weights is None else weights
 
-    def forward(self, pred: torch.Tensor, batch: Dict[str, Any], device: torch.device):
+    def forward(self, outputs: Dict[str, Any], batch: Dict[str, Any], device: torch.device):
         """
         Forward inputs and targets through multiple losses
         """
         total_loss = 0
         total_loss_dict = {}
 
         for weight, loss_fn in zip(self.weights, self.losses):
-          loss, loss_dict = loss_fn(pred, batch, device)
+          loss, loss_dict = loss_fn(outputs, batch, device)
           total_loss += (weight*loss)
           total_loss_dict.update(loss_dict)
 

theseus/classification/callbacks/__init__.py

@@ -1,4 +1,6 @@
 from theseus.classification.callbacks.visualize_callbacks import ClassificationVisualizerCallbacks
+from theseus.classification.callbacks.gradcam_callbacks import GradCAMVisualizationCallbacks 
 from theseus.base.callbacks import CALLBACKS_REGISTRY
 
-CALLBACKS_REGISTRY.register(ClassificationVisualizerCallbacks)
+CALLBACKS_REGISTRY.register(ClassificationVisualizerCallbacks)
+CALLBACKS_REGISTRY.register(GradCAMVisualizationCallbacks)

theseus/classification/callbacks/gradcam_callbacks.py

@@ -0,0 +1,97 @@
+from typing import Dict
+import matplotlib.pyplot as plt
+import torch
+from torchvision.transforms import functional as TFF
+
+from theseus.base.callbacks.base_callbacks import Callbacks
+from theseus.utilities.loggers.observer import LoggerObserver
+from theseus.classification.utilities.gradcam import CAMWrapper, show_cam_on_image
+from theseus.utilities.visualization.visualizer import Visualizer
+
+LOGGER = LoggerObserver.getLogger("main")
+
+class GradCAMVisualizationCallbacks(Callbacks):
+    """
+    Callbacks for visualizing stuff during training
+    Features:
+        - Visualize datasets; plot model architecture, analyze datasets in sanity check
+        - Visualize prediction at every end of validation
+
+    """
+
+    def __init__(self, **kwargs) -> None:
+        super().__init__()
+        self.visualizer = Visualizer()
+
+    @torch.enable_grad() #enable grad for CAM
+    def on_val_epoch_end(self, logs: Dict=None):
+        """
+        After finish validation
+        """
+
+        iters = logs['iters']
+        last_batch = logs['last_batch']
+        model = self.params['trainer'].model
+        valloader = self.params['trainer'].valloader
+        optimizer = self.params['trainer'].optimizer
+
+        # Zeroing gradients in model and optimizer for supress warning
+        optimizer.zero_grad()
+        model.zero_grad()
+
+        # Vizualize Grad Class Activation Mapping and model predictions
+        LOGGER.text("Visualizing model predictions...", level=LoggerObserver.DEBUG)
+
+        images = last_batch["inputs"]
+        targets = last_batch["targets"]
+        model.eval()
+        
+        ## Calculate GradCAM and Grad Class Activation Mapping and
+        model_name = model.model.name
+
+        try:
+            grad_cam = CAMWrapper.get_method(
+                name='gradcam', 
+                model=model.model.get_model(), 
+                model_name=model_name, use_cuda=next(model.parameters()).is_cuda)
+            
+            grayscale_cams, label_indices, scores = grad_cam(images, return_probs=True)
+
+        except:
+            LOGGER.text("Cannot calculate GradCAM", level=LoggerObserver.ERROR)
+            return
+            
+        gradcam_batch = []
+        for idx in range(len(grayscale_cams)):
+            image = images[idx]
+            target = targets[idx].item()
+            label = label_indices[idx]
+            grayscale_cam = grayscale_cams[idx, :]
+
+            img_show = self.visualizer.denormalize(image)
+            if valloader.dataset.classnames is not None:
+                label = valloader.dataset.classnames[label]
+                target = valloader.dataset.classnames[target]
+
+            img_cam =show_cam_on_image(img_show, grayscale_cam, use_rgb=True)
+
+            img_cam = TFF.to_tensor(img_cam)
+            gradcam_batch.append(img_cam)
+
+            if idx == 63: # limit number of images
+                break
+        
+        # GradCAM images
+        gradcam_grid_img = self.visualizer.make_grid(gradcam_batch)
+        fig = plt.figure(figsize=(8,8))
+        plt.imshow(gradcam_grid_img)
+        plt.axis("off")
+        plt.tight_layout(pad=0)
+        LOGGER.log([{
+            'tag': "Validation/gradcam",
+            'value': fig,
+            'type': LoggerObserver.FIGURE,
+            'kwargs': {
+                'step': iters
+            }
+        }])

theseus/classification/callbacks/visualize_callbacks.py

@@ -150,7 +150,7 @@ def analyze_gt(self, trainset, valset, iters):
         plt.clf()   # Clear figure
         plt.close()
 
-    @torch.enable_grad() #enable grad for CAM
+    @torch.no_grad() #enable grad for CAM
     def on_val_epoch_end(self, logs: Dict=None):
         """
         After finish validation
@@ -160,41 +160,24 @@ def on_val_epoch_end(self, logs: Dict=None):
         last_batch = logs['last_batch']
         model = self.params['trainer'].model
         valloader = self.params['trainer'].valloader
-        optimizer = self.params['trainer'].optimizer
 
-        # Zeroing gradients in model and optimizer for supress warning
-        optimizer.zero_grad()
-        model.zero_grad()
-
-        # Vizualize Grad Class Activation Mapping and model predictions
+        # Vizualize model predictions
         LOGGER.text("Visualizing model predictions...", level=LoggerObserver.DEBUG)
 
         images = last_batch["inputs"]
         targets = last_batch["targets"]
         model.eval()
 
-        ## Calculate GradCAM
-        model_name = model.model.name
-
-        try:
-            grad_cam = CAMWrapper.get_method(
-                name='gradcam', 
-                model=model.model.get_model(), 
-                model_name=model_name, use_cuda=next(model.parameters()).is_cuda)
-            
-            grayscale_cams, label_indices, scores = grad_cam(images, return_probs=True)
-
-        except:
-            LOGGER.text("Cannot calculate GradCAM", level=LoggerObserver.ERROR)
-            return
+        ## Get prediction on last batch
+        outputs = model.model.get_prediction(last_batch, device=model.device)
+        label_indices = outputs['labels']
+        scores = outputs['confidences']
 
-        gradcam_batch = []
         pred_batch = []
-        for idx in range(len(grayscale_cams)):
+        for idx in range(len(images)):
             image = images[idx]
             target = targets[idx].item()
             label = label_indices[idx]
-            grayscale_cam = grayscale_cams[idx, :]
             score = scores[idx]
 
             img_show = self.visualizer.denormalize(image)
@@ -217,32 +200,12 @@ def on_val_epoch_end(self, logs: Dict=None):
                 offset=100
             )
 
-            img_cam =show_cam_on_image(img_show, grayscale_cam, use_rgb=True)
-
-            img_cam = TFF.to_tensor(img_cam)
-            gradcam_batch.append(img_cam)
-
             pred_img = self.visualizer.get_image()
             pred_img = TFF.to_tensor(pred_img)
             pred_batch.append(pred_img)
 
             if idx == 63: # limit number of images
                 break
-        
-        # GradCAM images
-        gradcam_grid_img = self.visualizer.make_grid(gradcam_batch)
-        fig = plt.figure(figsize=(8,8))
-        plt.imshow(gradcam_grid_img)
-        plt.axis("off")
-        plt.tight_layout(pad=0)
-        LOGGER.log([{
-            'tag': "Validation/gradcam",
-            'value': fig,
-            'type': LoggerObserver.FIGURE,
-            'kwargs': {
-                'step': iters
-            }
-        }])
 
         # Prediction images
         pred_grid_img = self.visualizer.make_grid(pred_batch)
@@ -261,8 +224,4 @@ def on_val_epoch_end(self, logs: Dict=None):
 
         plt.cla()   # Clear axis
         plt.clf()   # Clear figure
-        plt.close()
-
-        # Zeroing gradients in model and optimizer for safety
-        optimizer.zero_grad()
-        model.zero_grad()
+        plt.close()

theseus/classification/metrics/__init__.py

@@ -6,10 +6,12 @@
 from .confusion_matrix import *
 from .errorcases import *
 from .projection import *
+from .precision_recall import *
 
 METRIC_REGISTRY.register(Accuracy)
 METRIC_REGISTRY.register(BalancedAccuracyMetric)
 METRIC_REGISTRY.register(F1ScoreMetric)
 METRIC_REGISTRY.register(ConfusionMatrix)
 METRIC_REGISTRY.register(ErrorCases)
-METRIC_REGISTRY.register(EmbeddingProjection)
+METRIC_REGISTRY.register(EmbeddingProjection)
+METRIC_REGISTRY.register(PrecisionRecall)

theseus/classification/metrics/precision_recall.py

@@ -0,0 +1,38 @@
+from sklearn.metrics import precision_score, recall_score
+from typing import Any, Dict
+
+from theseus.base.metrics.metric_template import Metric
+from theseus.classification.utilities.logits import logits2labels
+
+class PrecisionRecall(Metric):
+    """
+    F1 Score Metric (including macro, micro)
+    """
+    def __init__(self, average = 'weighted', label_type:str = 'multiclass', **kwargs):
+        super().__init__(**kwargs)
+        self.average = average
+        self.type =label_type
+        self.threshold = kwargs.get('threshold', 0.5)
+        self.reset()
+
+    def update(self, outputs: Dict[str, Any], batch: Dict[str, Any]):
+        """
+        Perform calculation based on prediction and targets
+        """
+        targets = batch["targets"] 
+        outputs = outputs["outputs"] 
+
+        outputs = logits2labels(outputs, label_type=self.type, threshold=self.threshold)
+        targets = targets.squeeze()
+    
+        self.preds +=  outputs.numpy().tolist()
+        self.targets +=  targets.numpy().tolist()
+
+    def value(self):
+        precision = precision_score(self.targets, self.preds, average=self.average, zero_division=1)
+        recall = recall_score(self.targets, self.preds, average=self.average, zero_division=1)
+        return {f"{self.average}-precision": precision, f"{self.average}-recall": recall}
+
+    def reset(self):
+        self.targets = []
+        self.preds = []