A linear block that contains linear/norm/activation layers. For low level vision, we add spectral norm and padding layer. Args: in_features (int): Same as nn.Linear. out_features (int): Same as nn.Linear. bias (bool): Same as nn.Linear. act_cfg (dict): Confi
| 4 | |
| 5 | |
| 6 | class LinearModule(nn.Module): |
| 7 | """A linear block that contains linear/norm/activation layers. |
| 8 | |
| 9 | For low level vision, we add spectral norm and padding layer. |
| 10 | |
| 11 | Args: |
| 12 | in_features (int): Same as nn.Linear. |
| 13 | out_features (int): Same as nn.Linear. |
| 14 | bias (bool): Same as nn.Linear. |
| 15 | act_cfg (dict): Config dict for activation layer, "relu" by default. |
| 16 | inplace (bool): Whether to use inplace mode for activation. |
| 17 | with_spectral_norm (bool): Whether use spectral norm in linear module. |
| 18 | order (tuple[str]): The order of linear/activation layers. It is a |
| 19 | sequence of "linear", "norm" and "act". Examples are |
| 20 | ("linear", "act") and ("act", "linear"). |
| 21 | """ |
| 22 | |
| 23 | def __init__(self, |
| 24 | in_features, |
| 25 | out_features, |
| 26 | bias=True, |
| 27 | act_cfg=dict(type='ReLU'), |
| 28 | inplace=True, |
| 29 | with_spectral_norm=False, |
| 30 | order=('linear', 'act')): |
| 31 | super().__init__() |
| 32 | assert act_cfg is None or isinstance(act_cfg, dict) |
| 33 | self.act_cfg = act_cfg |
| 34 | self.inplace = inplace |
| 35 | self.with_spectral_norm = with_spectral_norm |
| 36 | self.order = order |
| 37 | assert isinstance(self.order, tuple) and len(self.order) == 2 |
| 38 | assert set(order) == set(['linear', 'act']) |
| 39 | |
| 40 | self.with_activation = act_cfg is not None |
| 41 | self.with_bias = bias |
| 42 | |
| 43 | # build linear layer |
| 44 | self.linear = nn.Linear(in_features, out_features, bias=bias) |
| 45 | # export the attributes of self.linear to a higher level for |
| 46 | # convenience |
| 47 | self.in_features = self.linear.in_features |
| 48 | self.out_features = self.linear.out_features |
| 49 | |
| 50 | if self.with_spectral_norm: |
| 51 | self.linear = nn.utils.spectral_norm(self.linear) |
| 52 | |
| 53 | # build activation layer |
| 54 | if self.with_activation: |
| 55 | act_cfg_ = act_cfg.copy() |
| 56 | act_cfg_.setdefault('inplace', inplace) |
| 57 | self.activate = build_activation_layer(act_cfg_) |
| 58 | |
| 59 | # Use msra init by default |
| 60 | self.init_weights() |
| 61 | |
| 62 | def init_weights(self): |
| 63 | if self.with_activation and self.act_cfg['type'] == 'LeakyReLU': |