My model demonstrates exceptional performance when trained on subsets of up to 10 features but struggles to generalize effectively when utilizing the full dataset of 414 features. Would it be feasible and methodologically sound to design a wrapper framework that trains multiple smaller models, each using 10 features, to collectively capture the data's complexity? Additionally, is this a recommended approach in such scenarios?
# Weight normzlizatioin 22~23, Layer 24~25.

# Chomp1d removes extra time steps from the end of the sequence
class Chomp1d(nn.Module):
    def __init__(self, chomp_size):
        super(Chomp1d, self).__init__()
        self.chomp_size = chomp_size
    def forward(self, x):
        if self.chomp_size > 0:
            return x[:, :, :-self.chomp_size]
        else:
            return x

# Residual Block with causal convolutions and Chomp1d
class ResidualBlock(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, dilation, dropout):
        super(ResidualBlock, self).__init__()
        padding = (kernel_size - 1) * dilation

        self.conv1 = nn.utils.weight_norm(
            nn.Conv1d(in_channels, out_channels, kernel_size,
                      padding=padding, dilation=dilation)
        )
        self.chomp1 = Chomp1d(padding)
        self.relu1 = nn.ReLU()
        self.dropout1 = nn.Dropout(dropout)
        self.conv2 = nn.utils.weight_norm(
            nn.Conv1d(out_channels, out_channels, kernel_size,
                      padding=padding, dilation=dilation)
        )
        self.chomp2 = Chomp1d(padding)
        self.relu2 = nn.ReLU()
        self.dropout2 = nn.Dropout(dropout)
        # If the input and output channels are not the same, adjust with a 1x1 convolution
        self.downsample = nn.Conv1d(in_channels, out_channels, kernel_size=1) \
            if in_channels != out_channels else None

        self.relu = nn.ReLU()
        self.init_weights()

    def init_weights(self):
        nn.init.normal_(self.conv1.weight, std=0.01)
        nn.init.normal_(self.conv2.weight, std=0.01)
        if self.downsample is not None:
            nn.init.normal_(self.downsample.weight, std=0.01)

    def forward(self, x):
        residual = x  # Save the input for the residual connection
        out = self.conv1(x)
        out = self.chomp1(out)  # Remove extra padding
        out = self.relu1(out)
        out = self.dropout1(out)
        out = self.conv2(out)
        out = self.chomp2(out)  # Remove extra padding
        out = self.relu2(out)
        out = self.dropout2(out)
        if self.downsample is not None:
            residual = self.downsample(residual)

        return self.relu(out + residual)  # Residual connection with addition

# Updated TCN Encoder with Residual Blocks 0.4
class TCNEncoder(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, kernel_size, dropout):
        super(TCNEncoder, self).__init__()
        layers = []
        in_channels = input_size

        for i in range(num_layers):
            dilation = 2 ** i
            out_channels = hidden_size

            layers.append(ResidualBlock(in_channels, out_channels, kernel_size, dilation, dropout))

            in_channels = out_channels  # For next layer

        self.network = nn.Sequential(*layers)

    def forward(self, x):
        # x shape: (batch_size, seq_length, input_size)
        x = x.transpose(1, 2)  # Change to (batch_size, input_size, seq_length)
        x = self.network(x)
        x = x.transpose(1, 2)  # Back to (batch_size, seq_length, hidden_size)
        return x

    
# Transformer Decoder with masked multi-head attention
class TransformerDecoder(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, num_heads, dropout):
        super(TransformerDecoder, self).__init__()
        self.embedding = nn.Linear(input_size, hidden_size)
        self.positional_encoding = PositionalEncoding(hidden_size, dropout)
        decoder_layer = nn.TransformerDecoderLayer(d_model=hidden_size, nhead=num_heads, dropout=dropout)
        self.transformer_decoder = nn.TransformerDecoder(decoder_layer, num_layers=num_layers)
        self.fc_out = nn.Linear(hidden_size, input_size)

    def forward(self, tgt, memory, tgt_mask):
        # tgt: (batch_size, tgt_seq_length, input_size)
        tgt_emb = self.embedding(tgt)
        tgt_emb = self.positional_encoding(tgt_emb)
        memory = memory.transpose(0, 1)  # (src_seq_length, batch_size, hidden_size)
        tgt_emb = tgt_emb.transpose(0, 1)  # (tgt_seq_length, batch_size, hidden_size)
        output = self.transformer_decoder(tgt_emb, memory, tgt_mask=tgt_mask)
        output = output.transpose(0, 1)  # Back to (batch_size, tgt_seq_length, hidden_size)
        return output

    
# TCN-Encoder-Transformer-Decoder Model
class TCNEncoderTransformerDecoder(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, kernel_size, num_heads, dropout):
        super(TCNEncoderTransformerDecoder, self).__init__()
        self.encoder = TCNEncoder(input_size, hidden_size, num_layers, kernel_size, dropout)
        # Additional Attention Layer (Fusion Cross-Attention) ver = 0.4
        self.attention_layer = nn.MultiheadAttention(embed_dim=hidden_size, num_heads=num_heads, dropout=dropout)
        self.decoder = TransformerDecoder(input_size, hidden_size, num_layers, num_heads, dropout)
        self.positional_encoding = PositionalEncoding(hidden_size, dropout)  # For encoder outputs

        # Fusion Layer ver = 0.4 if Fusion startgy 1 then multiply hidden_size by 2
        self.final_fc_layer = nn.Linear(hidden_size, input_size)
        self.dropout = nn.Dropout(dropout)

        # Alpha parameter for weighted sum fusion ver = 0.4
        # Initialize alpha as a learnable parameter with value 0.5
        self.alpha = nn.Parameter(torch.full((hidden_size,), 0.5))

    def forward(self, src, tgt):
        # src: (batch_size, src_seq_length, input_size)
        # tgt: (batch_size, tgt_seq_length, input_size)

        # Encoder
        memory = self.encoder(src)  # (batch_size, src_seq_length, hidden_size)
        memory = self.positional_encoding(memory)  # Add positional encoding to memory

        # Decoder
        tgt_mask = nn.Transformer.generate_square_subsequent_mask(tgt.size(1)).to(tgt.device)  # (tgt_seq_length, tgt_seq_length)
        decoder_output = self.decoder(tgt, memory, tgt_mask)  # (batch_size, tgt_seq_length, hidden_size)

        # Reshape for MultiheadAttention (seq_length, batch_size, embed_dim)
        decoder_output_t = decoder_output.transpose(0, 1)  # (tgt_seq_length, batch_size, hidden_size)
        memory_t = memory.transpose(0, 1)  # (src_seq_length, batch_size, hidden_size)

        # Additional Attention Layer
        attn_output, _ = self.attention_layer(decoder_output_t, memory_t, memory_t)
        attn_output = attn_output.transpose(0, 1)  # (batch_size, tgt_seq_length, hidden_size)

        # Fusion Strategy 1: Concatenate decoder output and attention output
        #combined = torch.cat((decoder_output, attn_output), dim=-1)  # (batch_size, tgt_seq_length, hidden_size * 2)
        #combined = self.dropout(combined)#

        # Fusion Strategy 2: Weighted Sum of decoder output and attention output
        # I must always ensure alpha is in [0,1] using sigmoid or similar functions
        # Tested scaler alpha and it perfromed worse than 4% worse than Strategy 1 so switched it to vector
        alpha = torch.sigmoid(self.alpha).to(decoder_output.device).view(1, 1, -1)  # Shape: (1, 1, hidden_size)
        combined = alpha * decoder_output + (1 - alpha) * attn_output  # (batch_size, tgt_seq_length, hidden_size)
        combined = self.dropout(combined)

        # Fusion Strategy 2: Multiplication
        #final_output = self.multi_scale_fusion(decoder_output, memory)
        
        # Final Output Layer
        final_output = self.final_fc_layer(combined)  # (batch_size, tgt_seq_length, input_size)

        return final_output


# Positional Encoding class (unchanged)
class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout=0.1, max_len=5000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)
        
        pe = torch.zeros(max_len, d_model)  # Create positional encoding matrix
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-np.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)  # Apply sine to even indices
        pe[:, 1::2] = torch.cos(position * div_term)  # Apply cosine to odd indices
        pe = pe.unsqueeze(0)
        self.register_buffer('pe', pe)  # Save positional encoding matrix

    def forward(self, x):
        x = x + self.pe[:, :x.size(1), :].to(x.device)  # Add positional encoding
        return self.dropout(x)


# Learnable Positional Encoding class
class LearnablePositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000):
        super(LearnablePositionalEncoding, self).__init__()
        self.pe = nn.Parameter(torch.randn(1, max_len, d_model))

    def forward(self, x):
        x = x + self.pe[:, :x.size(1), :].to(x.device)
        return x