base

base#

class fkat.pytorch.schedule.base.Always(*args, **kwargs)[source]#

A schedule for an event that always happens.

check(*, stage: Optional[str] = None, batch_idx: Optional[int] = None, step: Optional[int] = None, trainer: Optional[Trainer] = None) → bool[source]#

Checks the schedule for the given moment.

Parameters:

  • stage (str) – current trainer stage. eg. train/test/validate/predict

  • batch_idx (Optional[int]) – current batch_idx

  • step (Optional[int]) – current step

  • trainer (Optional[Trainer]) – lightning trainer of callback

Returns:

True if schedule passed the check, False otherwise

Return type:

bool

class fkat.pytorch.schedule.base.CombinedSchedule(fn: Callable[[bool, bool], bool], schedules: Sequence[Schedule])[source]#
check(*, stage: Optional[str] = None, batch_idx: Optional[int] = None, step: Optional[int] = None, trainer: Optional[Trainer] = None) → bool[source]#

Checks the schedule for the given moment.

Parameters:

  • stage (str) – current trainer stage. eg. train/test/validate/predict

  • batch_idx (Optional[int]) – current batch_idx

  • step (Optional[int]) – current step

  • trainer (Optional[Trainer]) – lightning trainer of callback

Returns:

True if schedule passed the check, False otherwise

Return type:

bool

class fkat.pytorch.schedule.base.Elapsed(interval: timedelta)[source]#

A schedule for an event that happens after the provided time interval has elapsed.

check(*, stage: Optional[str] = None, batch_idx: Optional[int] = None, step: Optional[int] = None, trainer: Optional[Trainer] = None) → bool[source]#

Checks the schedule for the given moment.

Parameters:

  • stage (str) – current trainer stage. eg. train/test/validate/predict

  • batch_idx (Optional[int]) – current batch_idx

  • step (Optional[int]) – current step

  • trainer (Optional[Trainer]) – lightning trainer of callback

Returns:

True if schedule passed the check, False otherwise

Return type:

bool

class fkat.pytorch.schedule.base.Every(*, n_batches: int = 0, n_steps: int = 0, stage: Optional[str] = None)[source]#

A schedule for an event that happens every specified number of batches and/or steps.

n_batches#

A positive number of batches between logging events. Defaults to 0 - use only n_steps

Type:

Optional[int]

n_steps#

A positive number of (train) steps between logging events. Defaults to 0 - use only n_batches

Type:

Optional[int]

stage#

The stage this schedule applies to (‘train’, ‘validation’, ‘test’, ‘predict’). If None, applies to all stages.

Type:

Optional[str]

check(*, stage: Optional[str] = None, batch_idx: Optional[int] = None, step: Optional[int] = None, trainer: Optional[Trainer] = None) → bool[source]#

Checks the schedule for the given moment.

Parameters:

  • stage (str) – current trainer stage. eg. train/test/validate/predict

  • batch_idx (Optional[int]) – current batch_idx

  • step (Optional[int]) – current step

  • trainer (Optional[Trainer]) – lightning trainer of callback

Returns:

True if schedule passed the check, False otherwise

Return type:

bool

class fkat.pytorch.schedule.base.Fixed(warmup_steps: int, active_steps: int)[source]#

A schedule for an event that happens after a warmup period and lasts for a fixed number of steps.

warmup_steps#

Number of initial steps to skip before logging starts.

Type:

int

active_steps#

Number of steps to log after warmup period.

Type:

int

check(*, stage: Optional[str] = None, batch_idx: Optional[int] = None, step: Optional[int] = None, trainer: Optional[Trainer] = None) → bool[source]#

Checks the schedule for the given moment.

Parameters:

  • stage (str) – current trainer stage. eg. train/test/validate/predict

  • batch_idx (Optional[int]) – current batch_idx

  • step (Optional[int]) – current step

  • trainer (Optional[Trainer]) – lightning trainer of callback

Returns:

True if schedule passed the check, False otherwise

Return type:

bool

class fkat.pytorch.schedule.base.GlobalRank(*ranks: int)[source]#

A schedule that only executes on specific global ranks in a distributed training setup.

This is useful for operations that should only be performed on certain ranks, such as logging, checkpointing, or other operations that would be redundant or conflicting if performed on all ranks.

ranks#

The global ranks on which this schedule should execute.

Type:

tuple[int, …]

check(*, stage: Optional[str] = None, batch_idx: Optional[int] = None, step: Optional[int] = None, trainer: Optional[Trainer] = None) → bool[source]#

Check if the current global rank is in the specified ranks.

Parameters:

  • stage – Current trainer stage (ignored by this schedule).

  • batch_idx – Current batch index (ignored by this schedule).

  • step – Current step (ignored by this schedule).

  • trainer – Lightning trainer instance, used to get the global rank.

Returns:

True if the trainer’s global rank is in the specified ranks, False otherwise.

Always returns False if trainer is None.

Return type:

bool

class fkat.pytorch.schedule.base.InvertedSchedule(other: Schedule)[source]#
check(*, stage: Optional[str] = None, batch_idx: Optional[int] = None, step: Optional[int] = None, trainer: Optional[Trainer] = None) → bool[source]#

Checks the schedule for the given moment.

Parameters:

  • stage (str) – current trainer stage. eg. train/test/validate/predict

  • batch_idx (Optional[int]) – current batch_idx

  • step (Optional[int]) – current step

  • trainer (Optional[Trainer]) – lightning trainer of callback

Returns:

True if schedule passed the check, False otherwise

Return type:

bool

class fkat.pytorch.schedule.base.LocalRank(*ranks: int)[source]#

A schedule that only executes on specific local ranks in a distributed training setup.

This is useful for node-specific operations that should only be performed on certain ranks within each node, such as local logging or monitoring.

ranks#

The local ranks on which this schedule should execute.

Type:

tuple[int, …]

check(*, stage: Optional[str] = None, batch_idx: Optional[int] = None, step: Optional[int] = None, trainer: Optional[Trainer] = None) → bool[source]#

Check if the current local rank is in the specified ranks.

Parameters:

  • stage – Current trainer stage (ignored by this schedule).

  • batch_idx – Current batch index (ignored by this schedule).

  • step – Current step (ignored by this schedule).

  • trainer – Lightning trainer instance, used to get the local rank.

Returns:

True if the trainer’s local rank is in the specified ranks, False otherwise.

Always returns False if trainer is None.

Return type:

bool

class fkat.pytorch.schedule.base.Never(*args, **kwargs)[source]#

A schedule for an event that never happens.

check(*, stage: Optional[str] = None, batch_idx: Optional[int] = None, step: Optional[int] = None, trainer: Optional[Trainer] = None) → bool[source]#

Checks the schedule for the given moment.

Parameters:

  • stage (str) – current trainer stage. eg. train/test/validate/predict

  • batch_idx (Optional[int]) – current batch_idx

  • step (Optional[int]) – current step

  • trainer (Optional[Trainer]) – lightning trainer of callback

Returns:

True if schedule passed the check, False otherwise

Return type:

bool

class fkat.pytorch.schedule.base.Schedule(*args, **kwargs)[source]#

Protocol defining a generic PyTorch schedule

check(*, stage: Optional[str] = None, batch_idx: Optional[int] = None, step: Optional[int] = None, trainer: Optional[Trainer] = None) → bool[source]#

Checks the schedule for the given moment.

Parameters:

  • stage (str) – current trainer stage. eg. train/test/validate/predict

  • batch_idx (Optional[int]) – current batch_idx

  • step (Optional[int]) – current step

  • trainer (Optional[Trainer]) – lightning trainer of callback

Returns:

True if schedule passed the check, False otherwise

Return type:

bool

Contents