chore: 更新项目名称

2026-03-31 09:34:11 +08:00
parent 780b9e1855
commit 2e009cf59a
51 changed files with 118 additions and 118 deletions
@@ -0,0 +1,99 @@
+import torch.nn as nn
+from torch.optim import Optimizer
+from torch.optim.lr_scheduler import LRScheduler
+from torch.utils.data import DataLoader
+
+from astrai.data import ResumableDistributedSampler
+from astrai.data.serialization import Checkpoint
+from astrai.trainer.strategy import StrategyFactory, BaseStrategy
+from astrai.config.train_config import TrainConfig
+from astrai.parallel.setup import get_current_device, get_world_size, get_rank
+
+from dataclasses import dataclass, field
+from typing import Optional, Self
+
+
+@dataclass
+class TrainContext:
+    model: nn.Module = field(default=None)
+    strategy: BaseStrategy = field(default=None)
+    dataloader: DataLoader = field(default=None)
+    optimizer: Optimizer = field(default=None)
+    scheduler: LRScheduler = field(default=None)
+    checkpoint: Checkpoint = field(default=None)
+
+    epoch: int = field(default=0)
+    iteration: int = field(default=0)
+    loss: float = field(default=0.0)
+
+    world_size: int = field(default=1)
+    rank: int = field(default=0)
+    kwargs: dict = field(default_factory=dict)
+
+
+class TrainContextBuilder:
+    def __init__(self, config: TrainConfig):
+        self.config = config
+        self._context = TrainContext(
+            model=config.model,
+            world_size=get_world_size(),
+            rank=get_rank(),
+        )
+
+        device = get_current_device()
+        self._context.model = self._context.model.to(device=device)
+
+        if self.config.nprocs > 1:
+            fn = self.config.parallel_wrapper
+            self._context.model = fn(self._context.model)
+
+        self._context.optimizer = self.config.optimizer_fn(self._context.model)
+        self._context.scheduler = self.config.scheduler_fn(self._context.optimizer)
+
+    def with_checkpoint(self, checkpoint: Optional[Checkpoint]) -> Self:
+        if checkpoint is None:
+            checkpoint = Checkpoint(
+                state_dict=self._context.model.state_dict(),
+            )
+        else:
+            # resume from the assigned checkpoint or assigned iteration
+            self._context.epoch = max(checkpoint.epoch, self.config.start_epoch)
+            self._context.iteration = max(checkpoint.iteration, self.config.start_batch)
+            self._context.model.load_state_dict(checkpoint.state_dict)
+
+        self._context.checkpoint = checkpoint
+        return self
+
+    def with_dataloader(self) -> Self:
+        # fix: change batch level iteration to sample level offset
+        config = self.config
+        sampler_offset = self._context.iteration * config.batch_size
+        resumeable_sampler = ResumableDistributedSampler(
+            data_source=config.dataset,
+            start_epoch=self._context.epoch,
+            start_iter=sampler_offset,
+            seed=config.random_seed,
+        )
+
+        dataloader = DataLoader(
+            config.dataset,
+            batch_size=config.batch_size,
+            sampler=resumeable_sampler,
+            num_workers=config.num_workers,
+            pin_memory=config.pin_memory,
+            prefetch_factor=config.prefetch_factor,
+        )
+        self._context.dataloader = dataloader
+        return self
+
+    def with_strategy(self) -> Self:
+        self._context.strategy = StrategyFactory.create(
+            model=self._context.model,
+            train_type=self.config.strategy,
+            device=get_current_device(),
+            **self.config.extra_kwargs,
+        )
+        return self
+
+    def build(self) -> TrainContext:
+        return self._context