时间衰减+GRU pytorch代码

下面是一个简单的时间衰减GRU模型的PyTorch代码：

import torch
import torch.nn as nn

class TimeDecayGRU(nn.Module):

    def __init__(self, input_size, hidden_size, time_decay=0.1):
        super(TimeDecayGRU, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.time_decay = time_decay

        self.weight_ih = nn.Parameter(torch.Tensor(3 * hidden_size, input_size))
        self.weight_hh = nn.Parameter(torch.Tensor(3 * hidden_size, hidden_size))
        self.bias_ih = nn.Parameter(torch.Tensor(3 * hidden_size))
        self.bias_hh = nn.Parameter(torch.Tensor(3 * hidden_size))

        self.reset_parameters()

    def reset_parameters(self):
        nn.init.kaiming_uniform_(self.weight_ih, a=math.sqrt(5))
        nn.init.kaiming_uniform_(self.weight_hh, a=math.sqrt(5))
        nn.init.zeros_(self.bias_ih)
        nn.init.zeros_(self.bias_hh)

    def forward(self, input, hidden, time_diff):
        gi = torch.matmul(self.weight_ih, input.t())
        gh = torch.matmul(self.weight_hh, hidden.t())
        bias = self.bias_ih + self.bias_hh
        i_r, i_i, i_n = gi.chunk(3)
        h_r, h_i, h_n = gh.chunk(3)
        reset_gate = torch.sigmoid(i_r + h_r + bias[:self.hidden_size])
        input_gate = torch.sigmoid(i_i + h_i + bias[self.hidden_size:2 * self.hidden_size])
        new_input = torch.tanh(i_n + reset_gate * h_n + bias[2 * self.hidden_size:])
        hidden = (1 - input_gate) * hidden + input_gate * new_input * torch.exp(-time_diff * self.time_decay)
        return hidden

    def init_hidden(self, batch_size):
        return torch.zeros(batch_size, self.hidden_size)

该模型继承自PyTorch的`nn.Module`类，其中包含了输入大小，隐藏状态大小和时间衰减值作为超参数。在`__init__`方法中，定义了GRU层的参数：输入权重，隐藏状态权重和偏置项。`reset_parameters`方法用于对这些参数进行初始化。

`forward`方法实现了GRU层的前向传播，并且使用时间衰减计算当前时刻的隐藏状态。在实现中，首先计算输入和隐藏状态的线性组合，然后分别计算重置门，输入门和新输入。最后根据时间衰减值更新隐藏状态，并将其返回。

`init_hidden`方法用于初始化初始隐藏状态。它接受一个批量大小，返回一个形状为(batch_size, hidden_size)的全零张量。

该模型可以通过以下方式进行实例化和使用：

model = TimeDecayGRU(input_size=10, hidden_size=20, time_decay=0.1)
input = torch.randn(5, 10)  # batch_size=5, input_size=10
hidden = model.init_hidden(5)  # batch_size=5, hidden_size=20
time_diff = torch.FloatTensor([1, 2, 3, 4, 5])  # time differences
output = model(input, hidden, time_diff)