GRU

class torchrl.modules.GRU(input_size: int, hidden_size: int, num_layers: int = 1, bias: bool =True, batch_first: bool =True, dropout: float =0.0, bidirectional: bool =False, device=None, dtype=None)

一个用于执行多层 GRU 多步操作的 PyTorch 模块。该模块的行为与 torch.nn.GRU 完全一致，但此实现完全使用 Python 编写。

注意

此类实现不依赖 CuDNN，这使其与 torch.vmap() 和 torch.compile() 兼容。

示例

>>> import torch
>>> from torchrl.modules.tensordict_module.rnn import GRU

>>> device = torch.device("cuda") if torch.cuda.device_count() else torch.device("cpu")
>>> B = 2
>>> T = 4
>>> N_IN = 10
>>> N_OUT = 20
>>> N_LAYERS = 2
>>> V = 4  # vector size
>>> gru = GRU(
...     input_size=N_IN,
...     hidden_size=N_OUT,
...     device=device,
...     num_layers=N_LAYERS,
... )

# 单次调用 >>> x = torch.randn(B, T, N_IN, device=device) >>> h0 = torch.zeros(N_LAYERS, B, N_OUT, device=device) >>> with torch.no_grad(): … h1 = gru(x, h0)

# 向量化调用 - nn.GRU 不支持 >>> def call_gru(x, h): … h_out = gru(x, h) … return h_out >>> batched_call = torch.vmap(call_gru) >>> x = torch.randn(V, B, T, 10, device=device) >>> h0 = torch.zeros(V, N_LAYERS, B, N_OUT, device=device) >>> with torch.no_grad(): … h1 = batched_call(x, h0)

__init__(input_size,hidden_size,num_layers=1,bias=True,batch_first=False,dropout=0.0,bidirectional=False,device=None,dtype=None)

将多层门控循环单元 (GRU) RNN 应用于输入序列。对于输入序列中的每个元素，每层计算以下函数

\[\begin{split}\begin{array}{ll} r_t = \sigma(W_{ir} x_t + b_{ir} + W_{hr} h_{(t-1)} + b_{hr}) \\ z_t = \sigma(W_{iz} x_t + b_{iz} + W_{hz} h_{(t-1)} + b_{hz}) \\ n_t = \tanh(W_{in} x_t + b_{in} + r_t \odot (W_{hn} h_{(t-1)}+ b_{hn})) \\ h_t = (1 - z_t) \odot n_t + z_t \odot h_{(t-1)} \end{array}\end{split}\]

其中 \(h_t\) 是时间 t 的隐藏状态，\(x_t\) 是时间 t 的输入，\(h_{(t-1)}\) 是时间 t-1 的层隐藏状态或时间 0 的初始隐藏状态，而 \(r_t\)、\(z_t\)、\(n_t\) 分别是重置门、更新门和新门。\(\sigma\) 是 sigmoid 函数，\(\odot\) 是 Hadamard 积。

在多层 GRU 中，第 \(l\) 层 (\(l \ge 2\)) 的输入 \(x^{(l)}_t\) 是前一层隐藏状态 \(h^{(l-1)}_t\) 乘以 dropout \(\delta^{(l-1)}_t\)，其中每个 \(\delta^{(l-1)}_t\) 是一个 Bernoulli 随机变量，其取值为 \(0\) 的概率为 dropout。

参数:

• input_size – 输入 x 中预期的特征数量

• hidden_size – 隐藏状态 h 中的特征数量

• num_layers – 循环层的数量。例如，设置 num_layers=2 意味着将两个 GRU 堆叠在一起形成一个 堆叠 GRU，其中第二个 GRU 接收第一个 GRU 的输出并计算最终结果。默认值：1

• bias – 如果为 False，则该层不使用偏置权重 b_ih 和 b_hh。默认值：True

• batch_first – 如果为 True，则输入和输出张量以 (batch, seq, feature) 的形式提供，而不是 (seq, batch, feature)。请注意，这不适用于隐藏状态或单元状态。有关详细信息，请参阅下面的输入/输出部分。默认值：False

• dropout – 如果非零，则在除最后一层之外的每个 GRU 层的输出上引入一个 Dropout 层，其 dropout 概率等于 dropout。默认值：0

• bidirectional – 如果为 True，则成为双向 GRU。默认值：False

输入: input, h_0

• input: 对于无批次输入，形状为 \((L, H_{in})\)；当 batch_first=False 时，形状为 \((L, N, H_{in})\)；当 batch_first=True 时，形状为 \((N, L, H_{in})\) 的张量，包含输入序列的特征。输入也可以是打包的可变长度序列。有关详细信息，请参阅 torch.nn.utils.rnn.pack_padded_sequence() 或 torch.nn.utils.rnn.pack_sequence()。

• h_0: 形状为 \((D * \text{num\_layers}, H_{out})\) 或 \((D * \text{num\_layers}, N, H_{out})\) 的张量，包含输入序列的初始隐藏状态。如果未提供，默认为零。

其中

\[\begin{split}\begin{aligned} N ={} & \text{批次大小} \\ L ={} & \text{序列长度} \\ D ={} & 2 \text{ 如果 bidirectional=True 否则为 } 1 \\ H_{in} ={} & \text{input\_size} \\ H_{out} ={} & \text{hidden\_size} \end{aligned}\end{split}\]

输出: output, h_n

• output: 对于无批次输入，形状为 \((L, D * H_{out})\)；当 batch_first=False 时，形状为 \((L, N, D * H_{out})\)；当 batch_first=True 时，形状为 \((N, L, D * H_{out})\) 的张量，包含 GRU 最后一层在每个时间步 t 的输出特征 (h_t)。如果输入是 torch.nn.utils.rnn.PackedSequence，则输出也将是打包序列。

• h_n: 形状为 \((D * \text{num\_layers}, H_{out})\) 或 \((D * \text{num\_layers}, N, H_{out})\) 的张量，包含输入序列的最终隐藏状态。

变量:

• weight_ih_l[k] – 第 \(\text{k}^{th}\) 层的可学习输入到隐藏权重 (W_ir|W_iz|W_in)，对于 k = 0，形状为 (3*hidden_size, input_size)。否则，形状为 (3*hidden_size, num_directions * hidden_size)

• weight_hh_l[k] – 第 \(\text{k}^{th}\) 层的可学习隐藏到隐藏权重 (W_hr|W_hz|W_hn)，形状为 (3*hidden_size, hidden_size)

• bias_ih_l[k] – 第 \(\text{k}^{th}\) 层的可学习输入到隐藏偏置 (b_ir|b_iz|b_in)，形状为 (3*hidden_size)

• bias_hh_l[k] – 第 \(\text{k}^{th}\) 层的可学习隐藏到隐藏偏置 (b_hr|b_hz|b_hn)，形状为 (3*hidden_size)

注意

所有权重和偏置均从 \(\mathcal{U}(-\sqrt{k}, \sqrt{k})\) 初始化，其中 \(k = \frac{1}{\text{hidden\_size}}\)

注意

对于双向 GRU，前向和后向分别是方向 0 和 1。当 batch_first=False 时，分割输出层的示例：output.view(seq_len, batch, num_directions, hidden_size)。

注意

对于无批次输入，batch_first 参数将被忽略。

注意

新门 \(n_t\) 的计算与原始论文和其他框架略有不同。在原始实现中，\(r_t\) 和前一个隐藏状态 \(h_{(t-1)}\) 之间的 Hadamard 积 \((\odot)\) 是在与权重矩阵 W 相乘并加上偏置之前进行的

\[\begin{aligned} n_t = \tanh(W_{in} x_t + b_{in} + W_{hn} ( r_t \odot h_{(t-1)} ) + b_{hn}) \end{aligned}\]

这与 PyTorch 实现不同，后者在 \(W_{hn} h_{(t-1)}\) 之后进行

\[\begin{aligned} n_t = \tanh(W_{in} x_t + b_{in} + r_t \odot (W_{hn} h_{(t-1)}+ b_{hn})) \end{aligned}\]

出于效率考虑，此实现故意有所不同。

示例

>>> rnn = nn.GRU(10, 20, 2)
>>> input = torch.randn(5, 3, 10)
>>> h0 = torch.randn(2, 3, 20)
>>> output, hn = rnn(input, h0)

forward(input, hx=None)

定义每次调用时执行的计算。

应由所有子类覆盖。

注意

尽管前向传播的逻辑需要在该函数中定义，但之后应该调用 Module 实例而不是此函数，因为前者会处理运行注册的 hook，而后者会静默忽略它们。