numpyrandompytorchinitializationdeterministic

Initialise Pytorch layer with local random number generator


When writing larger programs that require determinism for random processes, it is generally considered good practice to create function-specific random number generators (rngs) and pass those to the randomness-dependent functions (rather than setting a global seed and have the functions depend on this). See also here.

For example, when I have a function that generates some sample using numpy, I use a rng that I create at the beginning of the script:

# at the beginning of the script
import numpy as np
seed = 5465
rng = np.random.default_rng(seed)
# much later in the script
def generate_sample(rng, size):
  return rng.random(size)
generate_sample(rng, size=5)

I am trying to achieve the same when initialising a torch.nn.Linear layer, i.e. use a pre-defined rng to reproducibly initialise the layer.

Is this (reasonably) possible or am I forced to set the Pytorch global seed via torch.manual_seed() instead?


Solution

  • Generators are available in Pytorch, and in-place prngs can rely on this input, e.g. the normal distribution generator.

    When it comes to torch.nn.Linear, it has no such parameter, as you observed. You can still manually manage the internal state of the global prng by using a decorator that saves the internal state, loads the desired state, execute your nn.Linear, saves its state, and reset the global prng to its original state. By using this, you rely on roch.manual_seed() but also on the state getter and setter functions to reach states further from the seeded states.