chore: 更新项目名称
This commit is contained in:
@@ -0,0 +1,235 @@
|
||||
import os
|
||||
import json
|
||||
import time
|
||||
import torch.nn as nn
|
||||
|
||||
from pathlib import Path
|
||||
from tqdm import tqdm
|
||||
from torch.nn.utils import clip_grad_norm_
|
||||
from typing import Callable, List, Optional, Protocol
|
||||
|
||||
from astrai.parallel import only_on_rank
|
||||
from astrai.trainer.metric_util import (
|
||||
ctx_get_loss,
|
||||
ctx_get_lr,
|
||||
ctx_get_grad_max,
|
||||
ctx_get_grad_min,
|
||||
ctx_get_grad_norm,
|
||||
ctx_get_grad_mean,
|
||||
ctx_get_grad_std,
|
||||
ctx_get_grad_nan_num,
|
||||
)
|
||||
from astrai.data.serialization import Checkpoint
|
||||
from astrai.trainer.train_context import TrainContext
|
||||
|
||||
|
||||
class TrainCallback(Protocol):
|
||||
"""
|
||||
Callback interface for trainer.
|
||||
"""
|
||||
|
||||
def on_train_begin(self, context: TrainContext):
|
||||
"""Called at the beginning of training."""
|
||||
|
||||
def on_train_end(self, context: TrainContext):
|
||||
"""Called at the end of training."""
|
||||
|
||||
def on_epoch_begin(self, context: TrainContext):
|
||||
"""Called at the beginning of each epoch."""
|
||||
|
||||
def on_epoch_end(self, context: TrainContext):
|
||||
"""Called at the end of each epoch."""
|
||||
|
||||
def on_step_begin(self, context: TrainContext):
|
||||
"""Called at the beginning of each step."""
|
||||
|
||||
def on_step_end(self, context: TrainContext):
|
||||
"""Called at the end of each step."""
|
||||
|
||||
def on_batch_begin(self, context: TrainContext):
|
||||
"""Called at the beginning of each batch."""
|
||||
|
||||
def on_batch_end(self, context: TrainContext):
|
||||
"""Called at the end of each batch."""
|
||||
|
||||
def on_error(self, context: TrainContext):
|
||||
"""Called when an error occurs during training."""
|
||||
|
||||
|
||||
class GradientClippingCallback(TrainCallback):
|
||||
"""
|
||||
Gradient clipping callback for trainer.
|
||||
"""
|
||||
|
||||
def __init__(self, max_grad_norm: float):
|
||||
self.max_grad_norm = max_grad_norm
|
||||
|
||||
def on_step_begin(self, context: TrainContext):
|
||||
_ = context
|
||||
clip_grad_norm_(context.model.parameters(), self.max_grad_norm)
|
||||
|
||||
|
||||
class SchedulerCallback(TrainCallback):
|
||||
"""
|
||||
Scheduler callback for trainer.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def on_train_begin(self, context: TrainContext):
|
||||
for group in context.optimizer.param_groups:
|
||||
if "initial_lr" not in group:
|
||||
group["initial_lr"] = group["lr"]
|
||||
|
||||
def on_batch_end(self, context: TrainContext):
|
||||
if context.scheduler:
|
||||
context.scheduler.step()
|
||||
|
||||
|
||||
class CheckpointCallback(TrainCallback):
|
||||
"""
|
||||
Checkpoint callback for trainer.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
save_dir: str,
|
||||
interval: int,
|
||||
weight_only: bool = False,
|
||||
state_dict_fn: Optional[Callable[[nn.Module], dict]] = None,
|
||||
):
|
||||
self.save_dir = save_dir
|
||||
self.interval = interval
|
||||
self.weight_only = weight_only
|
||||
self.state_dict_fn = state_dict_fn
|
||||
self.last_ckpt_iter = 0
|
||||
|
||||
@only_on_rank(0)
|
||||
def _save_checkpoint(self, context: TrainContext):
|
||||
save_path = os.path.join(
|
||||
self.save_dir, f"epoch_{context.epoch}_iter_{context.iteration}"
|
||||
)
|
||||
state_dict = (
|
||||
self.state_dict_fn(context.model)
|
||||
if self.state_dict_fn
|
||||
else context.model.state_dict()
|
||||
)
|
||||
|
||||
context.checkpoint = Checkpoint(
|
||||
state_dict=state_dict, epoch=context.epoch, iteration=context.iteration
|
||||
)
|
||||
|
||||
context.checkpoint.save(save_path)
|
||||
self.last_ckpt_iter = context.iteration
|
||||
|
||||
def on_batch_end(self, context: TrainContext):
|
||||
if context.iteration - self.last_ckpt_iter >= self.interval:
|
||||
self._save_checkpoint(context)
|
||||
|
||||
def on_train_end(self, context: TrainContext):
|
||||
if context.iteration != self.last_ckpt_iter:
|
||||
self._save_checkpoint(context)
|
||||
|
||||
def on_error(self, context: TrainContext):
|
||||
self._save_checkpoint(context)
|
||||
|
||||
|
||||
class ProgressBarCallback(TrainCallback):
|
||||
"""
|
||||
Progress bar callback for trainer.
|
||||
"""
|
||||
|
||||
def __init__(self, num_epoch: int):
|
||||
self.num_epoch = num_epoch
|
||||
self.progress_bar: tqdm = None
|
||||
|
||||
@only_on_rank(0)
|
||||
def on_epoch_begin(self, context: TrainContext):
|
||||
self.progress_bar = tqdm(
|
||||
context.dataloader,
|
||||
desc=f"Epoch {context.epoch + 1}/{self.num_epoch}",
|
||||
dynamic_ncols=True,
|
||||
)
|
||||
|
||||
@only_on_rank(0)
|
||||
def on_batch_end(self, context: TrainContext):
|
||||
self.progress_bar.set_postfix(
|
||||
{
|
||||
"loss": f"{context.loss:.4f}",
|
||||
"lr": f"{context.optimizer.param_groups[-1]['lr']:.2e}",
|
||||
}
|
||||
)
|
||||
self.progress_bar.update(1)
|
||||
|
||||
@only_on_rank(0)
|
||||
def on_epoch_end(self, context: TrainContext):
|
||||
_ = context
|
||||
if self.progress_bar:
|
||||
self.progress_bar.close()
|
||||
|
||||
|
||||
class MetricLoggerCallback(TrainCallback):
|
||||
def __init__(
|
||||
self,
|
||||
log_dir: str,
|
||||
save_interval: int,
|
||||
log_interval: int = 10,
|
||||
metrics: List[str] = None,
|
||||
):
|
||||
self.last_log_iter = 0
|
||||
self.save_interval = save_interval
|
||||
self.log_interval = log_interval
|
||||
self.metrics = metrics or ["loss", "lr"]
|
||||
|
||||
self.log_dir = Path(log_dir) if log_dir else Path.cwd() / "logs"
|
||||
self.log_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
self.log_cache = []
|
||||
|
||||
self._metric_funcs = {
|
||||
"loss": ctx_get_loss,
|
||||
"lr": ctx_get_lr,
|
||||
"grad_norm": ctx_get_grad_norm,
|
||||
"grad_std": ctx_get_grad_std,
|
||||
"grad_max": ctx_get_grad_max,
|
||||
"grad_min": ctx_get_grad_min,
|
||||
"grad_mean": ctx_get_grad_mean,
|
||||
"grad_nan_num": ctx_get_grad_nan_num,
|
||||
}
|
||||
|
||||
def _get_log_data(self, context: TrainContext):
|
||||
return {
|
||||
"timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
|
||||
"epoch": context.epoch,
|
||||
"iter": context.iteration,
|
||||
**{m: self._metric_funcs[m](context) for m in self.metrics},
|
||||
}
|
||||
|
||||
@only_on_rank(0)
|
||||
def _add_log(self, log_data):
|
||||
self.log_cache.append(log_data)
|
||||
|
||||
@only_on_rank(0)
|
||||
def _save_log(self, epoch, iter):
|
||||
log_file = self.log_dir / f"epoch_{epoch}_iter_{iter}_metric.jsonl"
|
||||
|
||||
with open(log_file, "w") as f:
|
||||
for log in self.log_cache:
|
||||
f.write(json.dumps(log) + "\n")
|
||||
|
||||
def on_batch_end(self, context):
|
||||
if context.iteration % self.log_interval == 0:
|
||||
log_data = self._get_log_data(context)
|
||||
self._add_log(log_data)
|
||||
|
||||
if context.iteration - self.last_log_iter >= self.save_interval:
|
||||
self._save_log(context.epoch, context.iteration)
|
||||
self.last_log_iter = context.iteration
|
||||
|
||||
def on_train_end(self, context):
|
||||
if context.iteration != self.last_log_iter:
|
||||
self._save_log(context.epoch, context.iteration)
|
||||
|
||||
def on_error(self, context):
|
||||
self._save_log(context.epoch, context.iteration)
|
||||
Reference in New Issue
Block a user