Multivariate time-series transformer

Michelangelo_Balzano: <p>Hi everyone. I’m trying to implement a transformer for pre-training on partially masked multivariate time-series data. The model input is of the type (batch size, num signals, seq len), where there are 3 signals and they all have the same length (240 time steps). The loss function used is the RMSE between only the masked values ​​and the corresponding predictions. The problem is that the predictions are totally wrong. This is my model:</p>
<pre><code class="lang-auto">class LearnablePositionalEncoding(nn.Module):

    def __init__(self, segment_length, d_model, dropout):
        super(LearnablePositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)
        self.pe = nn.Parameter(torch.empty(segment_length, 1, d_model))
        nn.init.uniform_(self.pe, -0.02, 0.02)

    def forward(self, x):

        x = x + self.pe[:x.size(0), :]
        return self.dropout(x)

class MyEncoderLayer(nn.modules.Module):
    
    def __init__(self, d_model, nhead, dim_feedforward, dropout=0.1):
        super(MyEncoderLayer, self).__init__()
        self.self_attn = MultiheadAttention(d_model, nhead, dropout=dropout)
        
        self.linear1 = Linear(d_model, dim_feedforward)
        self.dropout = Dropout(dropout)
        self.linear2 = Linear(dim_feedforward, d_model)

        self.norm1 = BatchNorm1d(d_model, eps=1e-5)
        self.norm2 = BatchNorm1d(d_model, eps=1e-5)
        self.dropout1 = Dropout(dropout)
        self.dropout2 = Dropout(dropout)

        self.activation = F.gelu

    def forward(self, data: Tensor, src_mask: Optional[Tensor] = None,
                src_key_padding_mask: Optional[Tensor] = None, is_causal: Optional[bool] = None) -&gt; Tensor:
        
        # Self-attention
        data2 = self.self_attn(data, data, data, attn_mask=src_mask, key_padding_mask=src_key_padding_mask, is_causal=is_causal)[0]
        # Add &amp; Norm
        data = data + self.dropout1(data2)  # [segment_length, batch_size, d_model]
        data = data.permute(1, 2, 0)  # [batch_size, d_model, segment_length]
        data = self.norm1(data)
        data = data.permute(2, 0, 1)  # [segment_length, batch_size, d_model]
        # Feed-forward
        data2 = self.linear2(self.dropout(self.activation(self.linear1(data))))
        # Add &amp; Norm
        data = data + self.dropout2(data2)  # [segment_length, batch_size, d_model]
        data = data.permute(1, 2, 0)  # [batch_size, d_model, segment_length]
        data = self.norm2(data)
        data = data.permute(2, 0, 1)  # [segment_length, batch_size, d_model]

        return data

class TSTransformer(nn.Module):
    def __init__(self, num_signals, segment_length, iperparametri, device, pe_type='learnable'):
        super(TSTransformer, self).__init__()

        self.num_signals = num_signals
        self.segment_length = segment_length
        self.d_model = iperparametri['d_model']
        self.dropout = iperparametri['dropout']
        self.num_layers = iperparametri['num_layers']
        self.num_heads = iperparametri['num_heads']
        self.device = device

        self.project_inp = nn.Linear(self.num_signals, self.d_model)
        if pe_type == 'learnable':
            self.pos_enc = LearnablePositionalEncoding(self.segment_length, self.d_model, self.dropout)
        elif pe_type == 'fixed':
            self.pos_enc = FixedPositionalEncoding(self.segment_length, self.d_model, self.dropout, self.device)
        encoder_layer = MyEncoderLayer(self.d_model, self.num_heads, self.d_model, self.dropout)
        self.transformer_encoder = nn.TransformerEncoder(encoder_layer, self.num_layers)
        self.output_layer = nn.Linear(self.d_model, self.num_signals)
        self.act = F.gelu
        self.dropout1 = nn.Dropout(self.dropout)
        
    def forward(self, data):

        # [batch_size, num_signals, segment_length]
        inp = data.permute(2, 0, 1) # [segment_length, batch_size, num_signals]
        inp = self.project_inp(inp) * math.sqrt(self.d_model) # [segment_length, batch_size, d_model]
        inp = self.pos_enc(inp) # Aggiunta del positional encoding
        output = self.transformer_encoder(inp) # [segment_length, batch_size, d_model]
        output = self.act(output) # Funzione di attivazione
        output = output.permute(1, 0, 2)  # [batch_size, segment_length, d_model]
        output = self.dropout1(output) # Dropout
        output = self.output_layer(output) # [batch_size, segment_length, num_signals]
        output = output.permute(0, 2, 1) # [batch_size, num_signals, segment_length]

        return output # [batch_size, num_signals, segment_length]
</code></pre>
<p>For each epoch I carry out a training activity and a validation activity as shown below:</p>
<pre><code class="lang-auto">def train_pretrain_model(model, dataloader, num_signals, segment_length, iperparametri, optimizer, device):
    
    model.train()
    train_loss = 0.0
    num_batches = len(dataloader)

    progress_bar = tqdm(total=num_batches, desc="Train batch analizzati")
    for batch in dataloader:
        
        optimizer.zero_grad()
        masks = generate_masks(iperparametri['batch_size'], iperparametri['masking_ratio'], 
                               iperparametri['lm'], num_signals, segment_length, device) 
        masked_batch = batch * masks
        predictions = model(masked_batch) 
        loss = pretraining_loss(predictions, batch, masks)
        loss.backward() 
        optimizer.step()
        train_loss += loss.item() 

        progress_bar.update(1)
    progress_bar.close()

    return train_loss / num_batches, model

def validate_pretrain_model(model, dataloader, num_signals, segment_length, iperparametri, device):
    
    model.eval()
    val_loss = 0.0
    num_batches = len(dataloader)

    with torch.no_grad():
        
        progress_bar = tqdm(total=num_batches, desc="Val batch analizzati")
        for batch in dataloader:
            
            masks = generate_masks(iperparametri['batch_size'], iperparametri['masking_ratio'], 
                                   iperparametri['lm'], num_signals, segment_length, device)
            masked_batch = batch * masks 
            predictions = model(masked_batch) 
            loss = pretraining_loss(predictions, batch, masks) 
            val_loss += loss.item() 


            progress_bar.update(1)
        progress_bar.close()

    return val_loss / num_batches, model
</code></pre>
<p>The next image describes how the model’s predictions are incorrect. The blue lines are the actual values, the red lines are the predictions, the gray areas are the masked sections.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/8/2/82e33cb43cea441fd8826adfbbe0677c8989b684.png" data-download-href="https://discuss.pytorch.org/uploads/default/82e33cb43cea441fd8826adfbbe0677c8989b684" title="Figure_1"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/8/2/82e33cb43cea441fd8826adfbbe0677c8989b684_2_550x500.png" alt="Figure_1" data-base62-sha1="iFSVjBdZMnpUNXjBoRpOzwGNvWA" width="550" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/8/2/82e33cb43cea441fd8826adfbbe0677c8989b684_2_550x500.png, https://discuss.pytorch.org/uploads/default/optimized/3X/8/2/82e33cb43cea441fd8826adfbbe0677c8989b684_2_825x750.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/8/2/82e33cb43cea441fd8826adfbbe0677c8989b684.png 2x" data-dominant-color="F0EFF0"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Figure_1</span><span class="informations">880×800 55.4 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>