Analyze my timeseries forecasting model and suggest ways to optimize it


TCN Encoder
class TCNEncoder(nn.Module):
def init(self, input_size, hidden_size, num_layers, kernel_size, dropout):
super(TCNEncoder, self).init()
layers = []
num_levels = num_layers


    in_channels = input_size
    out_channels = hidden_size

    for i in range(num_levels):
        dilation_size = 2 ** i
        padding = (kernel_size - 1) * dilation_size
        conv = nn.Conv1d(in_channels, out_channels, kernel_size, padding=padding, dilation=dilation_size)
        conv = nn.utils.parametrizations.weight_norm(conv)
        layers += [conv, nn.ReLU(), nn.Dropout(dropout)]
        in_channels = hidden_size  # For next layer

    self.network = nn.Sequential(*layers)
    self.kernel_size = kernel_size
    self.num_levels = num_levels

def forward(self, x):
    # x shape: (batch_size, seq_length, input_size)
    x = x.permute(0, 2, 1)  # Reshape input for Conv1d: (batch_size, input_size, seq_length)
    x = self.network(x)
    total_padding = (self.kernel_size - 1) * (2 ** self.num_levels - 1)
    x = x[:, :, :-total_padding] if total_padding != 0 else x  # Remove extra padding
    x = x.permute(0, 2, 1)  # 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)
    output = self.fc_out(output)
    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)
self.decoder = TransformerDecoder(input_size, hidden_size, num_layers, num_heads, dropout)
self.positional_encoding = PositionalEncoding(hidden_size, dropout) # For encoder outputs


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)
    output = self.decoder(tgt, memory, tgt_mask)
    return 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)