From f3caa4576000ab74b15068adb113eb855b156d87 Mon Sep 17 00:00:00 2001 From: kunkunlin Date: Mon, 23 Dec 2024 19:31:22 +0800 Subject: [PATCH 1/2] [A] Add registry feature --- chameleon/__init__.py | 1 + chameleon/base/__init__.py | 10 +- chameleon/base/blocks/__init__.py | 24 ++-- chameleon/base/blocks/conv_block.py | 43 +++--- chameleon/base/components/__init__.py | 24 ---- chameleon/base/components/activation.py | 8 +- chameleon/base/components/dropout.py | 7 +- chameleon/base/components/loss.py | 7 +- chameleon/base/components/norm.py | 6 + chameleon/base/components/pooling.py | 6 + chameleon/base/layers/__init__.py | 17 --- chameleon/base/layers/aspp.py | 53 ++++---- chameleon/base/layers/grl.py | 2 + chameleon/base/layers/selayer.py | 39 +++--- chameleon/base/layers/vae.py | 2 + chameleon/base/layers/weighted_sum.py | 13 +- chameleon/base/ops/positional_encoding.py | 3 + chameleon/base/optim/__init__.py | 39 ++---- chameleon/base/optim/polynomial_lr_warmup.py | 3 + chameleon/base/optim/warm_up.py | 23 +--- chameleon/base/utils.py | 8 +- .../normalized_levenshtein_similarity.py | 3 + chameleon/modules/__init__.py | 4 +- chameleon/modules/backbones/__init__.py | 40 +----- chameleon/modules/backbones/gpunet.py | 45 +++++++ chameleon/modules/necks/__init__.py | 29 +--- chameleon/modules/necks/bifpn.py | 112 ++++++++------- chameleon/modules/necks/fpn.py | 82 ++++++----- chameleon/registry/__init__.py | 5 + chameleon/registry/registry.py | 127 ++++++++++++++++++ chameleon/registry/root.py | 51 +++++++ tests/base/components/test_activation.py | 8 +- tests/base/components/test_loss.py | 2 +- tests/base/components/test_norm.py | 2 +- tests/base/layers/test_aspp.py | 5 +- tests/modules/backbone/test_backbone.py | 32 ++--- tests/modules/neck/test_bifpn.py | 2 +- tests/modules/neck/test_neck.py | 25 +--- tests/registry/test_registry.py | 29 ++++ tests/registry/test_root.py | 57 ++++++++ tests/tools/test_calflops.py | 4 +- 41 files changed, 619 insertions(+), 383 deletions(-) create mode 100644 chameleon/registry/__init__.py create mode 100644 chameleon/registry/registry.py create mode 100644 chameleon/registry/root.py create mode 100644 tests/registry/test_registry.py create mode 100644 tests/registry/test_root.py diff --git a/chameleon/__init__.py b/chameleon/__init__.py index e377bf4..954cce1 100644 --- a/chameleon/__init__.py +++ b/chameleon/__init__.py @@ -1,6 +1,7 @@ from .base import * from .metrics import * from .modules import * +from .registry import * from .tools import * __version__ = '0.1.0' diff --git a/chameleon/base/__init__.py b/chameleon/base/__init__.py index 7c8749e..7e82999 100644 --- a/chameleon/base/__init__.py +++ b/chameleon/base/__init__.py @@ -1,8 +1,8 @@ -from .blocks import build_block, list_blocks -from .components import build_component, list_components -from .layers import build_layer, list_layers -from .optim import (build_lr_scheduler, build_optimizer, list_lr_schedulers, - list_optimizers) +from .blocks import * +from .components import * +from .layers import * +from .ops import * +from .optim import * from .power_module import PowerModule from .utils import (has_children, initialize_weights_, replace_module, replace_module_attr_value) diff --git a/chameleon/base/blocks/__init__.py b/chameleon/base/blocks/__init__.py index 2eeedf6..2428686 100644 --- a/chameleon/base/blocks/__init__.py +++ b/chameleon/base/blocks/__init__.py @@ -1,21 +1,19 @@ -import fnmatch - from .conv_block import Conv2dBlock, SeparableConv2dBlock # from .mamba_block import build_mamba_block # from .vit_block import build_vit_block -def build_block(name, **kwargs): - cls = globals().get(name, None) - if cls is None: - raise ValueError(f'Block named {name} is not support.') - return cls(**kwargs) +# def build_block(name, **kwargs): +# cls = globals().get(name, None) +# if cls is None: +# raise ValueError(f'Block named {name} is not support.') +# return cls(**kwargs) -def list_blocks(filter=''): - block_list = [k for k in globals().keys() if 'Block' in k] - if len(filter): - return fnmatch.filter(block_list, filter) # include these blocks - else: - return block_list +# def list_blocks(filter=''): +# block_list = [k for k in globals().keys() if 'Block' in k] +# if len(filter): +# return fnmatch.filter(block_list, filter) # include these blocks +# else: +# return block_list diff --git a/chameleon/base/blocks/conv_block.py b/chameleon/base/blocks/conv_block.py index e10b330..e012b07 100644 --- a/chameleon/base/blocks/conv_block.py +++ b/chameleon/base/blocks/conv_block.py @@ -3,10 +3,11 @@ import torch import torch.nn as nn -from ..components import build_component +from ...registry import BLOCKS, COMPONENTS from ..power_module import PowerModule +@BLOCKS.register_module() class SeparableConv2dBlock(PowerModule): def __init__( @@ -17,10 +18,10 @@ def __init__( stride: Union[int, Tuple[int, int]] = 1, padding: Union[int, Tuple[int, int]] = 1, bias: bool = False, - inner_norm: Optional[Union[dict, nn.Module]] = None, - inner_act: Optional[Union[dict, nn.Module]] = None, - norm: Optional[Union[dict, nn.Module]] = None, - act: Optional[Union[dict, nn.Module]] = None, + inner_norm: Optional[dict] = None, + inner_act: Optional[dict] = None, + norm: Optional[dict] = None, + act: Optional[dict] = None, init_type: str = 'normal', ): """ @@ -41,13 +42,13 @@ def __init__( Whether to include a bias term in the convolutional layer. Noted: if normalization layer is not None, bias will always be set to False. Defaults to False. - inner_norm (dict or nn.Module, optional): + inner_norm (dict, optional): Configuration of normalization layer between dw and pw layer. Defaults to None. - inner_act (dict or nn.Module, optional): + inner_act (dict, optional): Configuration of activation layer between dw and pw layer. Defaults to None. - norm (dict or nn.Module, optional): + norm (dict, optional): Configuration of normalization layer after pw layer. Defaults to None. - act (dict or nn.Module, optional): + act (dict, optional): Configuration of activation layer after pw layer. Defaults to None. init_type (str, optional): Initialization method for the model parameters. Defaults to 'normal'. @@ -77,13 +78,13 @@ def __init__( bias=bias, ) if inner_norm is not None: - self.block['inner_norm'] = build_component(**inner_norm) if isinstance(inner_norm, dict) else inner_norm + self.block['inner_norm'] = COMPONENTS.build(inner_norm) if isinstance(inner_norm, dict) else inner_norm if inner_act is not None: - self.block['inner_act'] = build_component(**inner_act) if isinstance(inner_act, dict) else inner_act + self.block['inner_act'] = COMPONENTS.build(inner_act) if isinstance(inner_act, dict) else inner_act if norm is not None: - self.block['norm'] = build_component(**norm) if isinstance(norm, dict) else norm + self.block['norm'] = COMPONENTS.build(norm) if isinstance(norm, dict) else norm if act is not None: - self.block['act'] = build_component(**act) if isinstance(act, dict) else act + self.block['act'] = COMPONENTS.build(act) if isinstance(act, dict) else act self.initialize_weights_(init_type) def forward(self, x: torch.Tensor) -> torch.Tensor: @@ -92,8 +93,8 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return x +@BLOCKS.register_module() class Conv2dBlock(PowerModule): - def __init__( self, in_channels: Union[float, int], @@ -105,8 +106,8 @@ def __init__( groups: int = 1, bias: bool = False, padding_mode: str = 'zeros', - norm: Union[dict, nn.Module] = None, - act: Union[dict, nn.Module] = None, + norm: Optional[dict] = None, + act: Optional[dict] = None, init_type: str = 'normal', ): """ @@ -136,13 +137,13 @@ def __init__( padding_mode (str, optional): Options = {'zeros', 'reflect', 'replicate', 'circular'}. Defaults to 'zeros'. - norm (Union[dict, nn.Module], optional): + norm (Optional[dict], optional): normalization layer or a dictionary of arguments for building a normalization layer. Default to None. - act (Union[dict, nn.Module], optional): + act (Optional[dict], optional): Activation function or a dictionary of arguments for building an activation function. Default to None. - pool (Union[dict, nn.Module], optional): + pool (Optional[dict], optional): pooling layer or a dictionary of arguments for building a pooling layer. Default to None. init_type (str): @@ -180,9 +181,9 @@ def __init__( padding_mode=padding_mode, ) if norm is not None: - self.block['norm'] = build_component(**norm) if isinstance(norm, dict) else norm + self.block['norm'] = COMPONENTS.build(norm) if isinstance(norm, dict) else norm if act is not None: - self.block['act'] = build_component(**act) if isinstance(act, dict) else act + self.block['act'] = COMPONENTS.build(act) if isinstance(act, dict) else act self.initialize_weights_(init_type) diff --git a/chameleon/base/components/__init__.py b/chameleon/base/components/__init__.py index 9502ffd..bfa7884 100644 --- a/chameleon/base/components/__init__.py +++ b/chameleon/base/components/__init__.py @@ -1,29 +1,5 @@ -import fnmatch - from .activation import * from .dropout import * from .loss import * from .norm import * from .pooling import * - - -def build_component_cls(name): - cls = globals().get(name, None) - if cls is None: - raise ValueError(f'Component named {name} is not support.') - return cls - - -def build_component(name, **options): - cls = globals().get(name, None) - if cls is None: - raise ValueError(f'Component named {name} is not support.') - return cls(**options) - - -def list_components(filter=''): - component_list = [k for k in globals().keys() if 'Component' in k] - if len(filter): - return fnmatch.filter(component_list, filter) # include these components - else: - return component_list diff --git a/chameleon/base/components/activation.py b/chameleon/base/components/activation.py index 82368c1..66b202c 100644 --- a/chameleon/base/components/activation.py +++ b/chameleon/base/components/activation.py @@ -1,5 +1,3 @@ -from typing import Union - import torch import torch.nn as nn import torch.nn.functional as F @@ -12,6 +10,8 @@ Softshrink, Softsign, Tanh, Tanhshrink, Threshold) +from ...registry import COMPONENTS + __all__ = [ 'Swish', 'Hsigmoid', 'Hswish', 'StarReLU', 'SquaredReLU', ] @@ -95,3 +95,7 @@ def forward(self, x): # Ref: https://pytorch.org/docs/stable/generated/torch.nn.SiLU.html Swish = nn.SiLU + + +for k in __all__: + COMPONENTS.register_module(name=k, module=globals()[k]) diff --git a/chameleon/base/components/dropout.py b/chameleon/base/components/dropout.py index 6232ba0..0ed65bd 100644 --- a/chameleon/base/components/dropout.py +++ b/chameleon/base/components/dropout.py @@ -1,6 +1,11 @@ -import torch.nn as nn from torch.nn import AlphaDropout, Dropout, Dropout2d, Dropout3d +from ...registry import COMPONENTS + __all__ = [ 'Dropout', 'Dropout2d', 'Dropout3d', 'AlphaDropout', ] + + +for k in __all__: + COMPONENTS.register_module(name=k, module=globals()[k]) diff --git a/chameleon/base/components/loss.py b/chameleon/base/components/loss.py index 61e48f9..b9f6d68 100644 --- a/chameleon/base/components/loss.py +++ b/chameleon/base/components/loss.py @@ -1,5 +1,4 @@ import math -from typing import Union import torch import torch.nn as nn @@ -7,6 +6,8 @@ CrossEntropyLoss, CTCLoss, KLDivLoss, L1Loss, MSELoss, SmoothL1Loss) +from ...registry import COMPONENTS + __all__ = [ 'AWingLoss', 'WeightedAWingLoss', 'BCELoss', 'BCEWithLogitsLoss', 'CrossEntropyLoss', @@ -139,3 +140,7 @@ def dice_loss(self, input, target): def forward(self, input, target): dice_loss = self.dice_loss(input, target) return torch.log(torch.cosh(dice_loss)) + + +for k in __all__: + COMPONENTS.register_module(name=k, module=globals()[k]) diff --git a/chameleon/base/components/norm.py b/chameleon/base/components/norm.py index 249e7f1..5f1a256 100644 --- a/chameleon/base/components/norm.py +++ b/chameleon/base/components/norm.py @@ -10,6 +10,8 @@ from torch.nn.modules.normalization import (CrossMapLRN2d, GroupNorm, LayerNorm, LocalResponseNorm) +from ...registry import COMPONENTS + __all__ = [ 'BatchNorm1d', 'BatchNorm2d', 'BatchNorm3d', 'SyncBatchNorm', 'InstanceNorm1d', 'InstanceNorm2d', 'InstanceNorm3d', 'CrossMapLRN2d', 'GroupNorm', 'LayerNorm', @@ -42,3 +44,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: self.weight, self.bias, self.eps) x = x.permute(0, 3, 1, 2) return x + + +for k in __all__: + COMPONENTS.register_module(name=k, module=globals()[k]) diff --git a/chameleon/base/components/pooling.py b/chameleon/base/components/pooling.py index f84723d..e800db9 100644 --- a/chameleon/base/components/pooling.py +++ b/chameleon/base/components/pooling.py @@ -8,6 +8,8 @@ AvgPool1d, AvgPool2d, AvgPool3d, MaxPool1d, MaxPool2d, MaxPool3d) +from ...registry import COMPONENTS + __all__ = [ 'AvgPool1d', 'AvgPool2d', 'AvgPool3d', 'MaxPool1d', 'MaxPool2d', 'MaxPool3d', 'AdaptiveAvgPool1d', 'AdaptiveAvgPool2d', @@ -44,3 +46,7 @@ def __init__(self): def forward(self, x: torch.Tensor) -> torch.Tensor: """Apply global max pooling on the input tensor.""" return self.pool(x) + + +for k in __all__: + COMPONENTS.register_module(name=k, module=globals()[k]) diff --git a/chameleon/base/layers/__init__.py b/chameleon/base/layers/__init__.py index 1d8a493..4aa47d6 100644 --- a/chameleon/base/layers/__init__.py +++ b/chameleon/base/layers/__init__.py @@ -1,22 +1,5 @@ -import fnmatch - from .aspp import ASPP from .grl import GradientReversalLayer from .selayer import SELayer from .vae import VAE from .weighted_sum import WeightedSum - - -def build_layer(name, **options): - cls = globals().get(name, None) - if cls is None: - raise ValueError(f'Layer named {name} is not support.') - return cls(**options) - - -def list_layers(filter=''): - layer_list = [k for k in globals().keys() if 'Layer' in k] - if len(filter): - return fnmatch.filter(layer_list, filter) # include these layers - else: - return layer_list diff --git a/chameleon/base/layers/aspp.py b/chameleon/base/layers/aspp.py index 952a2ce..f2358b5 100644 --- a/chameleon/base/layers/aspp.py +++ b/chameleon/base/layers/aspp.py @@ -1,10 +1,8 @@ -from copy import deepcopy import torch import torch.nn as nn -from ..blocks import build_block -from ..components import Hswish +from ...registry import BLOCKS, COMPONENTS, LAYERS from ..power_module import PowerModule __all__ = ['ASPP'] @@ -16,6 +14,7 @@ """ +@LAYERS.register_module() class ASPP(PowerModule): ARCHS = { @@ -30,7 +29,7 @@ def __init__( self, in_channels: int, out_channels: int, - output_activate: nn.Module = nn.ReLU(), + out_act: dict = {'name': 'ReLU', 'inplace': True}, ): """ Constructor for the ASPP class. @@ -40,35 +39,39 @@ def __init__( Number of input channels. out_channels (int): Number of output channels. - output_activate (nn.Module, optional): - Activation function to apply to the output. Defaults to nn.ReLU(). + out_act (dict, optional): + Activation function for the output layer. Defaults to ReLU. """ super().__init__() self.layers = nn.ModuleDict() for dilate_name, cfg in self.ARCHS.items(): ksize, stride, padding, dilation, use_hs = cfg - layer = build_block( - 'Conv2dBlock', - in_channels=in_channels, - out_channels=in_channels, - kernel=ksize, - stride=stride, - padding=padding, - dilation=dilation, - norm=nn.BatchNorm2d(in_channels), - act=Hswish() if use_hs else nn.ReLU(), + layer = BLOCKS.build( + { + 'name': 'Conv2dBlock', + 'in_channels': in_channels, + 'out_channels': in_channels, + 'kernel': ksize, + 'stride': stride, + 'padding': padding, + 'dilation': dilation, + 'norm': COMPONENTS.build({'name': 'BatchNorm2d', 'num_features': in_channels}), + 'act': COMPONENTS.build({'name': 'Hswish' if use_hs else 'ReLU'}), + } ) self.layers[dilate_name] = layer - self.output_layer = build_block( - 'Conv2dBlock', - in_channels=in_channels * len(self.layers), - out_channels=out_channels, - kernel=1, - stride=1, - padding=0, - norm=nn.BatchNorm2d(out_channels), - act=deepcopy(output_activate), + self.output_layer = BLOCKS.build( + { + 'name': 'Conv2dBlock', + 'in_channels': in_channels * len(self.layers), + 'out_channels': out_channels, + 'kernel': 1, + 'stride': 1, + 'padding': 0, + 'norm': COMPONENTS.build({'name': 'BatchNorm2d', 'num_features': out_channels}), + 'act': COMPONENTS.build(out_act), + } ) def forward(self, x: torch.Tensor) -> torch.Tensor: diff --git a/chameleon/base/layers/grl.py b/chameleon/base/layers/grl.py index b328b17..c585290 100644 --- a/chameleon/base/layers/grl.py +++ b/chameleon/base/layers/grl.py @@ -1,6 +1,7 @@ import torch from torch.autograd import Function +from ...registry import LAYERS from ..power_module import PowerModule __all__ = ['GradientReversalLayer'] @@ -24,6 +25,7 @@ def backward(ctx, grad_output): # pragma: no cover revgrad = RevGrad.apply +@LAYERS.register_module() class GradientReversalLayer(PowerModule): def __init__(self, warm_up=4000): diff --git a/chameleon/base/layers/selayer.py b/chameleon/base/layers/selayer.py index bbbd794..bc91684 100644 --- a/chameleon/base/layers/selayer.py +++ b/chameleon/base/layers/selayer.py @@ -1,12 +1,13 @@ import torch import torch.nn as nn -from ..blocks import build_block +from ...registry import BLOCKS, LAYERS from ..power_module import PowerModule __all__ = ['SELayer'] +@LAYERS.register_module() class SELayer(PowerModule): def __init__(self, in_channels: int, reduction: int = 4): @@ -26,23 +27,27 @@ def __init__(self, in_channels: int, reduction: int = 4): mid_channels = max(1, in_channels // reduction) self.avg_pool = nn.AdaptiveAvgPool2d(1) - self.fc1 = build_block( - 'Conv2dBlock', - in_channels=in_channels, - out_channels=mid_channels, - kernel=1, - stride=1, - padding=0, - act=nn.ReLU(True) + self.fc1 = BLOCKS.build( + { + 'name': 'Conv2dBlock', + 'in_channels': in_channels, + 'out_channels': mid_channels, + 'kernel': 1, + 'stride': 1, + 'padding': 0, + 'act': {'name': 'ReLU', 'inplace': True}, + } ) - self.fc2 = build_block( - 'Conv2dBlock', - in_channels=mid_channels, - out_channels=in_channels, - kernel=1, - stride=1, - padding=0, - act=nn.Sigmoid(), + self.fc2 = BLOCKS.build( + { + 'name': 'Conv2dBlock', + 'in_channels': mid_channels, + 'out_channels': in_channels, + 'kernel': 1, + 'stride': 1, + 'padding': 0, + 'act': {'name': 'Sigmoid'}, + } ) def forward(self, x: torch.Tensor) -> torch.Tensor: diff --git a/chameleon/base/layers/vae.py b/chameleon/base/layers/vae.py index 9396ee5..2912755 100644 --- a/chameleon/base/layers/vae.py +++ b/chameleon/base/layers/vae.py @@ -3,12 +3,14 @@ import torch import torch.nn as nn +from ...registry import LAYERS from ..components import GAP from ..power_module import PowerModule __all__ = ['VAE'] +@LAYERS.register_module() class VAE(PowerModule): def __init__(self, in_channels: int, out_channels: int, do_pooling: bool = False): diff --git a/chameleon/base/layers/weighted_sum.py b/chameleon/base/layers/weighted_sum.py index 5154322..9e31dc9 100644 --- a/chameleon/base/layers/weighted_sum.py +++ b/chameleon/base/layers/weighted_sum.py @@ -1,17 +1,17 @@ -from typing import List, Optional, Union +from typing import List, Optional import torch import torch.nn as nn -from ..components import build_component +from ...registry import COMPONENTS, LAYERS +@LAYERS.register_module() class WeightedSum(nn.Module): - def __init__( self, input_size: int, - act: Optional[Union[dict, nn.Module]] = None, + act: Optional[dict] = None, requires_grad: bool = True, ) -> None: """ @@ -20,7 +20,7 @@ def __init__( Args: input_size (int): The number of inputs to be summed. - act Optional[Union[dict, nn.Module]]: + act Optional[dict]: Optional activation function or dictionary of its parameters. Defaults to None. requires_grad (bool, optional): @@ -36,8 +36,7 @@ def __init__( if act is None: self.relu = nn.Identity() else: - self.relu = act if isinstance(act, nn.Module) \ - else build_component(**act) + self.relu = act if isinstance(act, nn.Module) else COMPONENTS.build(act) self.epsilon = 1e-4 def forward(self, x: List[torch.Tensor]) -> torch.Tensor: diff --git a/chameleon/base/ops/positional_encoding.py b/chameleon/base/ops/positional_encoding.py index b90acbd..089fc4d 100644 --- a/chameleon/base/ops/positional_encoding.py +++ b/chameleon/base/ops/positional_encoding.py @@ -2,9 +2,12 @@ import torch +from ...registry import OPS + __all__ = ['sinusoidal_positional_encoding_1d'] +@OPS.register_module() def sinusoidal_positional_encoding_1d(length, dim): """ Sinusoidal positional encoding for non-recurrent neural networks. REFERENCES: Attention Is All You Need diff --git a/chameleon/base/optim/__init__.py b/chameleon/base/optim/__init__.py index 1b59c24..6ec0a25 100644 --- a/chameleon/base/optim/__init__.py +++ b/chameleon/base/optim/__init__.py @@ -1,5 +1,3 @@ -import fnmatch - from torch.optim import (ASGD, LBFGS, SGD, Adadelta, Adagrad, Adam, Adamax, AdamW, RMSprop, Rprop, SparseAdam) from torch.optim.lr_scheduler import (CosineAnnealingLR, @@ -7,35 +5,20 @@ ExponentialLR, LambdaLR, MultiStepLR, OneCycleLR, ReduceLROnPlateau, StepLR) -from .polynomial_lr_warmup import PolynomialLRWarmup +from ...registry import OPTIMIZERS +from .polynomial_lr_warmup import * from .warm_up import * - -def build_optimizer(model_params, name, **optim_options): - cls_ = globals().get(name, None) - if cls_ is None: - raise ValueError(f'{name} is not supported optimizer.') - return cls_(model_params, **optim_options) - - -def build_lr_scheduler(optimizer, name, **lr_scheduler_options): - cls_ = globals().get(name, None) - if cls_ is None: - raise ValueError(f'{name} is not supported lr scheduler.') - return cls_(optimizer, **lr_scheduler_options) +__all__ = [ + 'ASGD', 'LBFGS', 'SGD', 'Adadelta', 'Adagrad', 'Adam', 'Adamax', 'AdamW', + 'RMSprop', 'Rprop', 'SparseAdam', 'CosineAnnealingLR', + 'CosineAnnealingWarmRestarts', 'CyclicLR', 'ExponentialLR', 'LambdaLR', + 'MultiStepLR', 'OneCycleLR', 'ReduceLROnPlateau', 'StepLR', +] -def list_optimizers(filter=''): - optimizer_list = [k for k in globals().keys() if k[0].isupper()] - if len(filter): - return [o for o in optimizer_list if filter in o.lower()] - else: - return optimizer_list +for k in __all__: + OPTIMIZERS.register_module(name=k, force=True, module=globals()[k]) -def list_lr_schedulers(filter=''): - lr_scheduler_list = [k for k in globals().keys() if 'LR' in k] - if len(filter): - return [o for o in lr_scheduler_list if filter in o.lower()] - else: - return lr_scheduler_list +__all__ += ['PolynomialLRWarmup', 'WrappedLRScheduler'] diff --git a/chameleon/base/optim/polynomial_lr_warmup.py b/chameleon/base/optim/polynomial_lr_warmup.py index 70fac78..e07d3f9 100644 --- a/chameleon/base/optim/polynomial_lr_warmup.py +++ b/chameleon/base/optim/polynomial_lr_warmup.py @@ -2,7 +2,10 @@ from torch.optim.lr_scheduler import _LRScheduler +from ...registry import OPTIMIZERS + +@OPTIMIZERS.register_module() class PolynomialLRWarmup(_LRScheduler): def __init__( diff --git a/chameleon/base/optim/warm_up.py b/chameleon/base/optim/warm_up.py index 2e97ddb..65ce76d 100644 --- a/chameleon/base/optim/warm_up.py +++ b/chameleon/base/optim/warm_up.py @@ -1,11 +1,14 @@ from typing import List from torch.optim import Optimizer -from torch.optim.lr_scheduler import MultiStepLR, _LRScheduler +from torch.optim.lr_scheduler import _LRScheduler -__all__ = ['WrappedLRScheduler', 'MultiStepLRWarmUp'] +from ...registry import OPTIMIZERS +__all__ = ['WrappedLRScheduler'] + +@OPTIMIZERS.register_module() class WrappedLRScheduler(_LRScheduler): """ Gradually warm-up(increasing) learning rate in optimizer. @@ -63,19 +66,3 @@ def step(self, epoch=None, metrics=None): self._last_lr = self.after_scheduler.get_last_lr() else: return super().step() - - -def MultiStepLRWarmUp( - optimizer: Optimizer, - milestones: List[int], - warmup_milestone: int, - gamma: float = 0.1, - last_epoch: int = -1, - interval='step', - verbose: bool = False, -): - scheduler = MultiStepLR(optimizer, milestones, gamma, last_epoch, verbose) - return WrappedLRScheduler(optimizer, - warmup_milestone, - after_scheduler=scheduler, - interval=interval) diff --git a/chameleon/base/utils.py b/chameleon/base/utils.py index 1e87c42..b235e7a 100644 --- a/chameleon/base/utils.py +++ b/chameleon/base/utils.py @@ -2,7 +2,7 @@ import torch.nn as nn -from .components import build_component, build_component_cls +from ..registry import COMPONENTS __all__ = ['has_children', 'replace_module', 'replace_module_attr_value', 'initialize_weights_'] @@ -34,8 +34,8 @@ def replace_module( if not isinstance(dst_module, (nn.Module, dict)): raise ValueError(f'dst_module = {dst_module} should be an instance of Module or dict.') - target = build_component_cls(target) if isinstance(target, str) else target - dst_module = build_component(**dst_module) if isinstance(dst_module, dict) else dst_module + target = COMPONENTS.get(target) if isinstance(target, str) else target + dst_module = COMPONENTS.build(dst_module) if isinstance(dst_module, dict) else dst_module for name, m in model.named_children(): if has_children(m): @@ -60,7 +60,7 @@ def replace_module_attr_value( attr_name (str): The name of the attribute you want to modify. attr_value (Any): The new value of the attribute. """ - target = build_component_cls(target) if isinstance(target, str) else target + target = COMPONENTS.get(target) if isinstance(target, str) else target for module in model.modules(): if isinstance(module, target): setattr(module, attr_name, attr_value) diff --git a/chameleon/metrics/normalized_levenshtein_similarity.py b/chameleon/metrics/normalized_levenshtein_similarity.py index ee47ab4..accc73f 100644 --- a/chameleon/metrics/normalized_levenshtein_similarity.py +++ b/chameleon/metrics/normalized_levenshtein_similarity.py @@ -6,7 +6,10 @@ from torchmetrics.text import EditDistance from torchmetrics.utilities.data import dim_zero_cat +from ..registry import METRICS + +@METRICS.register_module() class NormalizedLevenshteinSimilarity(Metric): """ Normalized Levenshtein Similarity (NLS) is a metric that computes the diff --git a/chameleon/modules/__init__.py b/chameleon/modules/__init__.py index 5167faf..d538d3d 100644 --- a/chameleon/modules/__init__.py +++ b/chameleon/modules/__init__.py @@ -1,2 +1,2 @@ -from .backbones import build_backbone, list_backbones -from .necks import build_neck, list_necks +from .backbones import * +from .necks import * diff --git a/chameleon/modules/backbones/__init__.py b/chameleon/modules/backbones/__init__.py index ef12c5c..4c62191 100644 --- a/chameleon/modules/backbones/__init__.py +++ b/chameleon/modules/backbones/__init__.py @@ -1,40 +1,14 @@ -import fnmatch from functools import partial -from timm import create_model, list_models +import timm +from ...registry import BACKBONES from .gpunet import GPUNet -__all__ = [ - 'BACKBONE', 'build_backbone', 'list_backbones', -] +timm_models = timm.list_models() +for name in timm_models: + create_func = partial(timm.create_model, model_name=name) + BACKBONES.register_module(f'timm_{name}', module=create_func) -GPUNET_NAMES = [ - 'gpunet_0', - 'gpunet_1', - 'gpunet_2', - 'gpunet_p0', - 'gpunet_p1', - 'gpunet_d1', - 'gpunet_d2', -] - -BACKBONE = { - **{k: partial(create_model, model_name=k) for k in list_models()}, - **{name: partial(GPUNet.build_gpunet, name=name) for name in GPUNET_NAMES}, -} - - -def build_backbone(name: str, **kwargs): - if name not in BACKBONE: - raise ValueError(f'Backbone={name} is not supported.') - return BACKBONE[name](**kwargs) - - -def list_backbones(filter=''): - model_list = list(BACKBONE.keys()) - if len(filter): - return fnmatch.filter(model_list, filter) # include these models - else: - return model_list +__all__ = ['GPUNet'] diff --git a/chameleon/modules/backbones/gpunet.py b/chameleon/modules/backbones/gpunet.py index fad57f6..bccda37 100644 --- a/chameleon/modules/backbones/gpunet.py +++ b/chameleon/modules/backbones/gpunet.py @@ -4,10 +4,12 @@ import torch.nn as nn from ...base import PowerModule, has_children +from ...registry import BACKBONES __all__ = ['GPUNet'] +@BACKBONES.register_module() class GPUNet(PowerModule): MetaParams = { @@ -122,3 +124,46 @@ def _replace_padding(model): stages=stages, out_indices=out_indices, ) + + @classmethod + def build_gpunet_0(cls, **kwargs): + return cls.build_gpunet(name='gpunet_0', **kwargs) + + @classmethod + def build_gpunet_1(cls, **kwargs): + return cls.build_gpunet(name='gpunet_1', **kwargs) + + @classmethod + def build_gpunet_2(cls, **kwargs): + return cls.build_gpunet(name='gpunet_2', **kwargs) + + @classmethod + def build_gpunet_p0(cls, **kwargs): + return cls.build_gpunet(name='gpunet_p0', **kwargs) + + @classmethod + def build_gpunet_p1(cls, **kwargs): + return cls.build_gpunet(name='gpunet_p1', **kwargs) + + @classmethod + def build_gpunet_d1(cls, **kwargs): + return cls.build_gpunet(name='gpunet_d1', **kwargs) + + @classmethod + def build_gpunet_d2(cls, **kwargs): + return cls.build_gpunet(name='gpunet_d2', **kwargs) + + +GPUNETs = { + 'GPUNet_0': GPUNet.build_gpunet_0, + 'GPUNet_1': GPUNet.build_gpunet_1, + 'GPUNet_2': GPUNet.build_gpunet_2, + 'GPUNet_p0': GPUNet.build_gpunet_p0, + 'GPUNet_p1': GPUNet.build_gpunet_p1, + 'GPUNet_d1': GPUNet.build_gpunet_d1, + 'GPUNet_d2': GPUNet.build_gpunet_d2, +} + + +for k, v in GPUNETs.items(): + BACKBONES.register_module(name=k, module=v) diff --git a/chameleon/modules/necks/__init__.py b/chameleon/modules/necks/__init__.py index fe48a90..213eff1 100644 --- a/chameleon/modules/necks/__init__.py +++ b/chameleon/modules/necks/__init__.py @@ -1,27 +1,2 @@ -import fnmatch - -from .bifpn import BiFPN, BiFPNs -from .fpn import FPN - -NECK = { - 'fpn': FPN, - 'bifpn': BiFPN, - 'bifpns': BiFPNs, -} - - -def build_neck(name: str, **kwargs): - if name in NECK: - neck = NECK[name](**kwargs) - else: - raise ValueError(f'Neck={name} is not support.') - - return neck - - -def list_necks(filter=''): - model_list = list(NECK.keys()) - if len(filter): - return fnmatch.filter(model_list, filter) # include these models - else: - return model_list +from .bifpn import * +from .fpn import * diff --git a/chameleon/modules/necks/bifpn.py b/chameleon/modules/necks/bifpn.py index 6032bdd..dd19b39 100644 --- a/chameleon/modules/necks/bifpn.py +++ b/chameleon/modules/necks/bifpn.py @@ -1,14 +1,15 @@ -from copy import deepcopy from typing import List, Optional, Union import torch import torch.nn as nn -from ...base import PowerModule, build_block, build_layer +from ...base import PowerModule +from ...registry import BLOCKS, LAYERS, NECKS __all__ = ['BiFPN', 'BiFPNs'] +@NECKS.register_module() class BiFPN(PowerModule): def __init__( @@ -103,14 +104,16 @@ def __init__( in_channels = in_channels_list[i] if i < num_in_features else in_channels_list[-1] if in_channels != out_channels: conv1x1s.append( - build_block( - 'Conv2dBlock' if use_conv else 'SeparableConv2dBlock', - in_channels=in_channels, - out_channels=out_channels, - kernel=1, - stride=1, - padding=0, - norm=deepcopy(norm), + BLOCKS.build( + { + 'name': 'Conv2dBlock' if use_conv else 'SeparableConv2dBlock', + 'in_channels': in_channels, + 'out_channels': out_channels, + 'kernel': 1, + 'stride': 1, + 'padding': 0, + 'norm': norm, + } ) ) else: @@ -118,44 +121,50 @@ def __init__( self.conv1x1s = nn.ModuleList(conv1x1s) self.conv_up_3x3s = nn.ModuleList([ - build_block( - 'Conv2dBlock' if use_conv else 'SeparableConv2dBlock', - in_channels=out_channels, - out_channels=out_channels, - kernel=3, - stride=1, - padding=1, - norm=deepcopy(norm), - act=deepcopy(act), + BLOCKS.build( + { + 'name': 'Conv2dBlock' if use_conv else 'SeparableConv2dBlock', + 'in_channels': out_channels, + 'out_channels': out_channels, + 'kernel': 3, + 'stride': 1, + 'padding': 1, + 'norm': norm, + 'act': act, + } ) for _ in range(num_out_features - 1) ]) self.conv_down_3x3s = nn.ModuleList([ - build_block( - 'Conv2dBlock' if use_conv else 'SeparableConv2dBlock', - in_channels=out_channels, - out_channels=out_channels, - kernel=3, - stride=1, - padding=1, - norm=deepcopy(norm), - act=deepcopy(act), + BLOCKS.build( + { + 'name': 'Conv2dBlock' if use_conv else 'SeparableConv2dBlock', + 'in_channels': out_channels, + 'out_channels': out_channels, + 'kernel': 3, + 'stride': 1, + 'padding': 1, + 'norm': norm, + 'act': act, + } ) for _ in range(num_out_features - 1) ]) if extra_layers > 0: self.extra_conv_downs = nn.ModuleList([ - build_block( - 'Conv2dBlock' if use_conv else 'SeparableConv2dBlock', - in_channels=in_channels_list[-1], - out_channels=in_channels_list[-1], - kernel=3, - stride=2, - padding=1, - norm=nn.BatchNorm2d(in_channels_list[-1]) if norm is not None else None, - act=deepcopy(act), + BLOCKS.build( + { + 'name': 'Conv2dBlock' if use_conv else 'SeparableConv2dBlock', + 'in_channels': in_channels_list[-1], + 'out_channels': in_channels_list[-1], + 'kernel': 3, + 'stride': 2, + 'padding': 1, + 'norm': {'name': norm['name'], 'num_features': in_channels_list[-1]} if norm is not None else None, + 'act': act, + } ) for _ in range(extra_layers) ]) @@ -180,12 +189,26 @@ def __init__( # Weight self.weighted_sum_2_input = nn.ModuleList([ - build_layer('WeightedSum', input_size=2, act=nn.ReLU(False), requires_grad=attention) + LAYERS.build( + { + 'name': 'WeightedSum', + 'input_size': 2, + 'act': {'name': 'ReLU', 'inplace': True}, + 'requires_grad': attention, + } + ) for _ in range(num_out_features) ]) self.weighted_sum_3_input = nn.ModuleList([ - build_layer('WeightedSum', input_size=3, act=nn.ReLU(False), requires_grad=attention) + LAYERS.build( + { + 'name': 'WeightedSum', + 'input_size': 3, + 'act': {'name': 'ReLU', 'inplace': True}, + 'requires_grad': attention, + } + ) for _ in range(num_out_features-2) ]) @@ -278,9 +301,8 @@ def build_convbifpn( out_channels=out_channels, extra_layers=extra_layers, out_indices=out_indices, - norm=nn.BatchNorm2d(num_features=out_channels, - momentum=0.003, eps=1e-4), - act=nn.ReLU(False), + norm={'name': 'BatchNorm2d', 'num_features': out_channels, 'momentum': 0.003, 'eps': 1e-4}, + act={'name': 'ReLU', 'inplace': False}, upsample_mode=upsample_mode, use_conv=True, attention=attention, @@ -301,15 +323,15 @@ def build_bifpn( out_channels=out_channels, extra_layers=extra_layers, out_indices=out_indices, - norm=nn.BatchNorm2d(num_features=out_channels, - momentum=0.003, eps=1e-4), - act=nn.ReLU(False), + norm={'name': 'BatchNorm2d', 'num_features': out_channels, 'momentum': 0.003, 'eps': 1e-4}, + act={'name': 'ReLU', 'inplace': False}, upsample_mode=upsample_mode, use_conv=False, attention=attention, ) +@NECKS.register_module() class BiFPNs(PowerModule): def __init__( diff --git a/chameleon/modules/necks/fpn.py b/chameleon/modules/necks/fpn.py index 1e4cf5a..9216fc3 100644 --- a/chameleon/modules/necks/fpn.py +++ b/chameleon/modules/necks/fpn.py @@ -1,14 +1,15 @@ -from copy import deepcopy -from typing import List, Optional, Union +from typing import List, Optional import torch import torch.nn as nn -from ...base import PowerModule, build_block +from ...base import PowerModule +from ...registry import BLOCKS, NECKS __all__ = ['FPN'] +@NECKS.register_module() class FPN(PowerModule): def __init__( @@ -17,8 +18,8 @@ def __init__( out_channels: int, extra_layers: int = 0, out_indices: Optional[List[int]] = None, - norm: Optional[Union[dict, nn.Module]] = None, - act: Optional[Union[dict, nn.Module]] = None, + norm: Optional[dict] = None, + act: Optional[dict] = None, upsample_mode: str = 'bilinear', use_dwconv: bool = False, ) -> None: @@ -36,10 +37,10 @@ def __init__( out_indices (Optional[List[int]], optional): A list of integers indicating the indices of the feature maps to output. If None, all feature maps are output. Defaults to None. - norm Optional[Union[dict, nn.Module]]: + norm Optional[dict]: Optional normalization module or dictionary of its parameters. Defaults to None. - act Optional[Union[dict, nn.Module]]: + act Optional[dict]: Optional activation function or dictionary of its parameters. Defaults to None. upsample_mode (str, optional): @@ -75,14 +76,16 @@ def __init__( in_channels = in_channels_list[i] if i < num_in_features else in_channels_list[-1] if in_channels != out_channels: self.conv1x1s.append( - build_block( - "Conv2dBlock" if not use_dwconv else "SeparableConv2dBlock", - in_channels=in_channels, - out_channels=out_channels, - kernel=1, - stride=1, - padding=0, - norm=deepcopy(norm), + BLOCKS.build( + { + 'name': "Conv2dBlock" if not use_dwconv else "SeparableConv2dBlock", + 'in_channels': in_channels, + 'out_channels': out_channels, + 'kernel': 1, + 'stride': 1, + 'padding': 0, + 'norm': norm, + } ) ) else: @@ -90,30 +93,35 @@ def __init__( self.conv1x1s = nn.ModuleList(self.conv1x1s) self.smooth3x3s = nn.ModuleList([ - build_block( - "Conv2dBlock" if not use_dwconv else "SeparableConv2dBlock", - in_channels=out_channels, - out_channels=out_channels, - kernel=3, - stride=1, - padding=1, - norm=deepcopy(norm), - act=deepcopy(act), + BLOCKS.build( + { + 'name': "Conv2dBlock" if not use_dwconv else "SeparableConv2dBlock", + 'in_channels': out_channels, + 'out_channels': out_channels, + 'kernel': 3, + 'stride': 1, + 'padding': 1, + 'norm': norm, + 'act': act, + + } ) for _ in range(num_out_features - 1) ]) if extra_layers > 0: self.extra_conv_downs = nn.ModuleList([ - build_block( - "Conv2dBlock" if not use_dwconv else "SeparableConv2dBlock", - in_channels=in_channels_list[-1], - out_channels=in_channels_list[-1], - kernel=3, - stride=2, - padding=1, - norm=getattr(nn, norm.__class__.__name__)(in_channels_list[-1]) if norm is not None else None, - act=deepcopy(act), + BLOCKS.build( + { + 'name': "Conv2dBlock" if not use_dwconv else "SeparableConv2dBlock", + 'in_channels': in_channels_list[-1], + 'out_channels': in_channels_list[-1], + 'kernel': 3, + 'stride': 2, + 'padding': 1, + 'norm': {'name': norm['name'], 'num_features': in_channels_list[-1]} if norm is not None else None, + 'act': act, + } ) for _ in range(extra_layers) ]) @@ -176,8 +184,8 @@ def build_dwfpn( out_channels=out_channels, extra_layers=extra_layers, out_indices=out_indices, - norm=nn.BatchNorm2d(num_features=out_channels, momentum=0.003), - act=nn.ReLU(False), + norm={'name': 'BatchNorm2d', 'num_features': out_channels, 'momentum': 0.003}, + act={'name': 'ReLU', 'inplace': True}, upsample_mode=upsample_mode, use_dwconv=True, ) @@ -196,8 +204,8 @@ def build_fpn( out_channels=out_channels, extra_layers=extra_layers, out_indices=out_indices, - norm=nn.BatchNorm2d(num_features=out_channels, momentum=0.003), - act=nn.ReLU(False), + norm={'name': 'BatchNorm2d', 'num_features': out_channels, 'momentum': 0.003}, + act={'name': 'ReLU', 'inplace': True}, upsample_mode=upsample_mode, use_dwconv=False, ) diff --git a/chameleon/registry/__init__.py b/chameleon/registry/__init__.py new file mode 100644 index 0000000..28353a1 --- /dev/null +++ b/chameleon/registry/__init__.py @@ -0,0 +1,5 @@ +from .registry import Registry +from .root import (BACKBONES, BLOCKS, COMPONENTS, LAYERS, METRICS, NECKS, OPS, + OPTIMIZERS, build_backbone, build_block, build_component, + build_layer, build_metric, build_neck, build_ops, + build_optimizer) diff --git a/chameleon/registry/registry.py b/chameleon/registry/registry.py new file mode 100644 index 0000000..e092107 --- /dev/null +++ b/chameleon/registry/registry.py @@ -0,0 +1,127 @@ +import fnmatch +import inspect +from collections.abc import Callable +from typing import Any, Dict, List, Optional, Type, Union + +from rich.console import Console +from rich.table import Table + + +def build_from_cfg(cfg: dict, registry: "Registry") -> Any: + + if not isinstance(cfg, dict): + raise TypeError(f'`cfg` should be a dict, ConfigDict or Config, but got {type(cfg)}') + + if 'name' not in cfg: + raise KeyError('`cfg` must contain the key "name"') + + if not isinstance(registry, Registry): + raise TypeError(f'registry must be a chameleon.Registry object, but got {type(registry)}') + + kwargs = cfg.copy() + name = kwargs.pop('name') + obj_cls = registry.get(name) + if inspect.isclass(obj_cls) or inspect.ismethod(obj_cls): + obj = obj_cls(**kwargs) + else: + obj = obj_cls + return obj + + +class Registry: + + def __init__(self, name: str): + self._name = name + self._module_dict: Dict[str, Type] = dict() + + def __len__(self): + return len(self._module_dict) + + def __contains__(self, key): + return self.get(key) is not None + + def __repr__(self): + table = Table(title=f'Registry of {self._name}') + table.add_column('Names', justify='left', style='cyan') + table.add_column('Objects', justify='left', style='green') + + for name, obj in sorted(self._module_dict.items()): + table.add_row(name, str(obj)) + + console = Console() + with console.capture() as capture: + console.print(table, end='') + + return capture.get() + + @property + def name(self): + return self._name + + @property + def module_dict(self): + return self._module_dict + + def get(self, key: str) -> Optional[Type]: + if not isinstance(key, str): + raise TypeError(f'key must be a str, but got {type(key)}') + + obj_cls = self.module_dict.get(key, None) + + if obj_cls is None: + raise KeyError(f'{key} is not in the {self.name} registry') + + return obj_cls + + def build(self, cfg: dict) -> Any: + return build_from_cfg(cfg, registry=self) + + def _register_module( + self, + module: Type, + module_name: Optional[Union[str, List[str]]] = None, + force: bool = False + ) -> None: + if not callable(module): + raise TypeError(f'module must be a callable, but got {type(module)}') + + if module_name is None: + module_name = module.__name__ + if isinstance(module_name, str): + module_name = [module_name] + for name in module_name: + if not force and name in self._module_dict: + existed_module = self.module_dict[name] + raise KeyError(f'{name} is already registered in {self.name} at {existed_module.__module__}') + self._module_dict[name] = module + + def register_module( + self, + name: str = None, + force: bool = False, + module: Optional[Type] = None, + ) -> Union[type, Callable]: + + if not (name is None or isinstance(name, str)): + raise TypeError(f'`name` must be a str or None, but got {type(name)}') + + if not isinstance(force, bool): + raise TypeError(f'`force` must be a bool, but got {type(force)}') + + # use it as a normal method: x.register_module(module=SomeClass) + if module is not None: + self._register_module(module=module, module_name=name, force=force) + return module + + # use it as a decorator: @x.register_module() + def _register(module): + self._register_module(module=module, module_name=name, force=force) + return module + + return _register + + def list_module(self, filter: Optional[str] = None) -> List[str]: + list_modules = list(self._module_dict.keys()) + if filter is not None: + list_modules = fnmatch.filter(list_modules, filter) + return list_modules diff --git a/chameleon/registry/root.py b/chameleon/registry/root.py new file mode 100644 index 0000000..e787df6 --- /dev/null +++ b/chameleon/registry/root.py @@ -0,0 +1,51 @@ +from .registry import Registry + +BLOCKS = Registry('blocks') + +COMPONENTS = Registry('components') + +LAYERS = Registry('layers') + +OPS = Registry('ops') + +OPTIMIZERS = Registry('optim') + +METRICS = Registry('metrics') + +BACKBONES = Registry('backbones') + +NECKS = Registry('necks') + +FUNCTIONS = Registry('functions') + + +def build_block(name, **options): + return BLOCKS.build({'name': name, **options}) + + +def build_component(name, **options): + return COMPONENTS.build({'name': name, **options}) + + +def build_layer(name, **options): + return LAYERS.build({'name': name, **options}) + + +def build_ops(name, **options): + return OPS.build({'name': name, **options}) + + +def build_optimizer(name, **options): + return OPTIMIZERS.build({'name': name, **options}) + + +def build_metric(name, **options): + return METRICS.build({'name': name, **options}) + + +def build_backbone(name, **options): + return BACKBONES.build({'name': name, **options}) + + +def build_neck(name, **options): + return NECKS.build({'name': name, **options}) diff --git a/tests/base/components/test_activation.py b/tests/base/components/test_activation.py index 24d62ed..eb01512 100644 --- a/tests/base/components/test_activation.py +++ b/tests/base/components/test_activation.py @@ -2,7 +2,7 @@ import torch import torch.nn as nn -from chameleon.base.components import SquaredReLU, StarReLU, build_component +from chameleon import SquaredReLU, StarReLU, build_component test_build_component_data = [ ('ReLU', nn.ReLU), @@ -10,14 +10,14 @@ ('Swish', nn.SiLU), ('StarReLU', StarReLU), ('SquaredReLU', SquaredReLU), - ('FakeActivation', ValueError) + ('FakeActivation', KeyError) ] @pytest.mark.parametrize('name, expected_output', test_build_component_data) def test_build_component(name, expected_output): - if expected_output == ValueError: - with pytest.raises(ValueError): + if expected_output == KeyError: + with pytest.raises(KeyError): build_component(name) else: assert isinstance(build_component(name), expected_output) diff --git a/tests/base/components/test_loss.py b/tests/base/components/test_loss.py index 3d7a060..20ab5d4 100644 --- a/tests/base/components/test_loss.py +++ b/tests/base/components/test_loss.py @@ -1,7 +1,7 @@ import pytest import torch -from chameleon.base.components import AWingLoss, WeightedAWingLoss +from chameleon import AWingLoss, WeightedAWingLoss @pytest.fixture(scope='module') diff --git a/tests/base/components/test_norm.py b/tests/base/components/test_norm.py index 803559f..ba431be 100644 --- a/tests/base/components/test_norm.py +++ b/tests/base/components/test_norm.py @@ -7,7 +7,7 @@ from torch.nn.modules.normalization import (CrossMapLRN2d, GroupNorm, LayerNorm, LocalResponseNorm) -from chameleon.base.components import LayerNorm2d, build_component +from chameleon import LayerNorm2d, build_component NORM_CLASSES = { 'BatchNorm1d': BatchNorm1d, diff --git a/tests/base/layers/test_aspp.py b/tests/base/layers/test_aspp.py index 22959c8..9609bb0 100644 --- a/tests/base/layers/test_aspp.py +++ b/tests/base/layers/test_aspp.py @@ -1,8 +1,7 @@ import pytest import torch -from chameleon.base.components.activation import Hswish -from chameleon.base.layers import ASPP +from chameleon import ASPP @pytest.fixture @@ -20,7 +19,7 @@ def test_aspp_layer(input_tensor): assert output.size() == (1, out_channels, 32, 32) # Test with Hswish activation function - aspp_layer = ASPP(in_channels, out_channels, output_activate=Hswish()) + aspp_layer = ASPP(in_channels, out_channels, out_act={'name': 'Hswish'}) output = aspp_layer(input_tensor) assert output.size() == (1, out_channels, 32, 32) diff --git a/tests/modules/backbone/test_backbone.py b/tests/modules/backbone/test_backbone.py index 32ee631..9cf87b6 100644 --- a/tests/modules/backbone/test_backbone.py +++ b/tests/modules/backbone/test_backbone.py @@ -1,7 +1,7 @@ import pytest import torch -from chameleon import build_backbone, list_backbones +from chameleon import build_backbone INPUT1 = torch.rand(1, 3, 320, 320) INPUT2 = torch.rand(1, 6, 224, 224) @@ -9,7 +9,7 @@ # gpunet ( INPUT1, - {'name': 'gpunet_0', }, + {'name': 'GPUNet_0', }, { 'out_shapes': [ torch.Size([1, 32, 160, 160]), @@ -22,7 +22,7 @@ ), ( INPUT1, - {'name': 'gpunet_1', 'out_indices': [0, 2, 3]}, + {'name': 'GPUNet_1', 'out_indices': [0, 2, 3]}, { 'out_shapes': [ torch.Size([1, 24, 160, 160]), @@ -33,7 +33,7 @@ ), ( INPUT1, - {'name': 'gpunet_2', 'out_indices': [0, 1, 2]}, + {'name': 'GPUNet_2', 'out_indices': [0, 1, 2]}, { 'out_shapes': [ torch.Size([1, 32, 160, 160]), @@ -44,7 +44,7 @@ ), ( INPUT1, - {'name': 'gpunet_p0'}, + {'name': 'GPUNet_p0'}, { 'out_shapes': [ torch.Size([1, 32, 160, 160]), @@ -57,7 +57,7 @@ ), ( INPUT1, - {'name': 'gpunet_p1'}, + {'name': 'GPUNet_p1'}, { 'out_shapes': [ torch.Size([1, 32, 160, 160]), @@ -70,7 +70,7 @@ ), ( INPUT1, - {'name': 'gpunet_d1'}, + {'name': 'GPUNet_d1'}, { 'out_shapes': [ torch.Size([1, 33, 160, 160]), @@ -83,7 +83,7 @@ ), ( INPUT1, - {'name': 'gpunet_d2', 'out_indices': [3, 4]}, + {'name': 'GPUNet_d2', 'out_indices': [3, 4]}, { 'out_shapes': [ torch.Size([1, 272, 20, 20]), @@ -93,7 +93,7 @@ ), ( INPUT1, - {'name': 'gpunet_d2', 'out_indices': [-1]}, + {'name': 'GPUNet_d2', 'out_indices': [-1]}, { 'out_shapes': [torch.Size([1, 384, 10, 10])] } @@ -110,17 +110,3 @@ def test_build_backbone(in_tensor, build_kwargs, expected): else: out_shapes = outs.shape assert out_shapes == expected['out_shapes'] - - -data = [ - ( - '*gpunet*', - ['gpunet_0', 'gpunet_1', 'gpunet_2', 'gpunet_p0', - 'gpunet_p1', 'gpunet_d1', 'gpunet_d2'] - ), -] - - -@ pytest.mark.parametrize('filter,expected', data) -def test_list_backbones(filter, expected): - assert list_backbones(filter) == expected diff --git a/tests/modules/neck/test_bifpn.py b/tests/modules/neck/test_bifpn.py index 738e682..59d7d5c 100644 --- a/tests/modules/neck/test_bifpn.py +++ b/tests/modules/neck/test_bifpn.py @@ -1,6 +1,6 @@ import torch -from chameleon.modules.necks import BiFPN, BiFPNs +from chameleon import BiFPN, BiFPNs def test_bifpn(): diff --git a/tests/modules/neck/test_neck.py b/tests/modules/neck/test_neck.py index e3157a7..ffae842 100644 --- a/tests/modules/neck/test_neck.py +++ b/tests/modules/neck/test_neck.py @@ -1,7 +1,7 @@ import pytest import torch -from chameleon.modules import build_neck, list_necks +from chameleon import build_neck INPUT1 = [ torch.rand(1, 16, 80, 80), @@ -12,7 +12,7 @@ data = [ ( INPUT1, - {'name': 'fpn', 'in_channels_list': [16, 32, 64], 'out_channels': 24}, + {'name': 'FPN', 'in_channels_list': [16, 32, 64], 'out_channels': 24}, {'out_shapes': [ torch.Size((1, 24, 80, 80)), torch.Size((1, 24, 40, 40)), @@ -21,7 +21,7 @@ ), ( INPUT1, - {'name': 'bifpn', 'in_channels_list': [16, 32, 64], 'out_channels': 24, 'extra_layers': 2}, + {'name': 'BiFPN', 'in_channels_list': [16, 32, 64], 'out_channels': 24, 'extra_layers': 2}, {'out_shapes': [ torch.Size((1, 24, 80, 80)), torch.Size((1, 24, 40, 40)), @@ -32,7 +32,7 @@ ), ( INPUT1, - {'name': 'bifpn', 'in_channels_list': [16, 32, 64], 'out_channels': 24, 'out_indices': [0, 1, 2]}, + {'name': 'BiFPN', 'in_channels_list': [16, 32, 64], 'out_channels': 24, 'out_indices': [0, 1, 2]}, {'out_shapes': [ torch.Size((1, 24, 80, 80)), torch.Size((1, 24, 40, 40)), @@ -51,20 +51,3 @@ def test_build_backbone(in_tensor, build_kwargs, expected): else: out_shapes = outs.shape assert out_shapes == expected['out_shapes'] - - -data = [ - ( - '', - ['fpn', 'bifpn', 'bifpns'] - ), - ( - '*bi*', - ['bifpn', 'bifpns'] - ), -] - - -@ pytest.mark.parametrize('filter,expected', data) -def test_list_backbones(filter, expected): - assert list_necks(filter) == expected diff --git a/tests/registry/test_registry.py b/tests/registry/test_registry.py new file mode 100644 index 0000000..b88ce9f --- /dev/null +++ b/tests/registry/test_registry.py @@ -0,0 +1,29 @@ +import pytest + +from chameleon.registry import Registry + +TEST = Registry(name='test') + + +def test_Registry(): + class Test1: + def __init__(self): + self.test = 1 + + TEST.register_module('test1', module=Test1) + + test_obj = TEST.build({'name': 'test1'}) + assert test_obj.test == 1 + + +def test_Registry_repr(): + class Test2: + def __init__(self): + self.test = 2 + TEST.register_module('test2', module=Test2) + assert 'test2' in repr(TEST) + + +def test_Registry_list_module(): + modules = TEST.list_module("*2") + assert 'test2' in modules diff --git a/tests/registry/test_root.py b/tests/registry/test_root.py new file mode 100644 index 0000000..28f952b --- /dev/null +++ b/tests/registry/test_root.py @@ -0,0 +1,57 @@ +import torch +from torchmetrics.metric import Metric + +from chameleon import ASPP, FPN, AdamW, AWingLoss, Conv2dBlock, GPUNet +from chameleon.registry import (BACKBONES, BLOCKS, COMPONENTS, LAYERS, METRICS, + NECKS, OPS, OPTIMIZERS) + + +def test_COMPONENTS(): + loss = COMPONENTS.build({'name': 'AWingLoss'}) + assert isinstance(loss, AWingLoss) + + +def test_OPTIMIZERS(): + model = BLOCKS.build({'name': 'Conv2dBlock', 'in_channels': 3, + 'out_channels': 3, 'kernel': 3, 'stride': 1, 'padding': 1}) + optimizer = OPTIMIZERS.build({'name': 'AdamW', 'params': model.parameters(), 'lr': 1e-3}) + assert isinstance(optimizer, AdamW) + + +def test_BLOCKS(): + model = BLOCKS.build({'name': 'Conv2dBlock', 'in_channels': 3, + 'out_channels': 3, 'kernel': 3, 'stride': 1, 'padding': 1}) + assert isinstance(model, Conv2dBlock) + + +def test_LAYERS(): + model = LAYERS.build({'name': 'ASPP', 'in_channels': 3, 'out_channels': 3}) + assert isinstance(model, ASPP) + + +def test_METRICS(): + model = METRICS.build({'name': 'NormalizedLevenshteinSimilarity'}) + assert isinstance(model, Metric) + + +def test_BACKBONES(): + model = BACKBONES.build({'name': 'GPUNet_0'}) + assert isinstance(model, GPUNet) + + +def test_NECKS(): + model = NECKS.build({'name': 'FPN', 'in_channels_list': [3, 3, 3], 'out_channels': 3}) + assert isinstance(model, FPN) + + +def test_OPS(): + func = OPS.build({'name': 'sinusoidal_positional_encoding_1d'}) + assert isinstance(func(10, 2), torch.Tensor) + + +def test_add_to_registry(): + + @COMPONENTS.register_module() + class Test: + pass + assert 'Test' in COMPONENTS diff --git a/tests/tools/test_calflops.py b/tests/tools/test_calflops.py index a10e6df..d5c92e2 100644 --- a/tests/tools/test_calflops.py +++ b/tests/tools/test_calflops.py @@ -1,11 +1,11 @@ import pytest -from chameleon.modules import build_backbone +from chameleon import BACKBONES from chameleon.tools import calculate_flops def test_calcualte_flops(): - model = build_backbone('resnet50') + model = BACKBONES.build({'name': 'timm_resnet50'}) flops, macs, params = calculate_flops(model, (1, 3, 224, 224)) assert flops == '8.21 GFLOPS' assert macs == '4.09 GMACs' From 7405a96de98f8c2f5382794f4e97dc16a6a2fde6 Mon Sep 17 00:00:00 2001 From: kunkunlin Date: Mon, 23 Dec 2024 20:03:55 +0800 Subject: [PATCH 2/2] [C] Update for timm registrys --- chameleon/modules/backbones/__init__.py | 12 +----------- chameleon/modules/backbones/timm.py | 15 +++++++++++++++ tests/registry/test_root.py | 1 + tests/tools/test_calflops.py | 3 ++- 4 files changed, 19 insertions(+), 12 deletions(-) create mode 100644 chameleon/modules/backbones/timm.py diff --git a/chameleon/modules/backbones/__init__.py b/chameleon/modules/backbones/__init__.py index 4c62191..ff94830 100644 --- a/chameleon/modules/backbones/__init__.py +++ b/chameleon/modules/backbones/__init__.py @@ -1,14 +1,4 @@ -from functools import partial - -import timm - -from ...registry import BACKBONES from .gpunet import GPUNet - -timm_models = timm.list_models() -for name in timm_models: - create_func = partial(timm.create_model, model_name=name) - BACKBONES.register_module(f'timm_{name}', module=create_func) - +from .timm import * __all__ = ['GPUNet'] diff --git a/chameleon/modules/backbones/timm.py b/chameleon/modules/backbones/timm.py new file mode 100644 index 0000000..c17ff87 --- /dev/null +++ b/chameleon/modules/backbones/timm.py @@ -0,0 +1,15 @@ +import timm +import torch.nn as nn + +from ...registry import BACKBONES + + +class Timm: + @staticmethod + def build_model(*args, **kwargs) -> nn.Module: + return timm.create_model(*args, **kwargs) + + +timm_models = timm.list_models() +for name in timm_models: + BACKBONES.register_module(f'timm_{name}', module=Timm.build_model) diff --git a/tests/registry/test_root.py b/tests/registry/test_root.py index 28f952b..e15a3a0 100644 --- a/tests/registry/test_root.py +++ b/tests/registry/test_root.py @@ -1,4 +1,5 @@ import torch +import torch.nn as nn from torchmetrics.metric import Metric from chameleon import ASPP, FPN, AdamW, AWingLoss, Conv2dBlock, GPUNet diff --git a/tests/tools/test_calflops.py b/tests/tools/test_calflops.py index d5c92e2..17e1635 100644 --- a/tests/tools/test_calflops.py +++ b/tests/tools/test_calflops.py @@ -5,7 +5,8 @@ def test_calcualte_flops(): - model = BACKBONES.build({'name': 'timm_resnet50'}) + timm_create = BACKBONES.build({'name': 'timm_resnet50'}) + model = timm_create(model_name='resnet50') flops, macs, params = calculate_flops(model, (1, 3, 224, 224)) assert flops == '8.21 GFLOPS' assert macs == '4.09 GMACs'