File size: 974 Bytes
6b882c1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 | import torch
from torch import nn
from ..distill_layers import Conv1d
from . import base
class V3FirstBlock(base.FirstBlock): # (1, 5, 11, 21, 45)
def __init__(
self,
target_dim,
conv_kernels=(7, 7, 7, 7, 7),
pool_kernels=(1, 5, 11, 21, 45),
dilation_rate=7,
):
h_dim = len(pool_kernels) * 4
super().__init__(h_dim, conv_kernels, pool_kernels, dilation_rate=dilation_rate)
self.conv_1 = Conv1d(h_dim, h_dim * 4, kernel_size=1)
self.act = nn.GELU()
self.conv_2 = Conv1d(h_dim * 4 + 1, target_dim, kernel_size=1)
def forward(self, x):
h = super().forward(x)
h = self.conv_1(h)
h = self.act(h)
y = torch.cat([h, x], dim=1)
y = self.conv_2(y)
return y
FirstBlock = lambda dim: (
V3FirstBlock(
dim,
conv_kernels=(7, 7, 7, 7, 7),
pool_kernels=(1, 5, 11, 21, 45),
dilation_rate=99,
) # fv36
)
|