Trainer for basic training loop. Args: models (dict[str, nn.Module]): Models to train. dataset (torch.utils.data.Dataset): Dataset. output_dir (str): Output directory. load_dir (str): Load directory. step (int): Step to load. batch_size (
| 16 | |
| 17 | |
| 18 | class BasicTrainer(Trainer): |
| 19 | """ |
| 20 | Trainer for basic training loop. |
| 21 | |
| 22 | Args: |
| 23 | models (dict[str, nn.Module]): Models to train. |
| 24 | dataset (torch.utils.data.Dataset): Dataset. |
| 25 | output_dir (str): Output directory. |
| 26 | load_dir (str): Load directory. |
| 27 | step (int): Step to load. |
| 28 | batch_size (int): Batch size. |
| 29 | batch_size_per_gpu (int): Batch size per GPU. If specified, batch_size will be ignored. |
| 30 | batch_split (int): Split batch with gradient accumulation. |
| 31 | max_steps (int): Max steps. |
| 32 | optimizer (dict): Optimizer config. |
| 33 | lr_scheduler (dict): Learning rate scheduler config. |
| 34 | elastic (dict): Elastic memory management config. |
| 35 | grad_clip (float or dict): Gradient clip config. |
| 36 | ema_rate (float or list): Exponential moving average rates. |
| 37 | fp16_mode (str): FP16 mode. |
| 38 | - None: No FP16. |
| 39 | - 'inflat_all': Hold a inflated fp32 master param for all params. |
| 40 | - 'amp': Automatic mixed precision. |
| 41 | fp16_scale_growth (float): Scale growth for FP16 gradient backpropagation. |
| 42 | finetune_ckpt (dict): Finetune checkpoint. |
| 43 | log_param_stats (bool): Log parameter stats. |
| 44 | i_print (int): Print interval. |
| 45 | i_log (int): Log interval. |
| 46 | i_sample (int): Sample interval. |
| 47 | i_save (int): Save interval. |
| 48 | i_ddpcheck (int): DDP check interval. |
| 49 | """ |
| 50 | |
| 51 | def __str__(self): |
| 52 | lines = [] |
| 53 | lines.append(self.__class__.__name__) |
| 54 | lines.append(f' - Models:') |
| 55 | for name, model in self.models.items(): |
| 56 | lines.append(f' - {name}: {model.__class__.__name__}') |
| 57 | lines.append(f' - Dataset: {indent(str(self.dataset), 2)}') |
| 58 | lines.append(f' - Dataloader:') |
| 59 | lines.append(f' - Sampler: {self.dataloader.sampler.__class__.__name__}') |
| 60 | lines.append(f' - Num workers: {self.dataloader.num_workers}') |
| 61 | lines.append(f' - Number of steps: {self.max_steps}') |
| 62 | lines.append(f' - Number of GPUs: {self.world_size}') |
| 63 | lines.append(f' - Batch size: {self.batch_size}') |
| 64 | lines.append(f' - Batch size per GPU: {self.batch_size_per_gpu}') |
| 65 | lines.append(f' - Batch split: {self.batch_split}') |
| 66 | lines.append(f' - Optimizer: {self.optimizer.__class__.__name__}') |
| 67 | lines.append(f' - Learning rate: {self.optimizer.param_groups[0]["lr"]}') |
| 68 | if self.lr_scheduler_config is not None: |
| 69 | lines.append(f' - LR scheduler: {self.lr_scheduler.__class__.__name__}') |
| 70 | if self.elastic_controller_config is not None: |
| 71 | lines.append(f' - Elastic memory: {indent(str(self.elastic_controller), 2)}') |
| 72 | if self.grad_clip is not None: |
| 73 | lines.append(f' - Gradient clip: {indent(str(self.grad_clip), 2)}') |
| 74 | lines.append(f' - EMA rate: {self.ema_rate}') |
| 75 | lines.append(f' - FP16 mode: {self.fp16_mode}') |
nothing calls this directly
no outgoing calls
no test coverage detected