时间衰减+biGRU pytorch代码

以下是一个 biGRU 模型的 PyTorch 代码，其中包含时间衰减的实现。在此代码中，时间衰减通过计算每个时间步的权重来实现，这些权重随着时间步数的增加而减小。这种权重的计算方法是通过将时间步数作为指数输入到一个指数函数中得到的。

import torch
import torch.nn as nn

class BiGRU(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim, n_layers, dropout, decay_rate):
        super(BiGRU, self).__init__()
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        self.output_dim = output_dim
        self.n_layers = n_layers
        self.dropout = dropout
        self.decay_rate = decay_rate
        
        self.gru = nn.GRU(input_dim, hidden_dim, n_layers, dropout=dropout, bidirectional=True)
        self.fc = nn.Linear(hidden_dim*2, output_dim)

    def forward(self, x, seq_lengths):
        # sort inputs by sequence length in descending order
        seq_lengths, idx_sort = torch.sort(seq_lengths, dim=0, descending=True)
        x = x[idx_sort]
        idx_unsort = torch.argsort(idx_sort, dim=0)
        
        # calculate time decay weights
        max_length = x.size(1)
        decay_weights = torch.pow(torch.tensor(self.decay_rate), torch.arange(max_length).float())
        decay_weights = decay_weights.unsqueeze(0).repeat(x.size(0), 1).to(device=x.device)
        x = x * decay_weights.unsqueeze(2)
        
        # pack padded sequence
        x = nn.utils.rnn.pack_padded_sequence(x, seq_lengths, batch_first=True)
        
        # pass through GRU
        _, h = self.gru(x)
        
        # concatenate forward and backward hidden states
        h = torch.cat((h[-2,:,:], h[-1,:,:]), dim=1)
        
        # unsort hidden states
        h = h[idx_unsort]
        
        # pass through linear layer
        out = self.fc(h)
        
        return out

在 `forward` 方法中，首先对输入按序列长度进行排序，并记录排序后的索引。然后计算时间衰减权重，并将其与输入相乘。接下来，使用 `nn.utils.rnn.pack_padded_sequence` 方法将输入打包，以便能够在 GRU 中有效地处理变长序列。最后，将前向和后向隐藏状态拼接在一起，并将其传递到全连接层以获得最终输出。

注意，在这个实现中，我们假设输入的第一个维度是 batch size，第二个维度是时间步数，第三个维度是输入特征的数量。此外，我们还假设每个 batch 中的输入长度已经按降序排序。如果这些假设不成立，可能需要进行一些额外的处理才能正确实现时间衰减。