From 75080a7287ecb04095e21caa98c4dd2d16e87cb4 Mon Sep 17 00:00:00 2001 From: Aditya Oke <47158509+oke-aditya@users.noreply.github.com> Date: Sun, 22 Nov 2020 17:34:40 +0530 Subject: [PATCH] Adds type annotations (#50) --- .../models/classification/cnn/engine.py | 23 +++++++++++-------- .../components/torchvision_backbones.py | 6 ++--- quickvision/models/detection/detr/engine.py | 13 +++++++---- .../models/detection/detr/model_factory.py | 4 ++-- .../models/detection/faster_rcnn/engine.py | 21 +++++++++-------- .../models/detection/retinanet/engine.py | 23 ++++++++++--------- 6 files changed, 49 insertions(+), 41 deletions(-) diff --git a/quickvision/models/classification/cnn/engine.py b/quickvision/models/classification/cnn/engine.py index 8bd3bca..5769582 100644 --- a/quickvision/models/classification/cnn/engine.py +++ b/quickvision/models/classification/cnn/engine.py @@ -4,6 +4,7 @@ # https://github.com/oke-aditya/pytorch_cnn_trainer import torch +from torch import nn from torch.cuda import amp from quickvision import utils from quickvision.metrics import accuracy @@ -15,7 +16,8 @@ "val_sanity_fit", "sanity_fit", ] -def train_step(model, train_loader, criterion, device, optimizer, +def train_step(model: nn.Module, train_loader, criterion, + device: str, optimizer, scheduler=None, num_batches: int = None, log_interval: int = 100, grad_penalty: bool = False, scaler=None,): @@ -141,7 +143,8 @@ def train_step(model, train_loader, criterion, device, optimizer, return metrics -def val_step(model, val_loader, criterion, device, num_batches=None, +def val_step(model: nn.Module, val_loader, criterion, + device: str, num_batches=None, log_interval: int = 100): """ @@ -214,8 +217,8 @@ def val_step(model, val_loader, criterion, device, num_batches=None, return metrics -def fit(model, epochs, train_loader, val_loader, criterion, - device, optimizer, scheduler=None, early_stopper=None, +def fit(model: nn.Module, epochs: int, train_loader, val_loader, criterion, + device: str, optimizer, scheduler=None, early_stopper=None, num_batches: int = None, log_interval: int = 100, grad_penalty: bool = False, fp16: bool = False, swa_start: int = None, swa_scheduler=None,): @@ -310,8 +313,8 @@ def fit(model, epochs, train_loader, val_loader, criterion, return history -def train_sanity_fit(model, train_loader, criterion, - device, num_batches: int = None, log_interval: int = 100, +def train_sanity_fit(model: nn.Module, train_loader, criterion, + device: str, num_batches: int = None, log_interval: int = 100, grad_penalty: bool = False, fp16: bool = False,): """ @@ -381,8 +384,8 @@ def train_sanity_fit(model, train_loader, criterion, return True -def val_sanity_fit(model, val_loader, - criterion, device, num_batches: int = None, +def val_sanity_fit(model: nn.Module, val_loader, + criterion, device: str, num_batches: int = None, log_interval: int = 100,): """ @@ -433,8 +436,8 @@ def val_sanity_fit(model, val_loader, return True -def sanity_fit(model, train_loader, val_loader, - criterion, device, num_batches: int = None, +def sanity_fit(model: nn.Module, train_loader, val_loader, + criterion, device: str, num_batches: int = None, log_interval: int = 100, grad_penalty: bool = False, fp16: bool = False,): diff --git a/quickvision/models/components/torchvision_backbones.py b/quickvision/models/components/torchvision_backbones.py index 6a9c255..c54dcc8 100644 --- a/quickvision/models/components/torchvision_backbones.py +++ b/quickvision/models/components/torchvision_backbones.py @@ -7,7 +7,7 @@ __all__ = ["create_torchvision_backbone"] -def _create_backbone_generic(model, out_channels: int): +def _create_backbone_generic(model: nn.Module, out_channels: int): """ Generic Backbone creater. It removes the last linear layer. Args: @@ -23,7 +23,7 @@ def _create_backbone_generic(model, out_channels: int): # Use this when you have Adaptive Pooling layer in End. # When Model.features is not applicable. -def _create_backbone_adaptive(model, out_channels: int = None): +def _create_backbone_adaptive(model: nn.Module, out_channels: int = None): """ Creates backbone by removing linear after Adaptive Pooling layer. Args: @@ -36,7 +36,7 @@ def _create_backbone_adaptive(model, out_channels: int = None): return _create_backbone_generic(model, out_channels=out_channels) -def _create_backbone_features(model, out_channels: int): +def _create_backbone_features(model: nn.Module, out_channels: int): """ Creates backbone from feature sequential block. Args: diff --git a/quickvision/models/detection/detr/engine.py b/quickvision/models/detection/detr/engine.py index 8b1bf17..5ffdddd 100644 --- a/quickvision/models/detection/detr/engine.py +++ b/quickvision/models/detection/detr/engine.py @@ -1,4 +1,5 @@ import torch +from torch import nn from torch.cuda import amp from quickvision import utils from tqdm import tqdm @@ -9,8 +10,9 @@ "val_sanity_fit", "sanity_fit", ] -def train_step(model, train_loader, criterion, device, optimizer, scheduler=None, - num_batches: int = None, log_interval: int = 100, scaler=None,): +def train_step(model: nn.Module, train_loader, criterion, device: str, + optimizer, scheduler=None, num_batches: int = None, + log_interval: int = 100, scaler=None,): """ Performs one step of training. Calculates loss, forward pass, computes gradient and returns metrics. Args: @@ -105,7 +107,7 @@ def train_step(model, train_loader, criterion, device, optimizer, scheduler=None return metrics -def val_step(model, val_loader, criterion, device, +def val_step(model: nn.Module, val_loader, criterion, device, num_batches: int = None, log_interval: int = 100): """ Performs one step of validation. Calculates loss, forward pass and returns metrics. @@ -177,8 +179,9 @@ def val_step(model, val_loader, criterion, device, return metrics -def fit(model, epochs, train_loader, val_loader, criterion, - device, optimizer, scheduler=None, +def fit(model: nn.Module, epochs: int, train_loader, + val_loader, criterion, + device: str, optimizer, scheduler=None, num_batches: int = None, log_interval: int = 100, fp16: bool = False, ): diff --git a/quickvision/models/detection/detr/model_factory.py b/quickvision/models/detection/detr/model_factory.py index 8457dce..b677d35 100644 --- a/quickvision/models/detection/detr/model_factory.py +++ b/quickvision/models/detection/detr/model_factory.py @@ -14,7 +14,7 @@ class vision_detr(nn.Module): num_queries: Number of queries for transformer in Detr. backbone: Backbone created from create_detr_backbone. """ - def __init__(self, num_classes, num_queries, backbone): + def __init__(self, num_classes: int, num_queries: int, backbone: str): super().__init__() self.num_classes = num_classes self.num_queries = num_queries @@ -29,7 +29,7 @@ def forward(self, images): return self.model(images) -def create_vision_detr(num_classes: int, num_queries: int, backbone): +def create_vision_detr(num_classes: int, num_queries: int, backbone: str): """ Creates Detr Model for Object Detection Args: diff --git a/quickvision/models/detection/faster_rcnn/engine.py b/quickvision/models/detection/faster_rcnn/engine.py index b87a588..1ecf43f 100644 --- a/quickvision/models/detection/faster_rcnn/engine.py +++ b/quickvision/models/detection/faster_rcnn/engine.py @@ -1,4 +1,5 @@ import torch +from torch import nn, Tensor from torch.cuda import amp from quickvision import utils from tqdm import tqdm @@ -11,7 +12,7 @@ "val_sanity_fit", "sanity_fit", ] -def train_step(model, train_loader, device, optimizer, +def train_step(model: nn.Module, train_loader, device: str, optimizer, scheduler=None, num_batches: int = None, log_interval: int = 100, scaler=None,): @@ -111,7 +112,7 @@ def train_step(model, train_loader, device, optimizer, return metrics -def val_step(model, val_loader, device, num_batches=None, +def val_step(model: nn.Module, val_loader, device: str, num_batches: int = None, log_interval: int = 100): """ @@ -171,8 +172,8 @@ def val_step(model, val_loader, device, num_batches=None, return metrics -def fit(model, epochs, train_loader, val_loader, - device, optimizer, scheduler=None, +def fit(model: nn.Module, epochs: int, train_loader, val_loader, + device: str, optimizer, scheduler=None, num_batches: int = None, log_interval: int = 100, fp16: bool = False, ): @@ -224,8 +225,8 @@ def fit(model, epochs, train_loader, val_loader, return history -def train_sanity_fit(model, train_loader, - device, num_batches: int = None, log_interval: int = 100, +def train_sanity_fit(model: nn.Module, train_loader, + device: str, num_batches: int = None, log_interval: int = 100, fp16: bool = False,): """ @@ -245,8 +246,8 @@ def train_sanity_fit(model, train_loader, pass -def val_sanity_fit(model, val_loader, - device, num_batches: int = None, +def val_sanity_fit(model: nn.Module, val_loader, + device: str, num_batches: int = None, log_interval: int = 100,): """ @@ -265,8 +266,8 @@ def val_sanity_fit(model, val_loader, pass -def sanity_fit(model, train_loader, val_loader, - device, num_batches: int = None, +def sanity_fit(model: nn.Module, train_loader, val_loader, + device: str, num_batches: int = None, log_interval: int = 100, fp16: bool = False,): """ diff --git a/quickvision/models/detection/retinanet/engine.py b/quickvision/models/detection/retinanet/engine.py index c5ac479..0e1bd8c 100644 --- a/quickvision/models/detection/retinanet/engine.py +++ b/quickvision/models/detection/retinanet/engine.py @@ -1,4 +1,5 @@ import torch +from torch import nn from torch.cuda import amp from quickvision import utils from tqdm import tqdm @@ -10,7 +11,7 @@ "val_sanity_fit", "sanity_fit", ] -def train_step(model, train_loader, device, optimizer, +def train_step(model: nn.Module, train_loader, device: str, optimizer, scheduler=None, num_batches: int = None, log_interval: int = 100, scaler=None,): @@ -102,8 +103,8 @@ def train_step(model, train_loader, device, optimizer, return metrics -def val_step(model, val_loader, device, num_batches=None, - log_interval: int = 100): +def val_step(model: nn.Module, val_loader, device: str, + num_batches=None, log_interval: int = 100): """ Performs one step of validation. Calculates loss, forward pass and returns metrics. @@ -162,8 +163,8 @@ def val_step(model, val_loader, device, num_batches=None, return metrics -def fit(model, epochs, train_loader, val_loader, - device, optimizer, scheduler=None, +def fit(model: nn.Module, epochs: int, train_loader, val_loader, + device: str, optimizer, scheduler=None, num_batches: int = None, log_interval: int = 100, fp16: bool = False, ): @@ -215,19 +216,19 @@ def fit(model, epochs, train_loader, val_loader, return history -def train_sanity_fit(model, train_loader, - device, num_batches: int = None, log_interval: int = 100, +def train_sanity_fit(model: nn.Module, train_loader, + device: str, num_batches: int = None, log_interval: int = 100, fp16: bool = False,): pass -def val_sanity_fit(model, val_loader, - device, num_batches: int = None, +def val_sanity_fit(model: nn.Module, val_loader, + device: str, num_batches: int = None, log_interval: int = 100,): pass -def sanity_fit(model, train_loader, val_loader, - device, num_batches: int = None, +def sanity_fit(model: nn.Module, train_loader, val_loader, + device: str, num_batches: int = None, log_interval: int = 100, fp16: bool = False,): pass