Spaces:
Running on Zero
Running on Zero
File size: 6,250 Bytes
9368ee7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 | # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations
from typing import Callable, Any, TYPE_CHECKING
import torch
from torch.utils.data import DataLoader
from lipforcing.utils import instantiate
import lipforcing.utils.logging_utils as logger
if TYPE_CHECKING:
from lipforcing.configs.config import BaseConfig
from lipforcing.trainer import Trainer
from lipforcing.methods import FastGenModel
class CallbackDict:
def __init__(self, config: BaseConfig, trainer: Trainer):
self._callbacks = {}
callback_configs = config.trainer.callbacks
if callback_configs:
if isinstance(callback_configs, list):
logger.warning(msg="The 'config.trainer.callbacks' parameter should be a dict instead of a list. ")
callback_configs = {f"callback_{k}": v for k, v in enumerate(callback_configs)}
for callback_name, current_callback_cfg in callback_configs.items():
if "_target_" not in current_callback_cfg:
logger.critical(
f"Callback {callback_name} is missing the '_target_' field. \n Skip {current_callback_cfg}"
)
continue
logger.critical(f"Instantiating callback {callback_name}: {current_callback_cfg}")
_callback = instantiate(current_callback_cfg)
assert isinstance(_callback, Callback), f"{current_callback_cfg} is not a valid callback."
_callback.config = config
_callback.trainer = trainer
_callback.on_app_begin()
self._callbacks[callback_name] = _callback
def __getattr__(self, method_name: str) -> Callable:
def load_state_dict(state_dict: dict[str, Any]) -> None:
for name, callback in self._callbacks.items():
if name in state_dict:
callback.load_state_dict(state_dict[name])
else:
logger.warning(f"Callback {name} not found in checkpoint.")
def state_dict() -> dict[str, Any]:
return {name: self._callbacks[name].state_dict() for name in self._callbacks}
def callbacks_wrapper(*args, **kwargs):
for callback in self._callbacks.values():
assert hasattr(callback, method_name)
method = getattr(callback, method_name)
assert callable(method), f"{method_name} is not callable."
method(*args, **kwargs)
if method_name == "state_dict":
return state_dict
if method_name == "load_state_dict":
return load_state_dict
return callbacks_wrapper
class Callback:
config: "BaseConfig"
trainer: "Trainer"
def on_app_begin(self) -> None:
pass
def on_model_init_start(self, model: FastGenModel) -> None:
pass
def on_model_init_end(self, model: FastGenModel | torch.nn.parallel.DistributedDataParallel) -> None:
pass
def on_optimizer_init_start(self, model: FastGenModel) -> None:
pass
def on_optimizer_init_end(self, model: FastGenModel) -> None:
pass
def on_load_checkpoint_start(self, model: FastGenModel) -> None:
pass
def on_load_checkpoint_end(self, model: FastGenModel, iteration: int = 0) -> None:
pass
def on_dataloader_init_start(self, model: FastGenModel, iteration: int = 0) -> None:
pass
def on_dataloader_init_end(
self, model: FastGenModel, dataloader_train: DataLoader, dataloader_val: DataLoader, iteration: int = 0
) -> None:
pass
def on_train_begin(self, model: FastGenModel, iteration: int = 0) -> None:
pass
def on_training_step_begin(
self,
model: FastGenModel,
iteration: int = 0,
) -> None:
pass
def on_training_accum_step_begin(
self,
model: FastGenModel,
data_batch: dict[str, torch.Tensor],
iteration: int = 0,
accum_iter: int = 0,
) -> None:
pass
def on_backward_begin(
self,
model: FastGenModel,
data_batch: dict[str, torch.Tensor],
output_batch: dict[str, torch.Tensor | Callable],
loss_dict: dict[str, torch.Tensor],
iteration: int = 0,
accum_iter: int = 0,
) -> None:
pass
def on_training_step_end(
self,
model: FastGenModel,
data_batch: dict[str, torch.Tensor],
output_batch: dict[str, torch.Tensor | Callable],
loss_dict: dict[str, torch.Tensor],
iteration: int = 0,
) -> None:
pass
def on_optimizer_step_begin(self, model: FastGenModel, iteration: int = 0) -> None:
pass
def on_train_end(self, model: FastGenModel, iteration: int = 0) -> None:
pass
def on_validation_begin(self, model: FastGenModel, iteration: int = 0, idx: int = 0) -> None:
pass
def on_validation_step_begin(
self, model: FastGenModel, data_batch: dict[str, torch.Tensor], step: int = 0, iteration: int = 0, idx: int = 0
) -> None:
pass
def on_validation_step_end(
self,
model: FastGenModel,
data_batch: dict[str, torch.Tensor],
output_batch: dict[str, torch.Tensor | Callable],
loss_dict: dict[str, torch.Tensor],
step: int = 0,
iteration: int = 0,
idx: int = 0,
) -> None:
pass
def on_validation_end(self, model: FastGenModel, iteration: int = 0, idx: int = 0) -> None:
pass
def on_save_checkpoint_start(self, model: FastGenModel, iteration: int = 0) -> None:
pass
def on_save_checkpoint_success(self, model: FastGenModel, iteration: int = 0, path: str = None) -> None:
pass
def on_save_checkpoint_end(self, model: FastGenModel, iteration: int = 0) -> None:
pass
def on_app_end(self, model: FastGenModel, iteration: int = 0) -> None:
pass
def state_dict(self) -> dict[str, Any]:
return {}
def load_state_dict(self, state_dict: dict[str, Any]) -> None:
pass
|