colossalai.context.random

colossalai.context.random.seed(parallel_mode)[source]

A context for seed switch

Examples:

with seed(ParallelMode.DATA):
    output = F.dropout(input)

Note

The parallel_mode should be concluded in ParallelMode. More details about ParallelMode could be found in parallel_mode.

colossalai.context.random.set_mode(parallel_mode)[source]

Sets the current mode of the seed manager.

Parameters: parallel_mode (colossalai.context.ParallelMode) – The chosen parallel mode.

Note

The parallel_mode should be concluded in ParallelMode. More details about ParallelMode could be found in parallel_mode.

colossalai.context.random.with_seed(func, parallel_mode)[source]

A function wrapper which executes the function with a specified seed.

Examples:

# use with decorator
@with_seed(ParallelMode.DATA)
def forward(input):
    return F.dropout(input)
out = forward(input)
# OR use it inline
def forward(input):
    return F.dropout(input)
wrapper_forward = with_seed(forward, ParallelMode.DATA)
out = wrapped_forward(input)

Note

The parallel_mode should be concluded in ParallelMode. More details about ParallelMode could be found in parallel_mode.

colossalai.context.random.add_seed(parallel_mode, seed, overwrite=False)[source]

Adds a seed to the seed manager for parallel_mode.

Parameters

parallel_mode (colossalai.context.ParallelMode) – The chosen parallel mode.
seed (int) – The seed to be added

Raises

AssertionError – Raises an AssertionError if parallel_mode is not an instance of colossalai.context.ParallelMode or the seed for parallel_mode has been added.

Note

The parallel_mode should be concluded in ParallelMode. More details about ParallelMode could be found in parallel_mode.

colossalai.context.random.get_seeds()[source]

Returns the seeds of the seed manager.

Returns: The seeds of the seed manager.
Return type: dict

colossalai.context.random.get_states(copy=False)[source]

Returns the seed states of the seed manager.

Returns: The seed states of the seed manager.
Return type: dict

colossalai.context.random.get_current_mode()[source]

Returns the current mode of the seed manager.

Returns: The current mode of the seed manager.
Return type: torch.ByteTensor

colossalai.context.random.set_seed_states(parallel_mode, state)[source]

Sets the state of the seed manager for parallel_mode.

Parameters

parallel_mode (colossalai.context.ParallelMode) – The chosen parallel mode.
state (torch.Tensor) – the state to be set.

Raises

AssertionError – Raises an AssertionError if parallel_mode is not found in the seed manager.

colossalai.context.random.seed_manager