diff options
Diffstat (limited to 'models')
| -rw-r--r-- | models/hpm.py | 4 | ||||
| -rw-r--r-- | models/layers.py | 18 | ||||
| -rw-r--r-- | models/model.py | 10 |
3 files changed, 15 insertions, 17 deletions
diff --git a/models/hpm.py b/models/hpm.py index f387154..1773f56 100644 --- a/models/hpm.py +++ b/models/hpm.py @@ -1,5 +1,3 @@ -from typing import Tuple - import torch import torch.nn as nn from torchvision.models import resnet50 @@ -10,7 +8,7 @@ from models.layers import HorizontalPyramidPooling class HorizontalPyramidMatching(nn.Module): def __init__( self, - scales: Tuple[int] = (1, 2, 4, 8), + scales: tuple[int, ...] = (1, 2, 4, 8), out_channels: int = 256, use_avg_pool: bool = False, **kwargs diff --git a/models/layers.py b/models/layers.py index 9b17205..cba6e47 100644 --- a/models/layers.py +++ b/models/layers.py @@ -1,8 +1,8 @@ -from typing import Union, Tuple +from typing import Union import torch -import torch.nn.functional as F import torch.nn as nn +import torch.nn.functional as F class BasicConv2d(nn.Module): @@ -10,7 +10,7 @@ class BasicConv2d(nn.Module): self, in_channels: int, out_channels: int, - kernel_size: Union[int, Tuple[int, int]], + kernel_size: Union[int, tuple[int, int]], **kwargs ): super().__init__() @@ -29,7 +29,7 @@ class VGGConv2d(BasicConv2d): self, in_channels: int, out_channels: int, - kernel_size: Union[int, Tuple[int, int]] = 3, + kernel_size: Union[int, tuple[int, int]] = 3, padding: int = 1, **kwargs ): @@ -47,7 +47,7 @@ class BasicConvTranspose2d(nn.Module): self, in_channels: int, out_channels: int, - kernel_size: Union[int, Tuple[int, int]], + kernel_size: Union[int, tuple[int, int]], **kwargs ): super().__init__() @@ -66,7 +66,7 @@ class DCGANConvTranspose2d(BasicConvTranspose2d): self, in_channels: int, out_channels: int, - kernel_size: Union[int, Tuple[int, int]] = 4, + kernel_size: Union[int, tuple[int, int]] = 4, stride: int = 2, padding: int = 1, is_last_layer: bool = False, @@ -88,7 +88,7 @@ class FocalConv2d(BasicConv2d): self, in_channels: int, out_channels: int, - kernel_size: Union[int, Tuple[int, int]], + kernel_size: Union[int, tuple[int, int]], halving: int, **kwargs ): @@ -108,7 +108,7 @@ class BasicConv1d(nn.Module): self, in_channels: int, out_channels: int, - kernel_size: Union[int, Tuple[int]], + kernel_size: Union[int, tuple[int]], **kwargs ): super(BasicConv1d, self).__init__() @@ -124,7 +124,7 @@ class HorizontalPyramidPooling(BasicConv2d): self, in_channels: int, out_channels: int, - kernel_size: Union[int, Tuple[int, int]] = 1, + kernel_size: Union[int, tuple[int, int]] = 1, use_avg_pool: bool = False, **kwargs ): diff --git a/models/model.py b/models/model.py index 369d6c2..cb0e756 100644 --- a/models/model.py +++ b/models/model.py @@ -1,4 +1,4 @@ -from typing import List, Dict, Union, Tuple +from typing import Union import torch from torch.utils.data.dataloader import default_collate @@ -7,15 +7,15 @@ from torch.utils.data.dataloader import default_collate class Model: def __init__( self, - batch_size: Tuple[int, int] + batch_size: tuple[int, int] ): (self.pr, self.k) = batch_size def _batch_splitter( self, - batch: List[Dict[str, Union[str, torch.Tensor]]] - ) -> List[Tuple[Dict[str, List[Union[str, torch.Tensor]]], - Dict[str, List[Union[str, torch.Tensor]]]]]: + batch: list[dict[str, Union[str, torch.Tensor]]] + ) -> list[tuple[dict[str, list[Union[str, torch.Tensor]]], + dict[str, list[Union[str, torch.Tensor]]]]]: """ Disentanglement cannot be processed on different subjects at the same time, we need to load `pr` subjects one by one. The batch |