| import einops |
| import torch |
| from torch import nn |
| import torch.nn.functional as F |
|
|
| from ..distill_layers import Conv1d |
|
|
|
|
| def trend_pool(x, kernel_size): |
| if kernel_size > 1: |
| pool_args = dict(kernel_size=kernel_size, stride=1, padding=kernel_size // 2) |
| return F.avg_pool1d(F.max_pool1d(x.abs(), **pool_args), **pool_args) |
| |
| else: |
| return x |
|
|
|
|
| class TrendPool(nn.Module): |
| def __init__(self, kernel_size=5): |
| super().__init__() |
| self.kernel_size = kernel_size |
|
|
| def forward(self, x): |
| return trend_pool(x, self.kernel_size) |
|
|
|
|
| class FirstBlock(nn.Module): |
| def __init__( |
| self, |
| target_dim, |
| conv_kernels=(7, 7, 7, 7), |
| pool_kernels=(1, 3, 5, 9), |
| dilation_rate=2, |
| ): |
| super().__init__() |
| assert target_dim % len(pool_kernels) == 0 |
| each_dim = target_dim // len(pool_kernels) |
| blocks = [] |
| for conv_kernel, pool_kernel in zip(conv_kernels, pool_kernels): |
| conv_dilation = pool_kernel // dilation_rate + 1 |
| conv_padding = (conv_kernel - 1) * conv_dilation // 2 |
| blocks.append( |
| nn.Sequential( |
| TrendPool(pool_kernel), |
| Conv1d( |
| 1, |
| each_dim, |
| kernel_size=conv_kernel, |
| dilation=conv_dilation, |
| padding=conv_padding, |
| ), |
| ) |
| ) |
| self.blocks = nn.ModuleList(blocks) |
|
|
| def forward(self, x): |
| return torch.cat([block(x) for block in self.blocks], dim=1) |
|
|
|
|
| class EnhanceBlock(FirstBlock): |
| def __init__(self, dim): |
| super().__init__(4, conv_kernels=(7, 7, 7, 7), pool_kernels=(1, 3, 5, 9)) |
| self.dim = dim |
| self.merge_layer = nn.Sequential( |
| |
| nn.InstanceNorm1d(4, affine=True), |
| nn.Conv1d(4, 1, kernel_size=1), |
| ) |
|
|
| def forward(self, x): |
| x = einops.rearrange(x, "b c t -> (b c) 1 t", c=self.dim) |
| y = super().forward(x) |
| y = self.merge_layer(y) |
| y = einops.rearrange(y, "(b c) 1 t -> b c t", c=self.dim) |
| return y |
|
|
|
|
| class SimpleEnhanceBlock(FirstBlock): |
| def __init__(self, dim): |
| super().__init__(4, conv_kernels=(7, 7, 7, 7), pool_kernels=(1, 3, 5, 9)) |
| self.dim = dim |
| self.merge_layer = nn.Sequential( |
| |
| nn.InstanceNorm1d(4, affine=True), |
| nn.Conv1d(4, self.dim, kernel_size=1), |
| ) |
|
|
| def forward(self, x): |
| xi = x[:, :1, :] |
| yi = super().forward(xi) |
| y = self.merge_layer(yi) |
| return x + y * x |
|
|
