trainer package

Subpackages

• trainer.optimization package
  • Submodules
  • trainer.optimization.adafactor module
  • trainer.optimization.adamw module
  • trainer.optimization.optimization module
  • Module contents
• trainer.utils package
  • Submodules
  • trainer.utils.layer_freezer module
  • trainer.utils.trainer_utils module
  • Module contents

Submodules

trainer.callback module

trainer.metrics module

trainer.modelcard module

trainer.scheduler module

Scheduler utilities for pytorch optimization.

trainer.scheduler.check_scheduler(scheduler_type: str) → bool

Check if the scheduler type is supported.

Parameters

scheduler_type (str) – the type of the scheduler.

Return (bool):

if the scheduler type is supported.

Example

>>> from towhee.trainer.scheduler import check_scheduler
>>> check_scheduler('constant')
True

trainer.scheduler.configure_constant_scheduler(optimizer: Optimizer, last_epoch: int = - 1)

Return a scheduler with a constant learning rate, using the learning rate set in optimizer.

Parameters

• optimizer (Optimizer) – The optimizer for which to schedule the learning rate.

• last_epoch (int) – The last epoch when resuming training.

Return (LambdaLR):

A constant scheduler

Example

>>> from towhee.trainer.scheduler import configure_constant_scheduler
>>> from towhee.trainer.optimization.adamw import AdamW
>>> from torch import nn
>>> def unwrap_scheduler(scheduler, num_steps=10):
>>>     lr_sch = []
>>>     for _ in range(num_steps):
>>>         lr_sch.append(scheduler.get_lr()[0])
>>>         scheduler.step()
>>>     return lr_sch
>>> mdl = nn.Linear(50, 50)
>>> optimizer = AdamW(mdl.parameters(), lr=10.0)
>>> num_steps = 2
>>> scheduler = configure_constant_scheduler(optimizer)
>>> lr_sch_1 = unwrap_scheduler(scheduler, num_steps)
[10.0, 10.0]

trainer.scheduler.configure_constant_scheduler_with_warmup(optimizer: Optimizer, num_warmup_steps: int, last_epoch: int = - 1)

Return a schedule with a constant learning rate preceded by a warmup period during which the learning rate increases linearly between 0 and the initial lr set in the optimizer.

Parameters

• optimizer (Optimizer) – The optimizer to be scheduled.

• num_warmup_steps (int) – Warmup steps.

• last_epoch (int) – The last epoch when training is resumed.

Return (LambdaLR):

A constant scheduler with warmup.

Example

>>> from towhee.trainer.scheduler import configure_constant_scheduler_with_warmup
>>> from towhee.trainer.optimization.adamw import AdamW
>>> from torch import nn
>>> def unwrap_scheduler(scheduler, num_steps=10):
>>>     lr_sch = []
>>>     for _ in range(num_steps):
>>>         lr_sch.append(scheduler.get_lr()[0])
>>>         scheduler.step()
>>>     return lr_sch
>>> mdl = nn.Linear(50, 50)
>>> optimizer = AdamW(mdl.parameters(), lr=10.0)
>>> num_steps = 10
>>> num_warmup_steps = 4
>>> scheduler = configure_constant_scheduler_with_warmup(optimizer, num_warmup_steps)
>>> lr_sch_1 = unwrap_scheduler(scheduler, num_steps)
[0.0, 2.5, 5.0, 7.5, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0]

trainer.scheduler.configure_cosine_scheduler_with_warmup(optimizer: Optimizer, num_warmup_steps: int, num_training_steps: int, num_cycles: float = 0.5, last_epoch: int = - 1)

Return a scheduler with a learning rate that decreases following the values of the cosine function between the initial lr set in the optimizer to 0, after a warmup period during which it increases linearly between 0 and the initial lr set in the optimizer.

Parameters

• optimizer (Optimizer) – The optimizer to be scheduled.

• num_warmup_steps (int) – The steps for the warmup phase.

• num_training_steps (int) – The number of training steps.

• num_cycles (int) – The number of periods in te cosine scheduler.

• last_epoch (int) – The last epoch when training is resumed.

Return (LambdaLR):

A cosine scheduler with warmup.

Example

>>> from towhee.trainer.scheduler import configure_cosine_scheduler_with_warmup
>>> from towhee.trainer.optimization.adamw import AdamW
>>> from torch import nn
>>> def unwrap_scheduler(scheduler, num_steps=10):
>>>     lr_sch = []
>>>     for _ in range(num_steps):
>>>         lr_sch.append(scheduler.get_lr()[0])
>>>         scheduler.step()
>>>     return lr_sch
>>> mdl = nn.Linear(50, 50)
>>> optimizer = AdamW(mdl.parameters(), lr=10.0)
>>> num_steps = 10
>>> num_warmup_steps = 4
>>> num_training_steps = 10
>>> scheduler = configure_cosine_scheduler_with_warmup(optimizer, num_warmup_steps, num_training_steps)
>>> lr_sch_1 = unwrap_scheduler(scheduler, num_steps)
[0.0, 5.0, 10.0, 9.61, 8.53, 6.91, 5.0, 3.08, 1.46, 0.38]

trainer.scheduler.configure_cosine_with_hard_restarts_scheduler_with_warmup(optimizer: Optimizer, num_warmup_steps: int, num_training_steps: int, num_cycles: int = 1, last_epoch: int = - 1)

Return a scheduler with a learning rate that decreases following the values of the cosine function between the initial lr set in the optimizer to 0, with several hard restarts, after a warmup period during which it increases linearly between 0 and the initial lr set in the optimizer.

Parameters

• optimizer (Optimizer) – The optimizer to be scheduled.

• num_warmup_steps (int) – The steps for the warmup phase.

• num_training_steps (int) – The number of training steps.

• num_cycles (int) – The number of hard restarts to be used.

• last_epoch (int) – The index of the last epoch when training is resumed.

Return (LambdaLR):

A cosine with hard restarts scheduler with warmup.

Example

>>> from towhee.trainer.scheduler import configure_cosine_with_hard_restarts_scheduler_with_warmup
>>> from towhee.trainer.optimization.adamw import AdamW
>>> from torch import nn
>>> def unwrap_scheduler(scheduler, num_steps=10):
>>>     lr_sch = []
>>>     for _ in range(num_steps):
>>>         lr_sch.append(scheduler.get_lr()[0])
>>>         scheduler.step()
>>>     return lr_sch
>>> mdl = nn.Linear(50, 50)
>>> optimizer = AdamW(mdl.parameters(), lr=10.0)
>>> num_steps = 10
>>> num_warmup_steps = 4
>>> num_training_steps = 10
>>> num_cycles = 2
>>> scheduler = configure_cosine_with_hard_restarts_scheduler_with_warmup(optimizer,
num_warmup_steps, num_training_steps, num_cycles)
>>> lr_sch_1 = unwrap_scheduler(scheduler, num_steps)
[0.0, 5.0, 10.0, 8.53, 5.0, 1.46, 10.0, 8.53, 5.0, 1.46]

trainer.scheduler.configure_linear_scheduler_with_warmup(optimizer, num_warmup_steps, num_training_steps, last_epoch=- 1)

Return a scheduler with a learning rate that decreases linearly from the initial lr set in the optimizer to 0, after a warmup period during which it increases linearly from 0 to the initial lr set in the optimizer.

Parameters

• optimizer (Optimizer) – The optimizer to be scheduled.

• num_warmup_steps (int) – Warmup steps.

• num_training_steps (int) – Training steps.

• last_epoch (int) – The last epoch when training is resumed.

Return (LambdaLR):

A linear scheduler with warmup.

Example

>>> from towhee.trainer.scheduler import configure_linear_scheduler_with_warmup
>>> from towhee.trainer.optimization.adamw import AdamW
>>> from torch import nn
>>> def unwrap_scheduler(scheduler, num_steps=10):
>>>     lr_sch = []
>>>     for _ in range(num_steps):
>>>         lr_sch.append(scheduler.get_lr()[0])
>>>         scheduler.step()
>>>     return lr_sch
>>> mdl = nn.Linear(50, 50)
>>> optimizer = AdamW(mdl.parameters(), lr=10.0)
>>> num_steps = 10
>>> num_warmup_steps = 4
>>> num_training_steps = 10
>>> scheduler = configure_constant_scheduler_with_warmup(optimizer, num_warmup_steps, num_training_steps)
>>> lr_sch_1 = unwrap_scheduler(scheduler, num_steps)
[0.0, 2.5, 5.0, 7.5, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0]

trainer.scheduler.configure_polynomial_decay_scheduler_with_warmup(optimizer, num_warmup_steps, num_training_steps, lr_end=1e-07, power=1.0, last_epoch=- 1)

Return a scheduler with a learning rate that decreases as a polynomial decay from the initial lr set in the optimizer to end lr defined by lr_end, after a warmup period during which it increases linearly from 0 to the initial lr set in the optimizer.

Parameters

• optimizer (Optimizer) – The optimizer to be scheduled.

• num_warmup_steps (int) – The steps for the warmup phase.

• num_training_steps (int) – The number of training steps

• lr_end (float) – The end LR.

• power (float) – Power factor.

• last_epoch (int) – The index of the last epoch when training is resumed.

Return (LambdaLR):

A polynomial decay scheduler with warmup.

Example

>>> from towhee.trainer.scheduler import configure_polynomial_decay_scheduler_with_warmup
>>> from towhee.trainer.optimization.adamw import AdamW
>>> from torch import nn
>>> def unwrap_scheduler(scheduler, num_steps=10):
>>>     lr_sch = []
>>>     for _ in range(num_steps):
>>>         lr_sch.append(scheduler.get_lr()[0])
>>>         scheduler.step()
>>>     return lr_sch
>>> mdl = nn.Linear(50, 50)
>>> optimizer = AdamW(mdl.parameters(), lr=10.0)
>>> num_steps = 10
>>> num_warmup_steps = 4
>>> num_training_steps = 10
>>> power = 2.0
>>> lr_end = 1e-7
>>> scheduler = configure_polynomial_decay_scheduler_with_warmup(optimizer,
num_warmup_steps, num_training_steps, num_cycles)
>>> lr_sch_1 = unwrap_scheduler(scheduler, num_steps)
[0.0, 5.0, 10.0, 7.656, 5.625, 3.906, 2.5, 1.406, 0.625, 0.156]

trainer.trainer module

trainer.training_config module

class trainer.training_config.TrainingConfig(output_dir: str = './output_dir', overwrite_output_dir: bool = True, eval_strategy: str = 'epoch', eval_steps: ~typing.Optional[int] = None, batch_size: ~typing.Optional[int] = 8, val_batch_size: ~typing.Optional[int] = -1, seed: int = 42, epoch_num: int = 2, dataloader_pin_memory: bool = True, dataloader_drop_last: bool = True, dataloader_num_workers: int = 0, lr: float = 5e-05, metric: ~typing.Optional[str] = 'Accuracy', print_steps: ~typing.Optional[int] = None, load_best_model_at_end: ~typing.Optional[bool] = False, early_stopping: ~typing.Union[dict, str] = <factory>, model_checkpoint: ~typing.Union[dict, str] = <factory>, tensorboard: ~typing.Optional[~typing.Union[dict, str]] = <factory>, loss: ~typing.Union[str, ~typing.Dict[str, ~typing.Any]] = 'CrossEntropyLoss', optimizer: ~typing.Union[str, ~typing.Dict[str, ~typing.Any]] = 'Adam', lr_scheduler_type: str = 'linear', warmup_ratio: float = 0.0, warmup_steps: int = 0, device_str: ~typing.Optional[str] = None, sync_bn: bool = False, freeze_bn: bool = False)

Bases: object

The training config, it can be defined in a yaml file

Parameters

• output_dir (str) – The output directory where the model predictions and checkpoints will be written.

• overwrite_output_dir (bool) – Overwrite the content of the output directory.

• eval_strategy (str) – The evaluation strategy.

• eval_steps (int) – Run an evaluation every X steps.

• batch_size (int) – Batch size for training.

• val_batch_size (int) – Batch size for evaluation.

• seed (int) – Random seed that will be set at the beginning of training.

• epoch_num (int) – Total number of training epochs to perform.

• dataloader_pin_memory (bool) – Drop the last incomplete batch if it is not divisible by the batch size.

• dataloader_drop_last (bool) – Drop the last incomplete batch if it is not divisible by the batch size.

• dataloader_num_workers (int) – Number of subprocesses to use for data loading.

• lr (float) – The initial learning rate for AdamW.

• metric (str) – The metric to use to compare two different models.

• print_steps (int) – If None, use the tqdm progress bar, otherwise it will print the logs on the screen every print_steps.

• load_best_model_at_end (bool) – Whether or not to load the best model found during training at the end of training.

• early_stopping (Union[dict, str]) – Early stopping.

• model_checkpoint (Union[dict, str]) – Model checkpoint.

• tensorboard (Union[dict, str]) – Tensorboard.

• loss (Union[str, Dict[str, Any]]) – Pytorch loss in torch.nn package.

• optimizer (Union[str, Dict[str, Any]]) – Pytorch optimizer Class name in torch.optim package.

• lr_scheduler_type (str) – The scheduler type to use.

• warmup_ratio (float) – Linear warmup over warmup_ratio fraction of total steps.

• device_str (str) – Device string.

• sync_bn (bool) – It will be work if device_str is cuda, the True sync_bn would make training slower but acc better.

• freeze_bn (bool) – It will completely freeze all BatchNorm layers during training.

load_from_yaml(path2yaml: Optional[str] = None) → TrainingConfig

Load training configuration from yaml.

Parameters

path2yaml (str) – The path to yaml.

Returns

(TrainingConfig).

TrainingConfig instance self.

Example

>>> from towhee.trainer.training_config import TrainingConfig
>>> from pathlib import Path
>>> conf = Path(__file__).parent / "config.yaml"
>>> ta = TrainingConfig()
>>> ta.save_to_yaml(conf)
>>> ta.load_from_yaml(conf)
>>> ta.epoch_num
2

save_to_yaml(path2yaml: Optional[str] = None)

Save training configuration to yaml.

Parameters

path2yaml (str) – The path to yaml.

Example

>>> from towhee.trainer.training_config import TrainingConfig
>>> from pathlib import Path
>>> conf = Path(__file__).parent / 'config.yaml'
>>> ta = TrainingConfig()
>>> ta.save_to_yaml(conf)
>>> ta.load_from_yaml(conf)
>>> ta.epoch_num
2

trainer.training_config.dump_default_yaml(yaml_path)

Dump a default yaml, which can be overridden by the custom operator.

trainer.training_config.get_config_help()

Get config setting infos. :returns:

(dict)

The help dict.

Module contents