Add scheduled sampling howto for seq2seq example.

docs/howtos.rst

@@ -45,6 +45,14 @@ Polyak averaging
 .. raw:: html
    :file: _formatted_howtos/polyak-averaging.diff.html
 
+Scheduled Sampling
+----------------
+
+⟶ `View as a side-by-side diff <https://github.com/google/flax/compare/master..howto/scheduled-sampling?diff=split>`_
+
+.. raw:: html
+   :file: _formatted_howtos/scheduled-sampling.diff.html
+
 
 How do HOWTOs work?
 -------------------

howtos/diffs/scheduled-sampling.diff

@@ -0,0 +1,138 @@
+diff --git a/examples/seq2seq/train.py b/examples/seq2seq/train.py
+index 9b64f88..69bed4f 100644
+--- a/examples/seq2seq/train.py
++++ b/examples/seq2seq/train.py
+@@ -186,17 +186,19 @@ class Encoder(nn.Module):
+ class Decoder(nn.Module):
+   """LSTM decoder."""
+
+-  def apply(self, init_state, inputs, teacher_force=False):
++  def apply(self, init_state, inputs, sample_probability=0.0):
+     # inputs.shape = (batch_size, seq_length, vocab_size).
+-    vocab_size = inputs.shape[2]
++    batch_size, _, vocab_size = inputs.shape
+     lstm_cell = nn.LSTMCell.partial(name='lstm')
+     projection = nn.Dense.partial(features=vocab_size, name='projection')
+
+     def decode_step_fn(carry, x):
+       rng, lstm_state, last_prediction = carry
+-      carry_rng, categorical_rng = jax.random.split(rng, 2)
+-      if not teacher_force:
+-        x = last_prediction
++      carry_rng, bernoulli_rng, categorical_rng = jax.random.split(rng, 3)
++      x = jnp.where(
++          jax.random.bernoulli(
++              bernoulli_rng, p=sample_probability, shape=(batch_size, 1)),
++          last_prediction, x)
+       lstm_state, y = lstm_cell(lstm_state, x)
+       logits = projection(y)
+       predicted_tokens = jax.random.categorical(categorical_rng, logits)
+@@ -223,7 +225,7 @@ class Seq2seq(nn.Module):
+   def apply(self,
+             encoder_inputs,
+             decoder_inputs,
+-            teacher_force=True,
++            sample_probability=0.0,
+             eos_id=1,
+             hidden_size=512):
+     """Run the seq2seq model.
+@@ -234,14 +236,15 @@ class Seq2seq(nn.Module):
+         `[batch_size, max(encoder_input_lengths), vocab_size]`.
+       decoder_inputs: padded batch of expected decoded sequences for teacher
+         forcing, shaped `[batch_size, max(decoder_inputs_length), vocab_size]`.
+-        When sampling (i.e., `teacher_force = False`), the initial time step is
+-        forced into the model and samples are used for the following inputs. The
+-        second dimension of this tensor determines how many steps will be
+-        decoded, regardless of the value of `teacher_force`.
+-      teacher_force: bool, whether to use `decoder_inputs` as input to the
+-        decoder at every step. If False, only the first input is used, followed
+-        by samples taken from the previous output logits.
+-      eos_id: int, the token signaling when the end of a sequence is reached.
++        The initial time step is forced into the model and samples are used for
++        the following inputs as determined by `sample_probability`. The second
++        dimension of this tensor determines how many steps will be
++        decoded, regardless of the value of `sample_probability`.
++      sample_probability: float in [0, 1], the probability of using the previous
++        sample as the next input instead of the value in `decoder_inputs` when
++        sampling. A value of 0 is equivalent to teacher forcing and 1 indicates
++        full sampling.
++      eos_id: int, the token signalling when the end of a sequence is reached.
+       hidden_size: int, the number of hidden dimensions in the encoder and
+         decoder LSTMs.
+     Returns:
+@@ -255,7 +258,7 @@ class Seq2seq(nn.Module):
+     logits, predictions = decoder(
+         init_decoder_state,
+         decoder_inputs[:, :-1],
+-        teacher_force=teacher_force)
++        sample_probability=sample_probability)
+
+     return logits, predictions
+
+@@ -325,14 +328,16 @@ def compute_metrics(logits, labels):
+
+
+ @jax.jit
+-def train_step(optimizer, batch, rng):
++def train_step(optimizer, batch, sample_probability, rng):
+   """Train one step."""
+   labels = batch['answer'][:, 1:]  # remove '=' start token
+
+   def loss_fn(model):
+     """Compute cross-entropy loss."""
+     with nn.stochastic(rng):
+-      logits, _ = model(batch['query'], batch['answer'])
++      logits, _ = model(
++          batch['query'], batch['answer'],
++          sample_probability=sample_probability)
+     loss = cross_entropy_loss(logits, labels, get_sequence_lengths(labels))
+     return loss, logits
+   grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
+@@ -356,7 +361,7 @@ def decode(model, inputs, rng):
+   init_decoder_inputs = jnp.tile(init_decoder_input,
+                                  (inputs.shape[0], get_max_output_len(), 1))
+   with nn.stochastic(rng):
+-    _, predictions = model(inputs, init_decoder_inputs, teacher_force=False)
++    _, predictions = model(inputs, init_decoder_inputs, sample_probability=1.0)
+   return predictions
+
+
+@@ -374,15 +379,22 @@ def decode_batch(model, batch_size):
+
+ def train_model():
+   """Train for a fixed number of steps and decode during training."""
++  def inv_sigmoid(i, k=500):
++    return k / (k + jnp.exp(i / k))
++
+   with nn.stochastic(jax.random.PRNGKey(0)):
+     model = create_model()
+     optimizer = create_optimizer(model, FLAGS.learning_rate)
+     for step in range(FLAGS.num_train_steps):
+       batch = get_batch(FLAGS.batch_size)
+-      optimizer, metrics = train_step(optimizer, batch, nn.make_rng())
++      sample_probability = 1 - inv_sigmoid(step)
++      optimizer, metrics = train_step(
++          optimizer, batch, sample_probability, nn.make_rng())
+       if step % FLAGS.decode_frequency == 0:
+-        logging.info('train step: %d, loss: %.4f, accuracy: %.2f', step,
+-                     metrics['loss'], metrics['accuracy'] * 100)
++        logging.info(
++            'train step: %d, sample prob: %.4f, loss: %.4f, accuracy: %.2f',
++            step, sample_probability, metrics['loss'],
++            metrics['accuracy'] * 100)
+         decode_batch(optimizer.target, 5)
+   return optimizer.target
+
+diff --git a/examples/seq2seq/train_test.py b/examples/seq2seq/train_test.py
+index 7efc1b1..122659e 100644
+--- a/examples/seq2seq/train_test.py
++++ b/examples/seq2seq/train_test.py
+@@ -103,7 +103,7 @@ class TrainTest(absltest.TestCase):
+       model = train.create_model()
+       optimizer = train.create_optimizer(model, 0.003)
+       optimizer, train_metrics = train.train_step(
+-          optimizer, batch, nn.make_rng())
++          optimizer, batch, 0.5, nn.make_rng())
+
+     self.assertLessEqual(train_metrics['loss'], 5)
+     self.assertGreaterEqual(train_metrics['accuracy'], 0)