| | from fastai.layers import * |
| | from fastai.torch_core import * |
| | from torch.nn.parameter import Parameter |
| | from torch.autograd import Variable |
| |
|
| |
|
| | |
| |
|
| |
|
| | def custom_conv_layer( |
| | ni: int, |
| | nf: int, |
| | ks: int = 3, |
| | stride: int = 1, |
| | padding: int = None, |
| | bias: bool = None, |
| | is_1d: bool = False, |
| | norm_type: Optional[NormType] = NormType.Batch, |
| | use_activ: bool = True, |
| | leaky: float = None, |
| | transpose: bool = False, |
| | init: Callable = nn.init.kaiming_normal_, |
| | self_attention: bool = False, |
| | extra_bn: bool = False, |
| | ): |
| | "Create a sequence of convolutional (`ni` to `nf`), ReLU (if `use_activ`) and batchnorm (if `bn`) layers." |
| | if padding is None: |
| | padding = (ks - 1) // 2 if not transpose else 0 |
| | bn = norm_type in (NormType.Batch, NormType.BatchZero) or extra_bn == True |
| | if bias is None: |
| | bias = not bn |
| | conv_func = nn.ConvTranspose2d if transpose else nn.Conv1d if is_1d else nn.Conv2d |
| | conv = init_default( |
| | conv_func(ni, nf, kernel_size=ks, bias=bias, stride=stride, padding=padding), |
| | init, |
| | ) |
| | if norm_type == NormType.Weight: |
| | conv = weight_norm(conv) |
| | elif norm_type == NormType.Spectral: |
| | conv = spectral_norm(conv) |
| | layers = [conv] |
| | if use_activ: |
| | layers.append(relu(True, leaky=leaky)) |
| | if bn: |
| | layers.append((nn.BatchNorm1d if is_1d else nn.BatchNorm2d)(nf)) |
| | if self_attention: |
| | layers.append(SelfAttention(nf)) |
| | return nn.Sequential(*layers) |
| |
|
