reformat: typehint and black

Author: bsridatta

.gitignore

@@ -11,5 +11,5 @@ __pycache__/
 # logs
 wandb/
 
-
+rotated_ship_data.py
 

setup.py

@@ -6,5 +6,5 @@
     author="Sri Datta Budaraju",
     author_email="b.sridatta@gmail.com",
     packages=setuptools.find_packages(),
-    python_requires='>=3.6'
-)
+    python_requires=">=3.6",
+)

src/callbacks/__init__.py

@@ -1,10 +1,10 @@
-from src.callbacks.base import CallbackList, Callback
-from src.callbacks.model_checkpoint import ModelCheckpoint
-from src.callbacks.logging import Logging
+from .base import CallbackList, Callback
+from .model_checkpoint import ModelCheckpoint
+from .logging import Logging
 
 
 __all__ = [
-    'CallbackList',
-    'ModelCheckpoint',
-    'Logging',
-]
+    "CallbackList",
+    "ModelCheckpoint",
+    "Logging",
+]

src/callbacks/base.py

@@ -1,7 +1,8 @@
-'''
+"""
+
 Callback inspritations from PyTorch Lightning - https://github.com/PyTorchLightning/PyTorch-Lightning
 and https://github.com/devforfu/pytorch_playground/blob/master/loop.ipynb
-'''
+"""
 
 import abc
 
@@ -73,7 +74,6 @@ def on_test_end(self, **kwargs):
 
 
 class CallbackList(Callback):
-
     def __init__(self, callbacks):
         self.callbacks = callbacks
 

src/callbacks/logging.py

@@ -1,52 +1,66 @@
-from src.callbacks.base import Callback
+from .base import Callback
 import torch
 
 
 class Logging(Callback):
     """Logging and printing metrics"""
 
     def setup(self, opt, model, **kwargs):
-        print(
-            f'[INFO]: Start training procedure using device: {opt.device}')
+        print(f"[INFO]: Start training procedure using device: {opt.device}")
         # log gradients and parameters of the model during training
         if opt.use_wandb:
-            opt.logger.watch(model, log='all')
+            opt.logger.watch(model, log="all")
 
-    def on_train_batch_end(self, opt, batch_idx, batch, dataloader, output, l_ship, l_bbox, **kwargs):
-        batch_len = len(batch['input'])
+    def on_train_batch_end(
+        self, opt, batch_idx, batch, dataloader, output, l_ship, l_bbox, **kwargs
+    ):
+        batch_len = len(batch["input"])
         dataset_len = len(dataloader.dataset)
         n_batches = len(dataloader)
 
         # print to console
-        print('Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.4f}\tL_ship: {:.4f}\tL_bbox: {:.4f}'.format(
-            opt.epoch, batch_idx * batch_len,
-            dataset_len, 100. * batch_idx / n_batches,
-            output, l_ship, l_bbox),
-            end='\n')
+        print(
+            "Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.4f}\tL_ship: {:.4f}\tL_bbox: {:.4f}".format(
+                opt.epoch,
+                batch_idx * batch_len,
+                dataset_len,
+                100.0 * batch_idx / n_batches,
+                output,
+                l_ship,
+                l_bbox,
+            ),
+            end="\n",
+        )
 
         # log to wandb
         if opt.use_wandb:
-            opt.logger.log({"train_loss": output,
-                            "l_ship": l_ship,
-                            "l_bbox": l_bbox})
+            opt.logger.log({"train_loss": output, "l_ship": l_ship, "l_bbox": l_bbox})
 
     def on_validation_end(self, opt, output, metrics, l_ship, l_bbox, **kwargs):
         # print and log metrics and loss after validation epoch
-        print("Valiation - Loss: {:.4f}\tL_ship: {:.4f}\tL_bbox: {:.4f}".format(
-            output, l_ship, l_bbox), end="\t")
+        print(
+            "Valiation - Loss: {:.4f}\tL_ship: {:.4f}\tL_bbox: {:.4f}".format(
+                output, l_ship, l_bbox
+            ),
+            end="\t",
+        )
 
         for k in metrics.keys():
             print(f"{k}: {metrics[k]}", end="\t")
             if opt.use_wandb:
                 opt.logger.log(metrics, commit=False)
-                opt.logger.log({"val_loss": output,
-                                "epoch": opt.epoch,
-                                "val_l_ship": l_ship,
-                                "val_l_bbox": l_bbox})
+                opt.logger.log(
+                    {
+                        "val_loss": output,
+                        "epoch": opt.epoch,
+                        "val_l_ship": l_ship,
+                        "val_l_bbox": l_bbox,
+                    }
+                )
         print("")
 
     def on_epoch_end(self, opt, optimizer, **kwargs):
-        lr = optimizer.param_groups[0]['lr']
+        lr = optimizer.param_groups[0]["lr"]
         if opt.use_wandb:
             opt.logger.log({f"LR": lr})
         print("lr @ ", lr)

src/callbacks/model_checkpoint.py

@@ -1,7 +1,9 @@
-from src.callbacks.base import Callback
-import torch
 import os
 
+import torch
+
+from callbacks.base import Callback
+
 
 class ModelCheckpoint(Callback):
     def __init__(self):
@@ -10,25 +12,27 @@ def __init__(self):
     def setup(self, opt, model, optimizer, **kwargs):
         # Save model code to wandb
         if opt.use_wandb:
-            opt.logger.save(
-                f"{os.path.dirname(os.path.abspath(__file__))}/models/*")
+            opt.logger.save(f"{os.path.dirname(os.path.abspath(__file__))}/models/*")
 
         # Resume training
         if opt.resume_run not in "None":
             state = torch.load(
-                f'{opt.save_dir}/{opt.resume_run}.pt', map_location=opt.device)
+                f"{opt.save_dir}/{opt.resume_run}.pt", map_location=opt.device
+            )
             print(
-                f'[INFO] Loaded Checkpoint {opt.resume_run}: @ epoch {state["epoch"]}')
-            model.load_state_dict(state['model_state_dict'])
+                f'[INFO] Loaded Checkpoint {opt.resume_run}: @ epoch {state["epoch"]}'
+            )
+            model.load_state_dict(state["model_state_dict"])
 
-        # Optimizers
+            # Optimizers
             optimizer_state_dic = torch.load(
-                f'{opt.save_dir}/{opt.resume_run}_optimizer.pt', map_location=opt.device)
+                f"{opt.save_dir}/{opt.resume_run}_optimizer.pt", map_location=opt.device
+            )
             optimizer.load_state_dict(optimizer_state_dic)
 
     def on_epoch_end(self, opt, val_loss, model, optimizer, epoch, **kwargs):
         # track val loss and save model when it decreases
-        if val_loss < self.val_loss_min and opt.device != 'cpu':
+        if val_loss < self.val_loss_min and opt.device != "cpu":
             self.val_loss_min = val_loss
 
             try:
@@ -37,28 +41,23 @@ def on_epoch_end(self, opt, val_loss, model, optimizer, epoch, **kwargs):
                 state_dict = model.state_dict()
 
             state = {
-                'epoch': epoch,
-                'val_loss': val_loss,
-                'model_state_dict': state_dict,
+                "epoch": epoch,
+                "val_loss": val_loss,
+                "model_state_dict": state_dict,
             }
 
             # model
-            torch.save(
-                state, f'{opt.save_dir}/{opt.run_name}.pt')
+            torch.save(state, f"{opt.save_dir}/{opt.run_name}.pt")
             if opt.use_wandb:
-                opt.logger.save(
-                    f'{opt.save_dir}/{opt.run_name}.pt')
-            print(
-                f'[INFO] Saved pt: {opt.save_dir}/{opt.run_name}.pt')
+                opt.logger.save(f"{opt.save_dir}/{opt.run_name}.pt")
+            print(f"[INFO] Saved pt: {opt.save_dir}/{opt.run_name}.pt")
 
             del state
 
             # Optimizer
             torch.save(
-                optimizer.state_dict(),
-                f'{opt.save_dir}/{opt.run_name}_optimizer.pt')
+                optimizer.state_dict(), f"{opt.save_dir}/{opt.run_name}_optimizer.pt"
+            )
             if opt.use_wandb:
-                opt.logger.save(
-                    f'{opt.save_dir}/{opt.run_name}_optimizer.pt')
-            print(
-                f'[INFO] Saved pt: {opt.save_dir}/{opt.run_name}_optimizer.pt')
+                opt.logger.save(f"{opt.save_dir}/{opt.run_name}_optimizer.pt")
+            print(f"[INFO] Saved pt: {opt.save_dir}/{opt.run_name}_optimizer.pt")

src/checkpoints/runX.pt

@@ -1,47 +1,56 @@
+from argparse import Namespace
+
 from torch.utils.data import DataLoader
+
 from src.dataset import Ships
 
 
-def train_dataloader(opt):
+def train_dataloader(opt: Namespace) -> DataLoader:
     print("[INFO]: Train dataloader called")
     dataset = Ships(n_samples=opt.train_len)
     sampler = None
     shuffle = True
-    loader = DataLoader(dataset=dataset,
-                        batch_size=opt.batch_size,
-                        num_workers=opt.num_workers,
-                        pin_memory=opt.pin_memory,
-                        sampler=sampler,
-                        shuffle=shuffle)
+    loader = DataLoader(
+        dataset=dataset,
+        batch_size=opt.batch_size,
+        num_workers=opt.num_workers,
+        pin_memory=opt.pin_memory,
+        sampler=sampler,
+        shuffle=shuffle,
+    )
     print("samples - ", len(dataset))
     return loader
 
 
-def val_dataloader(opt):
+def val_dataloader(opt: Namespace) -> DataLoader:
     print("[INFO]: Validation dataloader called")
     dataset = Ships(n_samples=opt.val_len)
     sampler = None
     shuffle = True
-    loader = DataLoader(dataset=dataset,
-                        batch_size=opt.batch_size,
-                        num_workers=opt.num_workers,
-                        pin_memory=opt.pin_memory,
-                        sampler=sampler,
-                        shuffle=shuffle)
+    loader = DataLoader(
+        dataset=dataset,
+        batch_size=opt.batch_size,
+        num_workers=opt.num_workers,
+        pin_memory=opt.pin_memory,
+        sampler=sampler,
+        shuffle=shuffle,
+    )
     print("samples - ", len(dataset))
     return loader
 
 
-def test_dataloader(opt):
+def test_dataloader(opt: Namespace) -> DataLoader:
     print("[INFO]: Test dataloader called")
     dataset = Ships(n_samples=opt.test_len)
     sampler = None
     shuffle = True
-    loader = DataLoader(dataset=dataset,
-                        batch_size=opt.batch_size,
-                        num_workers=opt.num_workers,
-                        pin_memory=opt.pin_memory,
-                        sampler=sampler,
-                        shuffle=shuffle)
+    loader = DataLoader(
+        dataset=dataset,
+        batch_size=opt.batch_size,
+        num_workers=opt.num_workers,
+        pin_memory=opt.pin_memory,
+        sampler=sampler,
+        shuffle=shuffle,
+    )
     print("samples - ", len(dataset))
     return loader

src/checkpoints/runX_optimizer.pt

@@ -1,9 +1,13 @@
-from torch.utils.data import Dataset
-from src.rotated_ship_data import make_data
-import torch
+from typing import Dict
+
 import numpy as np
+import torch
+from torch.functional import Tensor
+from torch.utils.data import Dataset
 from tqdm import tqdm
 
+from src.rotated_ship_data import make_data
+
 
 class Ships(Dataset):
     """ship datasets with has ship labels
@@ -15,7 +19,7 @@ class Ships(Dataset):
         sample {Tenosr} -- p_ship, x, y, yaw, h, w
     """
 
-    def __init__(self, n_samples=1000, pre_load=False):
+    def __init__(self, n_samples: int = 1000, pre_load: bool = False):
         self.n_samples = n_samples
         self.pre_load = pre_load
         if pre_load:
@@ -32,7 +36,7 @@ def __init__(self, n_samples=1000, pre_load=False):
     def __len__(self):
         return self.n_samples
 
-    def __getitem__(self, idx):
+    def __getitem__(self, idx: int) -> Dict[str, Tensor]:
         if self.pre_load:
             inp = self.inps[idx]
             target = self.targets[idx]
@@ -52,10 +56,11 @@ def __getitem__(self, idx):
 
         return sample
 
+
 # Used for simple experiment
 
 
-def make_batch(batch_size):
+def make_batch(batch_size: int):
     """Used only when pre_load = True
 
     Arguments: