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)[source]

用于执行多层 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{array}{r} \begin{array}{ll} r_{t} = σ (W_{i r} x_{t} + b_{i r} + W_{h r} h_{(t - 1)} + b_{h r}) \\ z_{t} = σ (W_{i z} x_{t} + b_{i z} + W_{h z} h_{(t - 1)} + b_{h z}) \\ n_{t} = \tanh (W_{i n} x_{t} + b_{i n} + r_{t} ⊙ (W_{h n} h_{(t - 1)} + b_{h n})) \\ h_{t} = (1 - z_{t}) ⊙ n_{t} + z_{t} ⊙ h_{(t - 1)} \end{array} \end{array}

其中 $h_{t}$ 是时间步 t 的隐藏状态， $x_{t}$ 是时间步 t 的输入， $h_{(t - 1)}$ 是时间步 t-1 的层的隐藏状态，或者时间步 0 的初始隐藏状态，并且 $r_{t}$ 、 $z_{t}$ 、 $n_{t}$ 分别是重置门、更新门和新门。 $σ$ 是 sigmoid 函数， $⊙$ 是 Hadamard 积。

在多层 GRU 中，第 $l$ 层的输入 $x_{t}^{(l)}$ ( $l \geq 2$ ) 是前一层的隐藏状态 $h_{t}^{(l - 1)}$ 乘以 dropout $δ_{t}^{(l - 1)}$ ，其中每个 $δ_{t}^{(l - 1)}$ 都是一个伯努利随机变量，其概率为 dropout 时为 $0$ 。

参数:

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_{i n})$ 的张量（用于非批量输入），当 batch_first=False 时为 $(L, N, H_{i n})$ ，或者当 batch_first=True 时为 $(N, L, H_{i n})$ ，包含输入序列的特征。输入也可以是打包的可变长度序列。有关详细信息，请参阅 torch.nn.utils.rnn.pack_padded_sequence() 或 torch.nn.utils.rnn.pack_sequence()。
h_0: 形状为 $(D * num\_layers, H_{o u t})$ 或 $(D * num\_layers, N, H_{o u t})$ 的张量，包含输入序列的初始隐藏状态。如果未提供，则默认为零。

其中

\begin{array}{r} \begin{aligned} N = & batch size \\ L = & sequence length \\ D = & 2 if bidirectional=True otherwise 1 \\ H_{i n} = & input\_size \\ H_{o u t} = & hidden\_size \end{aligned} \end{array}

输出：output, h_n

output: 形状为 $(L, D * H_{o u t})$ 的张量（用于非批量输入），当 batch_first=False 时为 $(L, N, D * H_{o u t})$ ，或者当 batch_first=True 时为 $(N, L, D * H_{o u t})$ ，包含来自 GRU 最后一层的输出特征 (h_t)，对于每个 t。如果 torch.nn.utils.rnn.PackedSequence 已作为输入给出，则输出也将是一个打包序列。
h_n: 形状为 $(D * num\_layers, H_{o u t})$ 或 $(D * num\_layers, N, H_{o u t})$ 的张量，包含输入序列的最终隐藏状态。

变量:

weight_ih_l[k] – 第 $k^{t h}$ 层的可学习输入到隐藏层权重 (W_ir|W_iz|W_in)，对于 k = 0，形状为 (3*hidden_size, input_size)。否则，形状为 (3*hidden_size, num_directions * hidden_size)
weight_hh_l[k] – 第 $k^{t h}$ 层的可学习隐藏层到隐藏层权重 (W_hr|W_hz|W_hn)，形状为 (3*hidden_size, hidden_size)
bias_ih_l[k] – 第 $k^{t h}$ 层的可学习输入到隐藏层偏置 (b_ir|b_iz|b_in)，形状为 (3*hidden_size)
bias_hh_l[k] – 第 $k^{t h}$ 层的可学习隐藏层到隐藏层偏置 (b_hr|b_hz|b_hn)，形状为 (3*hidden_size)

注意

所有权重和偏置都从 $U (- \sqrt{k}, \sqrt{k})$ 初始化，其中 $k = \frac{1}{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 积 $(⊙)$ 在与权重矩阵 W 相乘并添加偏置之前完成

\begin{array}{r} n_{t} = \tanh (W_{i n} x_{t} + b_{i n} + W_{h n} (r_{t} ⊙ h_{(t - 1)}) + b_{h n}) \end{array}

这与 PyTorch 实现形成对比，PyTorch 实现是在 $W_{h n} h_{(t - 1)}$ 之后完成的

\begin{array}{r} n_{t} = \tanh (W_{i n} x_{t} + b_{i n} + r_{t} ⊙ (W_{h n} h_{(t - 1)} + b_{h n})) \end{array}

此实现出于效率目的而有所不同。

示例

>>> 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)[source]

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

应由所有子类重写。

注意

虽然 forward 传递的配方需要在该函数中定义，但应在此之后调用 Module 实例而不是此函数，因为前者负责运行已注册的钩子，而后者会静默忽略它们。

GRU

文档

教程

资源