| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, |
| # software distributed under the License is distributed on an |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| # KIND, either express or implied. See the License for the |
| # specific language governing permissions and limitations |
| # under the License. |
| |
| import mxnet as mx |
| from mxnet import gluon, np |
| import numpy as _np |
| import copy |
| from functools import partial |
| from numpy.testing import assert_allclose |
| from mxnet.gluon.utils import split_rnn_params |
| import pytest |
| from mxnet.test_utils import almost_equal, assert_almost_equal, default_device |
| from common import assert_raises_cudnn_not_satisfied, retry |
| |
| |
| def check_rnn_states(fused_states, stack_states, num_layers, bidirectional=False, is_lstm=True): |
| directions = 2 if bidirectional else 1 |
| assert len(stack_states) / len(fused_states) == num_layers * directions |
| |
| fused_states = [state.asnumpy() for state in fused_states] |
| stack_states = [_np.expand_dims(state.asnumpy(), axis=0) for state in stack_states] |
| if is_lstm: |
| stack_states_h = stack_states[0::2] |
| stack_states_c = stack_states[1::2] |
| stack_states = [_np.concatenate(stack_states_h, axis=0), _np.concatenate(stack_states_c, axis=0)] |
| else: |
| stack_states = [_np.concatenate(stack_states, axis=0)] |
| |
| for f, s in zip(fused_states, stack_states): |
| assert f.shape == s.shape |
| assert_almost_equal(f, s, atol=1e-4, rtol=1e-4) |
| |
| |
| @mx.util.use_np |
| def test_rnn(): |
| cell = gluon.rnn.RNNCell(100) |
| inputs = [mx.np.ones((10,50)) for i in range(3)] |
| cell.infer_shape(0, inputs[0], False) |
| cell.initialize() |
| outputs, _ = cell.unroll(3, inputs) |
| assert sorted(cell.collect_params().keys()) == ['h2h_bias', 'h2h_weight', |
| 'i2h_bias', 'i2h_weight'] |
| |
| outs = [o.shape for o in outputs] |
| assert outs == [(10, 100), (10, 100), (10, 100)] |
| |
| |
| @mx.util.use_np |
| def test_lstm(): |
| cell = gluon.rnn.LSTMCell(100) |
| inputs = [mx.np.ones((10,50)) for i in range(3)] |
| cell.infer_shape(0, inputs[0], False) |
| cell.initialize() |
| outputs, _ = cell.unroll(3, inputs) |
| assert sorted(cell.collect_params().keys()) == ['h2h_bias', 'h2h_weight', 'i2h_bias', 'i2h_weight'] |
| |
| outs = [o.shape for o in outputs] |
| assert outs == [(10, 100), (10, 100), (10, 100)] |
| |
| |
| @assert_raises_cudnn_not_satisfied(min_version='7.2.1') |
| @pytest.mark.serial |
| def test_lstmp(): |
| hidden_size, projection_size = 512, 256 |
| rtol, atol = 1e-4, 1e-4 |
| batch_size, seq_len = 5, 3 |
| input_size = 128 |
| lstm_input = mx.np.random.uniform(size=(seq_len, batch_size, input_size)) |
| |
| # ==== Unidirectional Layer ==== |
| for num_layers in [1, 3]: |
| fused_layer = gluon.rnn.LSTM(hidden_size, projection_size=projection_size, |
| num_layers=num_layers, layout='TNC', bidirectional=False) |
| |
| stack_layer = mx.gluon.rnn.HybridSequentialRNNCell() |
| for _ in range(num_layers): |
| stack_layer.add(gluon.rnn.LSTMPCell(hidden_size, |
| projection_size=projection_size)) |
| fused_layer.initialize() |
| |
| fused_begin_state = fused_layer.begin_state(batch_size) |
| stack_begin_state = stack_layer.begin_state(batch_size=batch_size) |
| fused_layer.infer_shape(lstm_input, fused_begin_state) |
| stack_layer.infer_shape(0, lstm_input, False) |
| stack_layer.initialize() |
| fused_layer_params = fused_layer.collect_params() |
| stack_layer_params = stack_layer.collect_params() |
| |
| fused_weight_shape = fused_layer_params['rnn_param'].shape |
| print(fused_weight_shape) |
| w = mx.np.zeros(shape=fused_weight_shape) |
| fused_layer_params['rnn_param'].set_data(w) |
| fused_layer_params_split = split_rnn_params(w, 'lstm', num_layers, input_size,\ |
| hidden_size, False, projection_size=projection_size) |
| |
| for name, value in fused_layer_params_split.items(): |
| stack_layer_params[name[1:].replace('_', '.', 1)].set_data(value) |
| |
| fused_output, fused_states = fused_layer(lstm_input.copy(), fused_begin_state) |
| stack_output, stack_states = stack_layer.unroll(seq_len, lstm_input.copy(), begin_state=stack_begin_state, |
| layout='TNC', |
| merge_outputs=True) |
| |
| assert_almost_equal(fused_output.asnumpy(), stack_output.asnumpy(), rtol=rtol, atol=atol) |
| check_rnn_states(fused_states, stack_states, num_layers, False) |
| |
| # ==== Bidirectional Layer ==== |
| for num_layers in [1, 3]: |
| fused_layer = gluon.rnn.LSTM(hidden_size, projection_size=projection_size, |
| num_layers=num_layers, layout='TNC', bidirectional=True) |
| |
| stack_layer = mx.gluon.rnn.HybridSequentialRNNCell() |
| for _ in range(num_layers): |
| stack_layer.add( |
| gluon.rnn.BidirectionalCell(gluon.rnn.LSTMPCell(hidden_size, |
| projection_size=projection_size), |
| gluon.rnn.LSTMPCell(hidden_size, |
| projection_size=projection_size))) |
| fused_layer.initialize() |
| stack_layer.initialize() |
| |
| fused_begin_state = fused_layer.begin_state(batch_size) |
| stack_begin_state = stack_layer.begin_state(batch_size=batch_size) |
| fused_layer.infer_shape(lstm_input, fused_begin_state) |
| stack_layer.infer_shape(0, lstm_input, False) |
| stack_layer.initialize() |
| fused_layer_params = fused_layer.collect_params() |
| stack_layer_params = stack_layer.collect_params() |
| |
| fused_weight_shape = fused_layer_params['rnn_param'].shape |
| w = mx.np.zeros(shape=fused_weight_shape) |
| fused_layer_params['rnn_param'].set_data(w) |
| fused_layer_params_split = split_rnn_params(w, 'lstm', num_layers, input_size,\ |
| hidden_size, True, projection_size=projection_size) |
| |
| for name, value in fused_layer_params_split.items(): |
| cur = name.split("_")[0] |
| stack_layer_params["{}.{}_cell.{}".format(cur[1:], name[0], name[len(cur)+1:])].set_data(value) |
| |
| fused_output, fused_states = fused_layer(lstm_input.copy(), fused_begin_state) |
| stack_output, stack_states = stack_layer.unroll(seq_len, lstm_input.copy(), begin_state=stack_begin_state, |
| layout='TNC', |
| merge_outputs=True) |
| |
| assert_almost_equal(fused_output.asnumpy(), stack_output.asnumpy(), rtol=rtol, atol=atol) |
| check_rnn_states(fused_states, stack_states, num_layers, True) |
| |
| |
| @mx.util.use_np |
| @assert_raises_cudnn_not_satisfied(min_version='5.1.10') |
| def test_lstm_cpu_inference(): |
| # should behave the same as lstm cell |
| EXPECTED_LSTM_OUTPUT = np.array([[[0.72045636, 0.72045636, 0.95215213, 0.95215213], |
| [0.72045636, 0.72045636, 0.95215213, 0.95215213]], |
| [[0.95215213, 0.95215213, 0.72045636, 0.72045636], |
| [0.95215213, 0.95215213, 0.72045636, 0.72045636]]]) |
| x = mx.np.ones(shape=(2, 2, 2)) |
| model = mx.gluon.rnn.LSTM(2, num_layers=6, bidirectional=True) |
| model.initialize(mx.init.One()) |
| |
| y = model(x).asnumpy() |
| mx.test_utils.assert_almost_equal(y, EXPECTED_LSTM_OUTPUT, |
| rtol=1e-3, atol=1e-5) |
| |
| |
| @mx.util.use_np |
| def test_gru(): |
| cell = gluon.rnn.GRUCell(100, activation='relu', recurrent_activation='tanh') |
| inputs = [mx.np.ones((10,50)) for i in range(3)] |
| cell.infer_shape(0, inputs[0], False) |
| cell.initialize() |
| outputs, _ = cell.unroll(3, inputs) |
| assert sorted(cell.collect_params().keys()) == ['h2h_bias', 'h2h_weight', 'i2h_bias', 'i2h_weight'] |
| |
| outs = [o.shape for o in outputs] |
| assert outs == [(10, 100), (10, 100), (10, 100)] |
| |
| |
| @mx.util.use_np |
| @pytest.mark.serial |
| def test_residual(): |
| cell = gluon.rnn.ResidualCell(gluon.rnn.GRUCell(50)) |
| inputs = [mx.np.ones((10,50)) for i in range(2)] |
| cell.infer_shape(0, inputs[0], False) |
| cell.initialize() |
| outputs, _ = cell.unroll(2, inputs) |
| params = cell.collect_params() |
| assert sorted(params.keys()) == \ |
| ['base_cell.h2h_bias', 'base_cell.h2h_weight', 'base_cell.i2h_bias', 'base_cell.i2h_weight'] |
| |
| outs = [o.shape for o in outputs] |
| assert outs == [(10, 50), (10, 50)] |
| |
| |
| @mx.util.use_np |
| @pytest.mark.serial |
| def test_residual_bidirectional(): |
| cell = gluon.rnn.ResidualCell( |
| gluon.rnn.BidirectionalCell( |
| gluon.rnn.GRUCell(25), |
| gluon.rnn.GRUCell(25))) |
| inputs = [mx.np.ones((10,50)) for i in range(2)] |
| cell.infer_shape(0, inputs[0], True) |
| cell.initialize() |
| outputs, _ = cell.unroll(2, inputs, merge_outputs=False) |
| params = cell.collect_params() |
| assert sorted(params.keys()) == \ |
| ['base_cell.l_cell.h2h_bias', 'base_cell.l_cell.h2h_weight', |
| 'base_cell.l_cell.i2h_bias', 'base_cell.l_cell.i2h_weight', |
| 'base_cell.r_cell.h2h_bias', 'base_cell.r_cell.h2h_weight', |
| 'base_cell.r_cell.i2h_bias', 'base_cell.r_cell.i2h_weight'] |
| |
| outs = [o.shape for o in outputs] |
| assert outs == [(10, 50), (10, 50)] |
| |
| |
| @mx.util.use_np |
| def test_stack(): |
| cell = gluon.rnn.SequentialRNNCell() |
| for i in range(5): |
| if i == 1: |
| cell.add(gluon.rnn.ResidualCell(gluon.rnn.LSTMCell(100))) |
| else: |
| cell.add(gluon.rnn.LSTMCell(100)) |
| inputs = [mx.np.ones((10,50)) for i in range(3)] |
| cell.infer_shape(0, inputs[0], False) |
| cell.initialize() |
| outputs, _ = cell.unroll(3, inputs) |
| keys = sorted(cell.collect_params().keys()) |
| for i in range(5): |
| if i==1: |
| continue |
| assert f'{i}.h2h_weight' in keys |
| assert f'{i}.h2h_bias' in keys |
| assert f'{i}.i2h_weight' in keys |
| assert f'{i}.i2h_bias' in keys |
| assert '1.base_cell.h2h_weight' in keys |
| assert '1.base_cell.h2h_bias' in keys |
| assert '1.base_cell.i2h_weight' in keys |
| assert '1.base_cell.i2h_bias' in keys |
| |
| outs = [o.shape for o in outputs] |
| assert outs == [(10, 100), (10, 100), (10, 100)] |
| |
| |
| @mx.util.use_np |
| @pytest.mark.serial |
| def test_hybridstack(): |
| cell = gluon.rnn.HybridSequentialRNNCell() |
| for i in range(5): |
| if i == 1: |
| cell.add(gluon.rnn.ResidualCell(gluon.rnn.LSTMCell(100))) |
| else: |
| cell.add(gluon.rnn.LSTMCell(100)) |
| inputs = [mx.np.ones((10,50)) for i in range(3)] |
| cell.infer_shape(0, inputs[0], False) |
| cell.initialize() |
| outputs, _ = cell.unroll(3, inputs) |
| keys = sorted(cell.collect_params().keys()) |
| for i in range(5): |
| if i==1: |
| continue |
| assert f'{i}.h2h_weight' in keys |
| assert f'{i}.h2h_bias' in keys |
| assert f'{i}.i2h_weight' in keys |
| assert f'{i}.i2h_bias' in keys |
| assert '1.base_cell.h2h_weight' in keys |
| assert '1.base_cell.h2h_bias' in keys |
| assert '1.base_cell.i2h_weight' in keys |
| assert '1.base_cell.i2h_bias' in keys |
| |
| outs = [o.shape for o in outputs] |
| assert outs == [(10, 100), (10, 100), (10, 100)] |
| |
| # Test HybridSequentialRNNCell nested in nn.HybridBlock, SequentialRNNCell will fail in this case |
| class BidirectionalOfSequential(gluon.HybridBlock): |
| def __init__(self): |
| super(BidirectionalOfSequential, self).__init__() |
| |
| cell0 = gluon.rnn.HybridSequentialRNNCell() |
| cell0.add(gluon.rnn.LSTMCell(100)) |
| cell0.add(gluon.rnn.LSTMCell(100)) |
| |
| cell1 = gluon.rnn.HybridSequentialRNNCell() |
| cell1.add(gluon.rnn.LSTMCell(100)) |
| cell1.add(gluon.rnn.LSTMCell(100)) |
| |
| self.rnncell = gluon.rnn.BidirectionalCell(cell0, cell1) |
| |
| def forward(self, x): |
| return self.rnncell.unroll(3, x, layout="NTC", merge_outputs=True) |
| |
| def infer_shape(self, x, *args): |
| self.rnncell.infer_shape(0, x, True) |
| |
| x = mx.np.random.uniform(size=(10, 3, 100)) |
| net = BidirectionalOfSequential() |
| net.infer_shape(x) |
| net.initialize() |
| outs, _ = net(x) |
| |
| assert outs.shape == (10, 3, 200) |
| |
| |
| @mx.util.use_np |
| def test_bidirectional(): |
| cell = gluon.rnn.BidirectionalCell( |
| gluon.rnn.LSTMCell(100), |
| gluon.rnn.LSTMCell(100)) |
| inputs = [mx.np.ones((10,50)) for i in range(3)] |
| cell.infer_shape(0, inputs[0], False) |
| cell.initialize() |
| outputs, _ = cell.unroll(3, inputs) |
| |
| outs = [o.shape for o in outputs] |
| assert outs == [(10, 200), (10, 200), (10, 200)] |
| |
| |
| @mx.util.use_np |
| @assert_raises_cudnn_not_satisfied(min_version='5.1.10') |
| @pytest.mark.serial |
| def test_layer_bidirectional(): |
| class RefBiLSTM(gluon.Block): |
| def __init__(self, size, **kwargs): |
| super(RefBiLSTM, self).__init__(**kwargs) |
| self._lstm_fwd = gluon.rnn.LSTM(size, bidirectional=False) |
| self._lstm_bwd = gluon.rnn.LSTM(size, bidirectional=False) |
| |
| def forward(self, inpt): |
| fwd = self._lstm_fwd(inpt) |
| bwd_inpt = np.flip(inpt, 0) |
| bwd = self._lstm_bwd(bwd_inpt) |
| bwd = np.flip(bwd, 0) |
| return np.concatenate([fwd, bwd], axis=2) |
| |
| size = 7 |
| in_size = 5 |
| weights = {} |
| for d in ['l', 'r']: |
| weights['{}0_i2h_weight'.format(d)] = mx.np.random.uniform(size=(size*4, in_size)) |
| weights['{}0_h2h_weight'.format(d)] = mx.np.random.uniform(size=(size*4, size)) |
| weights['{}0_i2h_bias'.format(d)] = mx.np.random.uniform(size=(size*4,)) |
| weights['{}0_h2h_bias'.format(d)] = mx.np.random.uniform(size=(size*4,)) |
| |
| params = (weights['{}0_{}_{}'.format(d, g, t)].reshape(-1) |
| for t in ['weight', 'bias'] |
| for d in ['l', 'r'] |
| for g in ['i2h', 'h2h']) |
| net_params_concat = mx.np.concatenate(params) |
| params_left = (weights['l0_{}_{}'.format(g, t)].reshape(-1) |
| for t in ['weight', 'bias'] |
| for g in ['i2h', 'h2h']) |
| params_right = (weights['r0_{}_{}'.format(g, t)].reshape(-1) |
| for t in ['weight', 'bias'] |
| for g in ['i2h', 'h2h']) |
| net_ref_left_params = mx.np.concatenate(params_left) |
| net_ref_right_params = mx.np.concatenate(params_right) |
| net = gluon.rnn.LSTM(size, bidirectional=True) |
| ref_net = RefBiLSTM(size) |
| net.initialize() |
| ref_net.initialize() |
| net_params = net.collect_params() |
| ref_net_params = ref_net.collect_params() |
| net_params['rnn_param'].set_data(net_params_concat) |
| ref_net_params['_lstm_fwd.rnn_param'].set_data(net_ref_left_params) |
| ref_net_params['_lstm_bwd.rnn_param'].set_data(net_ref_right_params) |
| |
| data = mx.np.random.uniform(size=(11, 10, in_size)) |
| assert_allclose(net(data).asnumpy(), ref_net(data).asnumpy(), rtol=1e-04, atol=1e-02) |
| |
| |
| def test_zoneout(): |
| cell = gluon.rnn.ZoneoutCell(gluon.rnn.RNNCell(100), zoneout_outputs=0.5, |
| zoneout_states=0.5) |
| inputs = [mx.np.ones((10,50)) for i in range(3)] |
| cell.infer_shape(0, inputs[0], False) |
| cell.initialize() |
| outputs, _ = cell.unroll(3, inputs) |
| |
| outs = [o.shape for o in outputs] |
| assert outs == [(10, 100), (10, 100), (10, 100)] |
| |
| |
| @pytest.mark.serial |
| def test_unroll_layout(): |
| cell = gluon.rnn.HybridSequentialRNNCell() |
| for i in range(5): |
| if i == 1: |
| cell.add(gluon.rnn.ResidualCell(gluon.rnn.LSTMCell(100))) |
| else: |
| cell.add(gluon.rnn.LSTMCell(100)) |
| inputs = [mx.np.random.uniform(size=(10,50)) for _ in range(3)] |
| cell.infer_shape(0, inputs[0], False) |
| cell.initialize() |
| outputs, _ = cell.unroll(3, inputs, layout='TNC') |
| assert outputs[0].shape == (10, 100) |
| assert outputs[1].shape == (10, 100) |
| assert outputs[2].shape == (10, 100) |
| |
| outputs, _ = cell.unroll(3, inputs, layout='NTC') |
| assert outputs[0].shape == (10, 100) |
| assert outputs[1].shape == (10, 100) |
| assert outputs[2].shape == (10, 100) |
| |
| |
| def check_rnn_forward_backward(layer, merged_inputs, hybridize, merge_outputs, deterministic): |
| input_size = 5 |
| if merged_inputs: |
| inputs = mx.np.ones((8, 3, 5)) |
| inputs.attach_grad() |
| else: |
| inputs = [mx.np.ones((8, 5)) for _ in range(3)] |
| for x in inputs: |
| x.attach_grad() |
| |
| if hybridize: |
| layer.hybridize() |
| layer.initialize() |
| |
| with mx.autograd.record(): |
| out = layer.unroll(3, inputs, merge_outputs=merge_outputs)[0] |
| mx.autograd.backward(out) |
| |
| if hasattr(layer, 'i2h_weight'): |
| assert layer.i2h_weight.shape[1] == input_size, (layer.i2h_weight.shape[1], input_size) |
| |
| if merge_outputs: |
| np_out = out.asnumpy() |
| else: |
| np_out = _np.stack([x.asnumpy() for x in out], axis=1) |
| |
| if merged_inputs: |
| np_dx = inputs.grad.asnumpy() |
| else: |
| np_dx = _np.stack([x.grad.asnumpy() for x in inputs], axis=1) |
| |
| with mx.autograd.record(): |
| out = layer.unroll(3, inputs, merge_outputs=not merge_outputs)[0] |
| mx.autograd.backward(out) |
| |
| if merged_inputs: |
| input_grads = inputs.grad.asnumpy() |
| else: |
| input_grads = _np.stack([x.grad.asnumpy() for x in inputs], axis=1) |
| |
| if deterministic: |
| if not merge_outputs: |
| ref_np_out = out.asnumpy() |
| else: |
| ref_np_out = _np.stack([x.asnumpy() for x in out], axis=1) |
| mx.test_utils.assert_almost_equal(np_out, ref_np_out, rtol=1e-3, atol=1e-5) |
| mx.test_utils.assert_almost_equal(np_dx, input_grads, rtol=1e-3, atol=1e-5) |
| |
| |
| @retry(3) |
| @pytest.mark.parametrize('layer,determinism', [ |
| (gluon.rnn.LSTMCell(10, input_size=5), True), |
| (gluon.rnn.RNNCell(10, input_size=5), True), |
| (gluon.rnn.GRUCell(10, input_size=5), True), |
| (gluon.rnn.BidirectionalCell( |
| gluon.rnn.LSTMCell(10, input_size=5), |
| gluon.rnn.LSTMCell(10, input_size=5) |
| ), True), |
| (gluon.rnn.DropoutCell(0.5), False), |
| ]) |
| @pytest.mark.parametrize('merged_inputs', [True, False]) |
| @pytest.mark.parametrize('hybridize', [True, False]) |
| @pytest.mark.parametrize('merge_outputs', [True, False, None]) |
| @pytest.mark.skip(reason='https://github.com/apache/incubator-mxnet/issues/18225') |
| def test_rnn_forward_backward(layer, merged_inputs, hybridize, merge_outputs, determinism): |
| check_rnn_forward_backward(layer, merged_inputs, hybridize, merge_outputs, determinism) |
| |
| |
| @pytest.mark.parametrize('seq_rnn_type', [ |
| gluon.rnn.SequentialRNNCell, |
| gluon.rnn.HybridSequentialRNNCell |
| ]) |
| @pytest.mark.parametrize('determinism', [True, False]) |
| @pytest.mark.parametrize('merged_inputs', [True, False]) |
| @pytest.mark.parametrize('hybridize', [True, False]) |
| @pytest.mark.parametrize('merge_outputs', [True, False, None]) |
| @pytest.mark.skip(reason='https://github.com/apache/incubator-mxnet/issues/18291') |
| def test_sequential_rnn_cells(seq_rnn_type, determinism, merged_inputs, hybridize, merge_outputs): |
| net = gluon.rnn.SequentialRNNCell() |
| net.add(gluon.rnn.LSTMCell(10, input_size=5)) |
| net.add(gluon.rnn.RNNCell(10, input_size=10)) |
| net.add(gluon.rnn.GRUCell(10, input_size=10)) |
| if not determinism: |
| net.add(gluon.rnn.DropoutCell(0.5)) |
| check_rnn_forward_backward(net, merged_inputs, hybridize, merge_outputs, determinism) |
| |
| |
| @mx.util.use_np |
| def test_rnn_cells_export_import(): |
| class RNNLayer(gluon.HybridBlock): |
| def __init__(self): |
| super(RNNLayer, self).__init__() |
| self.cell = gluon.rnn.RNNCell(hidden_size=1) |
| |
| def forward(self, seq): |
| outputs, state = self.cell.unroll(inputs=seq, length=2, merge_outputs=True) |
| return outputs |
| |
| def infer_shape(self, input): |
| self.cell.infer_shape(0, input, False) |
| |
| class LSTMLayer(gluon.HybridBlock): |
| def __init__(self): |
| super(LSTMLayer, self).__init__() |
| self.cell = gluon.rnn.LSTMCell(hidden_size=1) |
| |
| def forward(self, seq): |
| outputs, state = self.cell.unroll(inputs=seq, length=2, merge_outputs=True) |
| return outputs |
| |
| def infer_shape(self, input): |
| self.cell.infer_shape(0, input, False) |
| |
| class GRULayer(gluon.HybridBlock): |
| def __init__(self): |
| super(GRULayer, self).__init__() |
| self.cell = gluon.rnn.GRUCell(hidden_size=1) |
| |
| def forward(self, seq): |
| outputs, state = self.cell.unroll(inputs=seq, length=2, merge_outputs=True) |
| return outputs |
| |
| def infer_shape(self, input): |
| self.cell.infer_shape(0, input, False) |
| |
| for hybrid in [RNNLayer(), LSTMLayer(), GRULayer()]: |
| input = mx.np.ones(shape=(1, 2, 1), device=mx.device.current_device()) |
| hybrid.infer_shape(input) |
| hybrid.initialize() |
| hybrid.hybridize() |
| output1 = hybrid(input) |
| hybrid.export(path="./model", epoch=0) |
| symbol = mx.gluon.SymbolBlock.imports( |
| symbol_file="./model-symbol.json", |
| input_names=["data"], |
| param_file="./model-0000.params", |
| device=mx.device.current_device() |
| ) |
| output2 = symbol(input) |
| assert_almost_equal(output1.asnumpy(), output2.asnumpy()) |
| |
| |
| def check_rnn_layer_forward(layer, inputs, states=None, run_only=False, device=mx.cpu()): |
| layer.initialize(device=device) |
| inputs = inputs.to_device(device) |
| inputs.attach_grad() |
| if states is not None: |
| if isinstance(states, (list, tuple)): |
| states = [s.to_device(device) for s in states] |
| else: |
| states = states.to_device(device) |
| with mx.autograd.record(): |
| if states is None: |
| out = layer(inputs) |
| else: |
| out = layer(inputs, states) |
| if states is not None: |
| assert isinstance(out, (list, tuple)) and len(out) == 2 |
| out = out[0] |
| else: |
| assert isinstance(out, mx.np.ndarray) |
| out.backward() |
| |
| np_out = out.asnumpy() |
| np_dx = inputs.grad.asnumpy() |
| |
| layer.hybridize() |
| |
| with mx.autograd.record(): |
| if states is not None: |
| out = layer(inputs, states) |
| assert isinstance(out, (list, tuple)) and len(out) == 2 |
| out = out[0] |
| else: |
| out = layer(inputs) |
| assert isinstance(out, mx.np.ndarray) |
| out.backward() |
| |
| if states is not None: |
| layer(inputs, states) # test is_training = false |
| else: |
| layer(inputs) |
| |
| if not run_only: |
| mx.test_utils.assert_almost_equal(np_out, out.asnumpy(), rtol=1e-3, atol=1e-5) |
| mx.test_utils.assert_almost_equal(np_dx, inputs.grad.asnumpy(), rtol=1e-3, atol=1e-5) |
| |
| |
| @mx.util.use_np |
| def run_rnn_layers(dtype, dtype2): |
| device = default_device() |
| |
| check_rnn_layer_forward(gluon.rnn.RNN(10, 2, dtype=dtype), mx.np.ones((8, 3, 20), dtype=dtype), device=device) |
| check_rnn_layer_forward(gluon.rnn.RNN(10, 2, dtype=dtype, bidirectional=True), mx.np.ones((8, 3, 20), dtype=dtype), mx.np.ones((4, 3, 10), dtype=dtype), device=device) |
| check_rnn_layer_forward(gluon.rnn.LSTM(10, 2,dtype=dtype), mx.np.ones((8, 3, 20), dtype=dtype), device=device) |
| check_rnn_layer_forward(gluon.rnn.LSTM(10, 2,dtype=dtype, bidirectional=True), mx.np.ones((8, 3, 20), dtype=dtype), [mx.np.ones((4, 3, 10), dtype=dtype), mx.np.ones((4, 3, 10), dtype=dtype)],device=device) |
| check_rnn_layer_forward(gluon.rnn.GRU(10, 2, dtype=dtype, ), mx.np.ones((8, 3, 20), dtype=dtype),device=device) |
| check_rnn_layer_forward(gluon.rnn.GRU(10, 2, dtype=dtype, bidirectional=True), mx.np.ones((8, 3, 20), dtype=dtype), mx.np.ones((4, 3, 10), dtype=dtype),device=device) |
| |
| |
| check_rnn_layer_forward(gluon.rnn.RNN(10, 2, dtype=dtype, dropout=0.5), mx.np.ones((8, 3, 20), dtype=dtype), |
| run_only=True, device=device) |
| check_rnn_layer_forward(gluon.rnn.RNN(10, 2, bidirectional=True, dropout=0.5, dtype=dtype), |
| mx.np.ones((8, 3, 20), dtype=dtype), mx.np.ones((4, 3, 10), dtype=dtype), run_only=True, device=device) |
| check_rnn_layer_forward(gluon.rnn.LSTM(10, 2, dropout=0.5, dtype=dtype), mx.np.ones((8, 3, 20), dtype=dtype), |
| run_only=True, device=device) |
| check_rnn_layer_forward(gluon.rnn.LSTM(10, 2, bidirectional=True, dropout=0.5, dtype=dtype), |
| mx.np.ones((8, 3, 20), dtype=dtype), |
| [mx.np.ones((4, 3, 10), dtype=dtype), mx.np.ones((4, 3, 10), dtype=dtype)], run_only=True, device=device) |
| check_rnn_layer_forward(gluon.rnn.GRU(10, 2, dropout=0.5, dtype=dtype), mx.np.ones((8, 3, 20), dtype=dtype), |
| run_only=True, device=device) |
| check_rnn_layer_forward(gluon.rnn.GRU(10, 2, bidirectional=True, dropout=0.5, dtype=dtype), |
| mx.np.ones((8, 3, 20), dtype=dtype), mx.np.ones((4, 3, 10), dtype=dtype), run_only=True, device=device) |
| |
| net = gluon.nn.Sequential() |
| net.add(gluon.rnn.LSTM(10, bidirectional=True, dtype=dtype2)) |
| net.add(gluon.nn.BatchNorm(axis=2)) |
| net.add(gluon.nn.Flatten()) |
| net.add(gluon.nn.Dense(3, activation='relu')) |
| net.initialize(device=device) |
| net.cast(dtype) |
| with mx.autograd.record(): |
| out = net(mx.np.ones((2, 3, 10), dtype=dtype, device=device)) |
| out.backward() |
| out = out.asnumpy() |
| |
| net2 = gluon.nn.HybridSequential() |
| net2.add(gluon.rnn.LSTM(10, bidirectional=True, dtype=dtype2)) |
| net2.add(gluon.nn.BatchNorm(axis=2)) |
| net2.add(gluon.nn.Flatten()) |
| net2.add(gluon.nn.Dense(3, activation='relu')) |
| net2.hybridize() |
| net2.initialize(device=device) |
| net2.cast(dtype) |
| with mx.autograd.record(): |
| out = net2(mx.np.ones((2, 3, 10), dtype=dtype, device=device)) |
| out.backward() |
| out = out.asnumpy() |
| |
| net3 = gluon.nn.HybridSequential() |
| net3.add(gluon.rnn.LSTM(10, bidirectional=True, dtype=dtype)) |
| net3.add(gluon.nn.BatchNorm(axis=2)) |
| net3.add(gluon.nn.Flatten()) |
| net3.add(gluon.nn.Dense(3, activation='relu')) |
| net3.hybridize() |
| net3.initialize(device=device) |
| net3.cast(dtype2) |
| with mx.autograd.record(): |
| out = net3(mx.np.ones((2, 3, 10), dtype=dtype2, device=device)) |
| out.backward() |
| out = out.asnumpy() |
| |
| @assert_raises_cudnn_not_satisfied(min_version='5.1.10') |
| @pytest.mark.serial |
| def test_rnn_layers_fp32(): |
| run_rnn_layers('float32', 'float32') |
| |
| @assert_raises_cudnn_not_satisfied(min_version='5.1.10') |
| @pytest.mark.serial |
| def test_rnn_layers_fp16(): |
| # Dynamic skip condition is best handled this way, rather than with pytest.mark.skipIf |
| if default_device().device_type == 'cpu': |
| pytest.skip('RNN FP16 only implemented for GPU for now') |
| run_rnn_layers('float16', 'float32') |
| |
| |
| def check_rnn_consistency(fused_layer, stack_layer, loss, mode, num_layers, input_size, hidden_size, bidirectional=False, rtol=1e-2, atol=1e-4): |
| x = mx.np.random.normal(size=(1, 5, input_size)) |
| fused_begin_state = fused_layer.begin_state(1) |
| stack_states = stack_layer.begin_state(batch_size=1) |
| fused_layer.infer_shape(x, fused_begin_state) |
| fused_layer_params = fused_layer.collect_params() |
| stack_layer.infer_shape(0, x, False) |
| stack_layer.initialize() |
| stack_layer_params = stack_layer.collect_params() |
| |
| fused_weight_shape = fused_layer_params['rnn_param'].shape |
| w = mx.np.zeros(shape=fused_weight_shape) |
| fused_layer_params_split = split_rnn_params(w, mode, num_layers, input_size, hidden_size, bidirectional) |
| for name, value in fused_layer_params_split.items(): |
| if 'bias' in name: |
| fused_layer_params_split[name] = mx.np.random.normal(size=value.shape) |
| _dir = 2 if bidirectional else 1 |
| params = (fused_layer_params_split['{}{}_{}_{}'.format(d, l, g, t)].reshape(-1) |
| for t in ['weight', 'bias'] |
| for l in range(num_layers) |
| for d in ['l', 'r'][:_dir] |
| for g in ['i2h', 'h2h']) |
| fused_params = mx.np.concatenate(params) |
| fused_layer_params['rnn_param'].set_data(fused_params) |
| for name, value in fused_layer_params_split.items(): |
| cur = name.split('_')[0] |
| num = cur[1:] |
| stack_name = ('{}.{}_cell.'.format(num, name[0]) if bidirectional else num + '.' ) + name[len(cur)+1:] |
| stack_layer_params[stack_name].set_data(value) |
| |
| fx = x.copy() |
| sx = x.copy() |
| y = mx.np.random.uniform(size=(1, 5, hidden_size * 2 if bidirectional else hidden_size)) |
| fx.attach_grad() |
| with mx.autograd.record(): |
| fused_out, fused_states = fused_layer(fx, fused_begin_state) |
| l = loss(fused_out, y).mean() |
| l.backward() |
| mx.npx.waitall() |
| fused_grads = split_rnn_params(fused_layer.collect_params()['rnn_param'].data().grad,\ |
| mode, num_layers, input_size, hidden_size, bidirectional) |
| fused_input_grad = fx.grad.asnumpy() |
| |
| sx.attach_grad() |
| with mx.autograd.record(): |
| stack_out, stack_states = stack_layer.unroll(5, sx, begin_state=stack_states, merge_outputs=True) |
| l = loss(stack_out, y).mean() |
| l.backward() |
| mx.npx.waitall() |
| stack_grads = dict([(name, p.grad()) for name, p in stack_layer.collect_params().items()]) |
| stack_input_grad = sx.grad.asnumpy() |
| |
| assert_allclose(fused_out.asnumpy(), stack_out.asnumpy(), rtol=rtol, atol=atol) |
| assert_allclose(fused_input_grad, stack_input_grad, rtol=rtol, atol=atol) |
| for name, value in fused_grads.items(): |
| cur = name.split('_')[0] |
| num = cur[1:] |
| stack_name = ('{}.{}_cell.'.format(num, name[0]) if bidirectional else num + '.' ) + name[len(cur)+1:] |
| assert_allclose(value.asnumpy(), stack_grads[stack_name].asnumpy(), rtol=rtol, atol=atol) |
| |
| num_layers = fused_begin_state[0].shape[0] // (2 if bidirectional else 1) |
| check_rnn_states(fused_states, stack_states, num_layers, bidirectional, len(fused_begin_state) == 2) |
| |
| |
| def create_op_by_mode(mode): |
| if mode == 'lstm': |
| fused_op = gluon.rnn.LSTM |
| stack_op = gluon.rnn.LSTMCell |
| recurrent_block_prefix = 'lstm0_' |
| elif mode == 'gru': |
| fused_op = gluon.rnn.GRU |
| stack_op = gluon.rnn.GRUCell |
| recurrent_block_prefix = 'gru0_' |
| elif mode == 'rnn_relu': |
| fused_op = partial(gluon.rnn.RNN, activation='relu') |
| stack_op = partial(gluon.rnn.RNNCell, activation='relu') |
| recurrent_block_prefix = 'rnn0_' |
| elif mode == 'rnn_tanh': |
| fused_op = partial(gluon.rnn.RNN, activation='tanh') |
| stack_op = partial(gluon.rnn.RNNCell, activation='tanh') |
| recurrent_block_prefix = 'rnn0_' |
| |
| return fused_op, stack_op, recurrent_block_prefix |
| |
| |
| def check_rnn_unidir_layer_gradients(mode, input_size, hidden_size, num_layers, loss): |
| fused_op, stack_op, recurrent_block_prefix = create_op_by_mode(mode) |
| |
| fused_layer = fused_op(hidden_size, num_layers=num_layers, layout='NTC', bidirectional=False) |
| fused_layer.initialize() |
| |
| stack_layer = mx.gluon.rnn.HybridSequentialRNNCell() |
| for _ in range(num_layers): |
| stack_layer.add(stack_op(hidden_size)) |
| stack_layer.initialize() |
| check_rnn_consistency(fused_layer, stack_layer, loss, mode, num_layers, input_size, hidden_size) |
| |
| |
| def check_rnn_bidir_layer_gradients(mode, input_size, hidden_size, num_layers, loss): |
| fused_op, stack_op, recurrent_block_prefix = create_op_by_mode(mode) |
| |
| fused_layer = fused_op(hidden_size, num_layers=num_layers, layout='NTC', bidirectional=True) |
| fused_layer.initialize() |
| |
| stack_layer = mx.gluon.rnn.HybridSequentialRNNCell() |
| for _ in range(num_layers): |
| stack_layer.add(gluon.rnn.BidirectionalCell(stack_op(hidden_size), |
| stack_op(hidden_size))) |
| stack_layer.initialize() |
| check_rnn_consistency(fused_layer, stack_layer, loss, mode, num_layers, input_size, hidden_size, bidirectional=True) |
| |
| |
| @mx.util.use_np |
| @pytest.mark.parametrize('input_size', [8]) |
| @pytest.mark.parametrize('hidden_size', [8, 16]) |
| @pytest.mark.parametrize('num_layers', [1, 2, 3, 4]) |
| @pytest.mark.parametrize('func', [check_rnn_unidir_layer_gradients, |
| check_rnn_bidir_layer_gradients]) |
| @pytest.mark.parametrize('mode', ['lstm', 'gru', 'rnn_relu', 'rnn_tanh']) |
| @assert_raises_cudnn_not_satisfied(min_version='5.1.10') |
| def test_fused_layer(input_size, hidden_size, num_layers, func, mode): |
| loss = mx.gluon.loss.L2Loss() |
| func(mode, input_size, hidden_size, num_layers, loss) |
| |
| |
| @pytest.mark.serial |
| def test_rnn_unroll_variant_length(): |
| # Test for imperative usage |
| cell_list = [] |
| for base_cell_class in [gluon.rnn.RNNCell, gluon.rnn.LSTMCell, gluon.rnn.GRUCell]: |
| cell_list.append(base_cell_class(20)) |
| cell_list.append(gluon.rnn.BidirectionalCell( |
| l_cell=base_cell_class(20), |
| r_cell=base_cell_class(20))) |
| cell_list.append(gluon.rnn.VariationalDropoutCell(base_cell=base_cell_class(20))) |
| stack_res_rnn_cell = gluon.rnn.SequentialRNNCell() |
| stack_res_rnn_cell.add(gluon.rnn.ResidualCell(base_cell=gluon.rnn.RNNCell(20))) |
| stack_res_rnn_cell.add(gluon.rnn.ResidualCell(base_cell=gluon.rnn.RNNCell(20))) |
| cell_list.append(stack_res_rnn_cell) |
| batch_size = 4 |
| max_length = 10 |
| valid_length = [3, 10, 5, 6] |
| valid_length_nd = mx.np.array(valid_length) |
| for cell in cell_list: |
| # Test for NTC layout |
| print(type(cell)) |
| data_nd = mx.np.random.normal(0, 1, size=(batch_size, max_length, 20)) |
| if isinstance(cell, (gluon.rnn.HybridSequentialRNNCell, gluon.rnn.SequentialRNNCell)): |
| cell.infer_shape(0, data_nd, False) |
| elif isinstance(cell, gluon.rnn.BidirectionalCell): |
| cell.infer_shape(0, data_nd, True) |
| else: |
| cell.infer_shape(0, data_nd, False) |
| cell.initialize() |
| cell.hybridize() |
| outs, states = cell.unroll(length=max_length, inputs=data_nd, |
| valid_length=valid_length_nd, |
| merge_outputs=True, |
| layout='NTC') |
| for i, ele_length in enumerate(valid_length): |
| # Explicitly unroll each sequence and compare the final states and output |
| ele_out, ele_states = cell.unroll(length=ele_length, |
| inputs=data_nd[i:(i+1), :ele_length, :], |
| merge_outputs=True, |
| layout='NTC') |
| assert_almost_equal(ele_out, outs[i:(i+1), :ele_length, :]) |
| if ele_length < max_length: |
| # Check the padded outputs are all zero |
| assert_almost_equal(outs[i:(i+1), ele_length:max_length, :], 0) |
| for valid_out_state, gt_state in zip(states, ele_states): |
| assert_almost_equal(valid_out_state[i:(i+1)], gt_state) |
| |
| # Test for TNC layout |
| data_nd = mx.np.random.normal(0, 1, size=(max_length, batch_size, 20)) |
| outs, states = cell.unroll(length=max_length, inputs=data_nd, |
| valid_length=valid_length_nd, |
| layout='TNC') |
| for i, ele_length in enumerate(valid_length): |
| # Explicitly unroll each sequence and compare the final states and output |
| ele_out, ele_states = cell.unroll(length=ele_length, |
| inputs=data_nd[:ele_length, i:(i+1), :], |
| merge_outputs=True, |
| layout='TNC') |
| assert_almost_equal(ele_out, outs[:ele_length, i:(i + 1), :]) |
| if ele_length < max_length: |
| # Check the padded outputs are all zero |
| assert_almost_equal(outs[ele_length:max_length, i:(i+1), :], 0) |
| for valid_out_state, gt_state in zip(states, ele_states): |
| assert_almost_equal(valid_out_state[i:(i+1)], gt_state) |
| |
| |
| def test_cell_fill_shape(): |
| cell = gluon.rnn.LSTMCell(10, input_size=7) |
| cell.hybridize() |
| assert cell.i2h_weight.shape[1] == 7, cell.i2h_weight.shape[1] |
| |
| def test_layer_fill_shape(): |
| layer = gluon.rnn.LSTM(10) |
| layer.hybridize() |
| check_rnn_layer_forward(layer, mx.np.ones((3, 2, 7))) |
| print(layer) |
| assert layer.rnn_param.shape[0] == 760 |
| |
| |
| @pytest.mark.serial |
| def test_bidirectional_unroll_valid_length(): |
| def _check_bidirectional_unroll_valid_length(length): |
| class BiLSTM(gluon.nn.HybridBlock): |
| def __init__(self, rnn_size, time_step, **kwargs): |
| super(BiLSTM, self).__init__(**kwargs) |
| self.time_step = time_step |
| self.bi_lstm = gluon.rnn.BidirectionalCell( |
| gluon.rnn.LSTMCell(rnn_size), |
| gluon.rnn.LSTMCell(rnn_size)) |
| |
| def forward(self, inputs, valid_len): |
| outputs, states = self.bi_lstm.unroll(self.time_step, inputs, valid_length=valid_len, |
| layout='NTC', merge_outputs=True) |
| return outputs, states |
| |
| def infer_shape(self, x, *args): |
| self.bi_lstm.infer_shape(0, x, True) |
| |
| rnn_size = 100 |
| net = BiLSTM(rnn_size, length) |
| inputs_data = mx.np.random.uniform(size=(10, length, 50)) |
| net.infer_shape(inputs_data) |
| net.initialize() |
| net.hybridize() |
| valid_len = mx.np.array([length]*10) |
| outputs, _ = net(inputs_data, valid_len) |
| assert outputs.shape == (10, length, 200) |
| |
| _check_bidirectional_unroll_valid_length(1) |
| _check_bidirectional_unroll_valid_length(3) |
| |
| |
| def check_rnn_forward(layer, inputs): |
| inputs.attach_grad() |
| layer.initialize() |
| with mx.autograd.record(): |
| layer.unroll(3, inputs, merge_outputs=True)[0].backward() |
| mx.autograd.backward(layer.unroll(3, inputs, merge_outputs=False)[0]) |
| mx.npx.waitall() |
| |
| |
| def test_rnn_cells(): |
| check_rnn_forward(gluon.rnn.Conv1DLSTMCell((5, 7), 10, (3,), (3,)), |
| mx.np.ones((8, 3, 5, 7))) |
| check_rnn_forward(gluon.rnn.Conv1DRNNCell((5, 7), 10, (3,), (3,)), |
| mx.np.ones((8, 3, 5, 7))) |
| check_rnn_forward(gluon.rnn.Conv1DGRUCell((5, 7), 10, (3,), (3,)), |
| mx.np.ones((8, 3, 5, 7))) |
| |
| net = mx.gluon.rnn.SequentialRNNCell() |
| net.add(gluon.rnn.Conv1DLSTMCell((5, 7), 10, (3,), (3,))) |
| net.add(gluon.rnn.Conv1DRNNCell((10, 5), 11, (3,), (3,))) |
| net.add(gluon.rnn.Conv1DGRUCell((11, 3), 12, (3,), (3,))) |
| check_rnn_forward(net, mx.np.ones((8, 3, 5, 7))) |
| |
| |
| @mx.util.use_np |
| def check_rnn_cell(cell, in_shape=(10, 50), out_shape=(10, 100), begin_state=None): |
| inputs = [mx.np.ones(shape=in_shape) for i in range(3)] |
| cell.infer_shape(0, inputs[0], False) |
| cell.initialize() |
| outputs, _ = cell.unroll(3, inputs, begin_state=begin_state) |
| assert sorted(cell.collect_params().keys()) == ['h2h_bias', 'h2h_weight', |
| 'i2h_bias', 'i2h_weight'] |
| |
| outs = [o.shape for o in outputs] |
| assert outs == [out_shape] * 3 |
| |
| |
| @mx.util.use_np |
| def test_convrnn(): |
| cell = gluon.rnn.Conv1DRNNCell((10, 50), 100, 3, 3) |
| check_rnn_cell(cell, in_shape=(1, 10, 50), out_shape=(1, 100, 48)) |
| |
| cell = gluon.rnn.Conv2DRNNCell((10, 20, 50), 100, 3, 3) |
| check_rnn_cell(cell, in_shape=(1, 10, 20, 50), out_shape=(1, 100, 18, 48)) |
| |
| cell = gluon.rnn.Conv3DRNNCell((10, 20, 30, 50), 100, 3, 3) |
| check_rnn_cell(cell, in_shape=(1, 10, 20, 30, 50), out_shape=(1, 100, 18, 28, 48)) |
| |
| |
| @mx.util.use_np |
| def test_convlstm(): |
| cell = gluon.rnn.Conv1DLSTMCell((10, 50), 100, 3, 3) |
| check_rnn_cell(cell, in_shape=(1, 10, 50), out_shape=(1, 100, 48)) |
| |
| cell = gluon.rnn.Conv2DLSTMCell((10, 20, 50), 100, 3, 3) |
| check_rnn_cell(cell, in_shape=(1, 10, 20, 50), out_shape=(1, 100, 18, 48)) |
| |
| cell = gluon.rnn.Conv3DLSTMCell((10, 20, 30, 50), 100, 3, 3) |
| check_rnn_cell(cell, in_shape=(1, 10, 20, 30, 50), out_shape=(1, 100, 18, 28, 48)) |
| |
| |
| @mx.util.use_np |
| def test_convgru(): |
| cell = gluon.rnn.Conv1DGRUCell((10, 50), 100, 3, 3) |
| check_rnn_cell(cell, in_shape=(1, 10, 50), out_shape=(1, 100, 48)) |
| |
| cell = gluon.rnn.Conv2DGRUCell((10, 20, 50), 100, 3, 3) |
| check_rnn_cell(cell, in_shape=(1, 10, 20, 50), out_shape=(1, 100, 18, 48)) |
| |
| cell = gluon.rnn.Conv3DGRUCell((10, 20, 30, 50), 100, 3, 3) |
| check_rnn_cell(cell, in_shape=(1, 10, 20, 30, 50), out_shape=(1, 100, 18, 28, 48)) |
| |
| |
| @mx.util.use_np |
| def test_conv_fill_shape(): |
| input = mx.np.ones((8, 3, 5, 7)) |
| cell = gluon.rnn.Conv1DLSTMCell((-1, 7), 10, (3,), (3,)) |
| cell.infer_shape(0, input, False) |
| cell.hybridize() |
| check_rnn_forward(cell, input) |
| assert cell.i2h_weight.shape[1] == 5, cell.i2h_weight.shape[1] |
| |
| |
| @mx.util.use_np |
| def test_lstmp(): |
| nhid = 100 |
| nproj = 64 |
| cell = gluon.rnn.LSTMPCell(nhid, nproj) |
| inputs = [mx.np.ones(shape=(10,50)) for i in range(3)] |
| cell.infer_shape(0, inputs[0], False) |
| cell.initialize() |
| outputs, _ = cell.unroll(3, inputs) |
| expected_params = ['h2h_bias', 'h2h_weight', 'h2r_weight', 'i2h_bias', 'i2h_weight'] |
| assert sorted(cell.collect_params().keys()) == expected_params |
| |
| outs = [o.shape for o in outputs] |
| assert outs == [(10, nproj)] * 3 |
| |
| |
| @mx.util.use_np |
| def test_vardrop(): |
| def check_vardrop(drop_inputs, drop_states, drop_outputs): |
| cell = gluon.rnn.VariationalDropoutCell(mx.gluon.rnn.RNNCell(100), |
| drop_outputs=drop_outputs, |
| drop_states=drop_states, |
| drop_inputs=drop_inputs) |
| |
| input_data = mx.np.random.uniform(size=(10, 3, 50), device=mx.device.current_device()) |
| cell.infer_shape(0, input_data, False) |
| cell.initialize(init='xavier') |
| with mx.autograd.record(): |
| outputs1, _ = cell.unroll(3, input_data, merge_outputs=True) |
| mx.npx.waitall() |
| outputs2, _ = cell.unroll(3, input_data, merge_outputs=True) |
| assert not almost_equal(outputs1.asnumpy(), outputs2.asnumpy()) |
| |
| inputs = [mx.np.ones(shape=(10,50)) for i in range(3)] |
| cell.infer_shape(0, inputs[0], False) |
| cell.initialize() |
| outputs, _ = cell.unroll(3, inputs, merge_outputs=False) |
| |
| outs = [o.shape for o in outputs] |
| assert outs == [(10, 100), (10, 100), (10, 100)] |
| |
| cell.reset() |
| cell.hybridize() |
| with mx.autograd.record(): |
| outputs3, _ = cell.unroll(3, input_data, merge_outputs=True) |
| mx.npx.waitall() |
| outputs4, _ = cell.unroll(3, input_data, merge_outputs=True) |
| assert not almost_equal(outputs3.asnumpy(), outputs4.asnumpy()) |
| assert not almost_equal(outputs1.asnumpy(), outputs3.asnumpy()) |
| |
| check_vardrop(0.5, 0.5, 0.5) |
| check_vardrop(0.5, 0, 0.5) |
| |
| |
| @mx.util.use_np |
| @pytest.mark.parametrize('cell_type,num_states', [ |
| (gluon.rnn.RNNCell, 1), |
| (gluon.rnn.LSTMCell, 2), |
| (gluon.rnn.GRUCell, 1) |
| ]) |
| @pytest.mark.parametrize('layout', ['NTC', 'TNC']) |
| def test_unroll(cell_type, num_states, layout): |
| class RNNLayer(gluon.HybridBlock): |
| def __init__(self, cell_type, hidden_size, layout): |
| super(RNNLayer, self).__init__() |
| self.cell = cell_type(hidden_size) |
| self.layout = layout |
| |
| def forward(self, inputs, states, valid_length): |
| if isinstance(valid_length, list) and len(valid_length) == 0: |
| valid_length = None |
| return gluon.rnn.rnn_cell.dynamic_unroll(self.cell, inputs, states, |
| valid_length=valid_length, |
| layout=self.layout) |
| |
| def infer_shape(self, x, *args): |
| self.cell.infer_shape(0, x, False) |
| |
| batch_size = 20 |
| input_size = 50 |
| hidden_size = 30 |
| seq_len = 10 |
| device = default_device() |
| if layout == 'TNC': |
| rnn_data = mx.np.random.normal(loc=0, scale=1, size=(seq_len, batch_size, input_size), device=device) |
| elif layout == 'NTC': |
| rnn_data = mx.np.random.normal(loc=0, scale=1, size=(batch_size, seq_len, input_size), device=device) |
| else: |
| print("Wrong layout") |
| return |
| valid_length = mx.np.round(mx.np.random.uniform(low=1, high=10, size=(batch_size), device=device)) |
| state_shape = (batch_size, hidden_size) |
| states = [mx.np.random.normal(loc=0, scale=1, size=state_shape, device=device) for i in range(num_states)] |
| |
| cell = cell_type(hidden_size) |
| if layout == 'TNC': |
| cell.infer_shape(0, rnn_data[0], False) |
| cell.initialize(device=default_device()) |
| cell(rnn_data[0], states) |
| else: |
| cell.infer_shape(0, rnn_data[:,0,:], False) |
| cell.initialize(device=default_device()) |
| cell(rnn_data[:,0,:], states) |
| params1 = cell.collect_params() |
| orig_params1 = copy.deepcopy(params1) |
| |
| trainer = gluon.Trainer(params1, 'sgd', {'learning_rate' : 0.03}) |
| with mx.autograd.record(): |
| res1, states1 = cell.unroll(seq_len, rnn_data, states, valid_length=valid_length, |
| layout=layout, merge_outputs=True) |
| res1.backward() |
| trainer.step(batch_size) |
| |
| configs = [ |
| lambda layer: None, |
| lambda layer: layer.hybridize(), |
| lambda layer: layer.hybridize({'inline_limit': 0}), |
| lambda layer: layer.hybridize({'static_alloc': True}), |
| lambda layer: layer.hybridize({'static_alloc': True, 'static_shape': True}) ] |
| # We can't pass None to a hybrid block, but it accepts an empty list. |
| # so we use an empty list to represent valid_length if it's None. |
| if valid_length is None: |
| valid_length = [] |
| for config in configs: |
| layer = RNNLayer(cell_type, hidden_size, layout) |
| layer.infer_shape(rnn_data) |
| layer.initialize(device=default_device()) |
| config(layer) |
| res2, states2 = layer(rnn_data, states, valid_length) |
| params2 = layer.collect_params() |
| for key, val in orig_params1.items(): |
| params2['cell.' + key].set_data(copy.deepcopy(val.data())) |
| |
| trainer = gluon.Trainer(params2, 'sgd', {'learning_rate' : 0.03}) |
| with mx.autograd.record(): |
| res2, states2 = layer(rnn_data, states, valid_length) |
| assert_almost_equal(res1, res2, rtol=0.001, atol=0.0001) |
| assert len(states1) == len(states2) |
| for i in range(len(states1)): |
| assert_almost_equal(states1[i], states2[i], rtol=0.001, atol=0.0001) |
| res2.backward() |
| trainer.step(batch_size) |
| |
| for key, val in params1.items(): |
| weight1 = val.data() |
| weight2 = params2['cell.' + key].data() |
| # Subgraph created from npx.foreach in deferred compute is |
| # little bit different from the legacy foreach operator. |
| assert_almost_equal(weight1, weight2, rtol=0.1, atol=0.1) |
