Fix parent_block.var(name) error in static mode for RNN

[0x45f]

PR types

Bug fixes

PR changes

Others

Describe

修复rnn在控制流中使用时，rnn调用parent_block.var(name)报错的问题。
问题描述：在静态图下网络参数都会在block0中，rnn的静态图逻辑中会调用parent_block.var去父block中找param，但是如果父block不是block不是block0则会报错。用户提供了如下的动转静代码，改用_find_var_recursive后动转静可以正常导出：

import paddle from paddle import nn import paddle.tensor as tensor import paddle.nn.functional as F import paddle.nn.initializer as I class LSTMCell(nn.RNNCellBase): def __init__(self, input_size: int, hidden_size: int, activation="tanh", weight_ih_attr=None, weight_hh_attr=None, bias_ih_attr=None, bias_hh_attr=None, guass_mean=0.0, guass_std=0.02, name=None): super(LSTMCell, self).__init__() self.weight_ih = self.create_parameter( (4 * hidden_size, input_size), weight_ih_attr, default_initializer=I.Normal(guass_mean, guass_std)) self.weight_hh = self.create_parameter( (4 * hidden_size, hidden_size), weight_hh_attr, default_initializer=I.Normal(guass_mean, guass_std)) self.bias_ih = self.create_parameter( (4 * hidden_size, ), bias_ih_attr, is_bias=True, default_initializer=I.Normal(guass_mean, guass_std)) self.bias_hh = self.create_parameter( (4 * hidden_size, ), bias_hh_attr, is_bias=True, default_initializer=I.Normal(guass_mean, guass_std)) self.hidden_size = hidden_size self.input_size = input_size self.gate_activation = F.sigmoid activation_dict = { 'tanh':paddle.tanh, 'relu':F.relu, 'gelu':F.gelu } if activation not in activation_dict: raise RuntimeError(f"{activation} is not supported in LSTMCell") self.activation = activation_dict[activation] def forward(self, inputs, states=None): # import pdb; pdb.set_trace() if states is None: states = self.get_initial_states(inputs, self.state_shape) prev_h, prev_c = states gates = paddle.matmul(inputs, self.weight_ih, transpose_y=True) if self.bias_ih is not None: gates = gates + self.bias_ih gates += paddle.matmul(prev_h, self.weight_hh, transpose_y=True) if self.bias_hh is not None: gates = gates + self.bias_hh chunked_gates = paddle.split(gates, num_or_sections=4, axis=-1) i = self.gate_activation(chunked_gates[0]) f = self.gate_activation(chunked_gates[1]) g = self.activation(chunked_gates[2]) o = self.gate_activation(chunked_gates[3]) c = f * prev_c + i * g h = o * self.activation(c) return h, (h, c) @property def state_shape(self): r"""  The `state_shape` of LSTMCell is a tuple with two shapes:   `((hidden_size, ), (hidden_size,))`. (-1 for batch size would be   automatically inserted into shape). These two shapes correspond   to :math:`h_{t-1}` and :math:`c_{t-1}` separately.  """ return ((self.hidden_size, ), (self.hidden_size, )) def extra_repr(self): return '{input_size}, {hidden_size}'.format(**self.__dict__) class Decoder(nn.Layer): def __init__(self, input_size, hidden_size): super().__init__() self.cell = LSTMCell(input_size, hidden_size) self.rnn = nn.RNN(self.cell) self.sos = paddle.ones(shape=[1, 1, 2], dtype='float32') self.init_states = (paddle.zeros(shape=[1, 1, 4], dtype='float32'), paddle.zeros(shape=[1, 1, 4], dtype='float32')) self.idx = paddle.zeros(shape=[1], dtype='int32') # self.idx = 0 self.states = self.rnn(self.sos, self.init_states) self.step = 0 def forward(self, inputs, hidden=None, cell=None): if hidden is None: states = self.states else: states = (hidden, cell) # import pdb; pdb.set_trace() if self.idx < 1: outs, states = self.rnn(inputs, states) self.idx += 1 final_states = states return outs, final_states def export(self): static_model = paddle.jit.to_static( self, input_spec=[ paddle.static.InputSpec( shape=[1, 1, 2], dtype='float32'), paddle.static.InputSpec( shape=[1, 1, 4], dtype='float32'), paddle.static.InputSpec( shape=[1, 1, 4], dtype='float32') ] ) return static_model model = Decoder(2, 4) model.eval() static_model = model.export() paddle.jit.save(static_model, "test_model")

[Aurelius84]

LGTM

[iclementine]

LGTM

[guoshengCS]

LGTM