| | |
| | import torch.nn as nn |
| |
|
| | from .registry import PADDING_LAYERS |
| |
|
| | PADDING_LAYERS.register_module('zero', module=nn.ZeroPad2d) |
| | PADDING_LAYERS.register_module('reflect', module=nn.ReflectionPad2d) |
| | PADDING_LAYERS.register_module('replicate', module=nn.ReplicationPad2d) |
| |
|
| |
|
| | def build_padding_layer(cfg, *args, **kwargs): |
| | """Build padding layer. |
| | |
| | Args: |
| | cfg (None or dict): The padding layer config, which should contain: |
| | - type (str): Layer type. |
| | - layer args: Args needed to instantiate a padding layer. |
| | |
| | Returns: |
| | nn.Module: Created padding layer. |
| | """ |
| | 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() |
| | padding_type = cfg_.pop('type') |
| | if padding_type not in PADDING_LAYERS: |
| | raise KeyError(f'Unrecognized padding type {padding_type}.') |
| | else: |
| | padding_layer = PADDING_LAYERS.get(padding_type) |
| |
|
| | layer = padding_layer(*args, **kwargs, **cfg_) |
| |
|
| | return layer |
| |
|
