Spaces:
Sleeping
Sleeping
| import torch | |
| import torch.nn as nn | |
| from .Modules.conformer import ConformerEncoder, ConformerDecoder | |
| from .Modules.mhsa_pro import RotaryEmbedding, ContinuousRotaryEmbedding | |
| from .kan.fasterkan import FasterKAN | |
| from kan import KAN | |
| class ConvBlock(nn.Module): | |
| def __init__(self, args, num_layer) -> None: | |
| super().__init__() | |
| if args.activation == 'silu': | |
| self.activation = nn.SiLU() | |
| else: | |
| self.activation = nn.ReLU() | |
| in_channels = args.encoder_dims[num_layer-1] if num_layer < len(args.encoder_dims) else args.encoder_dims[-1] | |
| out_channels = args.encoder_dims[num_layer] if num_layer < len(args.encoder_dims) else args.encoder_dims[-1] | |
| self.layers = nn.Sequential( | |
| nn.Conv1d(in_channels=in_channels, | |
| out_channels=out_channels, | |
| kernel_size=args.kernel_size, | |
| stride=1, padding='same', bias=False), | |
| nn.BatchNorm1d(num_features=out_channels), | |
| self.activation, | |
| ) | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| return self.layers(x) | |
| class CNNEncoder(nn.Module): | |
| def __init__(self, args) -> None: | |
| super().__init__() | |
| print("Using CNN encoder wit activation: ", args.activation, 'args avg_output: ', args.avg_output) | |
| if args.activation == 'silu': | |
| self.activation = nn.SiLU() | |
| else: | |
| self.activation = nn.ReLU() | |
| self.embedding = nn.Sequential(nn.Conv1d(in_channels = args.in_channels, | |
| kernel_size=3, out_channels = args.encoder_dims[0], stride=1, padding = 'same', bias = False), | |
| nn.BatchNorm1d(args.encoder_dims[0]), | |
| self.activation, | |
| ) | |
| self.layers = nn.ModuleList([ConvBlock(args, i+1) | |
| for i in range(args.num_layers)]) | |
| self.pool = nn.MaxPool1d(2) | |
| self.output_dim = args.encoder_dims[-1] | |
| self.min_seq_len = 2 | |
| self.avg_output = args.avg_output | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| if len(x.shape)==2: | |
| x = x.unsqueeze(1) | |
| if len(x.shape)==3 and x.shape[-1]==1: | |
| x = x.permute(0,2,1) | |
| x = self.embedding(x) | |
| for m in self.layers: | |
| x = m(x) | |
| if x.shape[-1] > self.min_seq_len: | |
| x = self.pool(x) | |
| if self.avg_output: | |
| x = x.mean(dim=-1) | |
| return x | |
| class MultiEncoder(nn.Module): | |
| def __init__(self, args, conformer_args): | |
| super().__init__() | |
| self.backbone = CNNEncoder(args) | |
| self.backbone.avg_output = False | |
| self.head_size = conformer_args.encoder_dim // conformer_args.num_heads | |
| self.rotary_ndims = int(self.head_size * 0.5) | |
| self.pe = RotaryEmbedding(self.rotary_ndims) | |
| self.encoder = ConformerEncoder(conformer_args) | |
| self.output_dim = conformer_args.encoder_dim | |
| self.avg_output = args.avg_output | |
| def forward(self, x): | |
| # Store backbone output in a separate tensor | |
| backbone_out = self.backbone(x) | |
| # Create x_enc from backbone_out | |
| if len(backbone_out.shape) == 2: | |
| x_enc = backbone_out.unsqueeze(1).clone() | |
| else: | |
| x_enc = backbone_out.permute(0,2,1).clone() | |
| RoPE = self.pe(x_enc, x_enc.shape[1]) | |
| x_enc = self.encoder(x_enc, RoPE) | |
| if len(x_enc.shape) == 3: | |
| if self.avg_output: | |
| x_enc = x_enc.sum(dim=1) | |
| else: | |
| x_enc = x_enc.permute(0,2,1) | |
| # Return x_enc and the original backbone output | |
| return x_enc, backbone_out | |
| class DualEncoder(nn.Module): | |
| def __init__(self, args_x, args_f, conformer_args) -> None: | |
| super().__init__() | |
| self.encoder_x = CNNEncoder(args_x) | |
| self.encoder_f = MultiEncoder(args_f, conformer_args) | |
| total_output_dim = args_x.encoder_dims[-1] + args_f.encoder_dims[-1] | |
| self.regressor = nn.Sequential( | |
| nn.Linear(total_output_dim, total_output_dim//2), | |
| nn.BatchNorm1d(total_output_dim//2), | |
| nn.SiLU(), | |
| nn.Linear(total_output_dim//2, 1) | |
| ) | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| x1 = self.encoder_x(x) | |
| x2, _ = self.encoder_f(x) | |
| logits = torch.cat([x1, x2], dim=-1) | |
| return self.regressor(logits).squeeze() | |
| class CNNKan(nn.Module): | |
| def __init__(self, args, conformer_args, kan_args): | |
| super().__init__() | |
| self.backbone = CNNEncoder(args) | |
| # self.kan = KAN(width=kan_args['layers_hidden']) | |
| self.kan = FasterKAN(**kan_args) | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| x = self.backbone(x) | |
| x = x.mean(dim=1) | |
| return self.kan(x) | |
| class KanEncoder(nn.Module): | |
| def __init__(self, args): | |
| super().__init__() | |
| self.kan_x = FasterKAN(**args) | |
| self.kan_f = FasterKAN(**args) | |
| self.kan_out = FasterKAN(layers_hidden=[args['layers_hidden'][-1]*2, 8,8,1]) | |
| def forward(self, x: torch.Tensor, f: torch.Tensor) -> torch.Tensor: | |
| x = self.kan_x(x) | |
| f = self.kan_f(f) | |
| out = torch.cat([x, f], dim=-1) | |
| return self.kan_out(out) | |