Method test_predict

axlearn/audio/decoder_asr_test.py:384–436

(self, blank_id)

    @parameterized.parameters([0, 1])
    def test_predict(self, blank_id):
        input_dim, vocab_size = 6, 8
        cfg = CTCDecoderModel.default_config().set(
            input_dim=input_dim,
            vocab_size=vocab_size,
            blank_id=blank_id,
        )
        # Initialize layer parameters.
        layer: CTCDecoderModel = cfg.set(name="test").instantiate(parent=None)
        prng_key = jax.random.PRNGKey(123)
        prng_key, init_key = jax.random.split(prng_key)
        layer_params = layer.initialize_parameters_recursively(init_key)
        self.assertEqual(
            {"lm_head": dict(weight=(input_dim, vocab_size), bias=(vocab_size,))},
            shapes(layer_params),
        )

        batch_size, max_seq_len = 2, 10
        seq_len = jnp.array([7, 5])
        # [batch_size, max_seq_len, dim] with the same data across sequences.
        inputs = jnp.tile(
            jax.random.normal(jax.random.PRNGKey(123), [1, max_seq_len, input_dim]),
            [batch_size, 1, 1],
        )
        # [batch_size, max_seq_len].
        paddings = jnp.arange(max_seq_len) >= seq_len[:, None]

        # Generate different padding data.
        padding_data = jax.random.normal(
            jax.random.PRNGKey(130), [batch_size, max_seq_len, input_dim]
        )
        # Generate input sequences with the same data at non-pad positions.
        inputs = jnp.where(paddings[..., None], padding_data, inputs)

        @jax.jit
        def jit_predict(input_batch):
            outputs, _ = F(
                layer,
                inputs=dict(input_batch=input_batch),
                is_training=True,
                prng_key=prng_key,
                state=layer_params,
                method="predict",
            )
            return outputs

        outputs = jit_predict(dict(inputs=inputs, paddings=paddings))
        self.assertSequenceEqual((batch_size, max_seq_len, vocab_size), outputs.shape)
        # Check that the outputs are the same in the non-padded positions.
        assert_allclose(outputs[0, : seq_len[1]], outputs[1, : seq_len[1]])
        outputs_at_padding = paddings[:, :, None] * outputs
        # Check that all padding position have 0 outputs.
        self.assertTrue(jnp.all(jnp.logical_not(outputs_at_padding)))

    @parameterized.parameters([0, 1])
    def test_forward(self, blank_id):