Linen api and randomness (dropout) example

[mortner31]

Hello,

I am trying to use dropout examples from the doc https://docs.google.com/document/d/1hYavTVPaKVVe9Be8pCB7yW7r6dDv3RALVNit8NZca4c/

The linen provided example is as follows.

class Dropout(nn.Module):
  rate: float

  def __call__(self, inputs, deterministic=False): # [1] [2]
    keep_prob = 1. - self.rate
    if deterministic:
      return inputs
    else:
      mask = random.bernoulli(
        self.make_rng('dropout'), p=keep_prob, shape=inputs.shape) 
      return lax.select(
        mask, inputs / keep_prob, jnp.zeros_like(inputs))

# [1] deterministic at application rather than at module creation ? 
# [2] missing nn.compact decorator ? 

For me it raises questions, I therefore changed the example like this :

class Dropout(nn.Module):
    rate: float
    deterministic: False # I put the deterministic here

    @nn.compact
    def __call__(self, inputs):
        keep_prob = 1. - self.rate

        if self.deterministic:
            return inputs
        else:
            mask = random.bernoulli(
                self.make_rng('dropout'), p=keep_prob, shape=inputs.shape)  # [1]
            return lax.select(
                mask, inputs / keep_prob, jnp.zeros_like(inputs))

My point is that I wanted to use the deterministic flag at module creation rather than usage like it is done in the ConvNet example:

class ResNet(nn.Module):
     """ResNetV1."""
     stage_sizes: Sequence[int]
     num_filters: int = 64
     train: bool = True

But maybe, I am shortsighted on something. Why would I apply a different logic for the ŧrain flag and the deterministic flag ?

[levskaya]

Hi @mortner31

I think we do have a typo in the doc, missing the @nn.compact decorator, thanks for pointing that out.

The answer to your general question is that it's a bit arbitrary, but depends on intended usage:

You could imagine wanting deterministic to be an apply-time argument, in case you wanted to use the same layer in multiple places both deterministically or not. But I agree that it's usually set by a global config hyperparameter, in which case having deterministic set at construction time would be fine. We could move it to be a construction arg if people find the current setup to be a weird choice, but it's just sort of a decision based on how you think the layer will be used. For the short-term, if our library version is awkward for your use case, I'd encourage forking it and moving it to a construction time argument.