|
|
|
from torch import nn |
|
|
|
from .registry import CONV_LAYERS |
|
|
|
CONV_LAYERS.register_module('Conv1d', module=nn.Conv1d) |
|
CONV_LAYERS.register_module('Conv2d', module=nn.Conv2d) |
|
CONV_LAYERS.register_module('Conv3d', module=nn.Conv3d) |
|
CONV_LAYERS.register_module('Conv', module=nn.Conv2d) |
|
|
|
|
|
def build_conv_layer(cfg, *args, **kwargs): |
|
"""Build convolution layer. |
|
|
|
Args: |
|
cfg (None or dict): The conv layer config, which should contain: |
|
- type (str): Layer type. |
|
- layer args: Args needed to instantiate an conv layer. |
|
args (argument list): Arguments passed to the `__init__` |
|
method of the corresponding conv layer. |
|
kwargs (keyword arguments): Keyword arguments passed to the `__init__` |
|
method of the corresponding conv layer. |
|
|
|
Returns: |
|
nn.Module: Created conv layer. |
|
""" |
|
if cfg is None: |
|
cfg_ = dict(type='Conv2d') |
|
else: |
|
if not isinstance(cfg, dict): |
|
raise TypeError('cfg must be a dict') |
|
if 'type' not in cfg: |
|
raise KeyError('the cfg dict must contain the key "type"') |
|
cfg_ = cfg.copy() |
|
|
|
layer_type = cfg_.pop('type') |
|
if layer_type not in CONV_LAYERS: |
|
raise KeyError(f'Unrecognized norm type {layer_type}') |
|
else: |
|
conv_layer = CONV_LAYERS.get(layer_type) |
|
|
|
layer = conv_layer(*args, **kwargs, **cfg_) |
|
|
|
return layer |
|
|
