chore: 修改类名，优化导入顺序

astrai/trainer/__init__.py

@@ -1,8 +1,8 @@
 from astrai.trainer.schedule import BaseScheduler, SchedulerFactory
 from astrai.trainer.strategy import BaseStrategy, StrategyFactory
 from astrai.trainer.train_callback import (
-    TrainCallback,
     CallbackFactory,
+    TrainCallback,
 )
 from astrai.trainer.trainer import Trainer
 

astrai/trainer/train_callback.py

@@ -8,8 +8,9 @@ import torch.nn as nn
 from torch.nn.utils import clip_grad_norm_
 from tqdm import tqdm
 
-from astrai.serialization import Checkpoint
+from astrai.factory import BaseFactory
 from astrai.parallel import only_on_rank
+from astrai.serialization import Checkpoint
 from astrai.trainer.metric_util import (
     ctx_get_grad_max,
     ctx_get_grad_mean,
@@ -21,7 +22,6 @@ from astrai.trainer.metric_util import (
     ctx_get_lr,
 )
 from astrai.trainer.train_context import TrainContext
-from astrai.factory import BaseFactory
 
 
 @runtime_checkable

astrai/trainer/train_context.py

@@ -8,8 +8,8 @@ from torch.utils.data import DataLoader
 
 from astrai.config.train_config import TrainConfig
 from astrai.dataset import ResumableDistributedSampler
-from astrai.serialization import Checkpoint
 from astrai.parallel.setup import get_current_device, get_rank, get_world_size
+from astrai.serialization import Checkpoint
 from astrai.trainer.strategy import BaseStrategy, StrategyFactory
 
 

astrai/trainer/trainer.py

@@ -2,11 +2,11 @@ import logging
 from typing import List, Optional
 
 from astrai.config import TrainConfig
-from astrai.serialization import Checkpoint
 from astrai.parallel.setup import spawn_parallel_fn
+from astrai.serialization import Checkpoint
 from astrai.trainer.train_callback import (
-    TrainCallback,
     CallbackFactory,
+    TrainCallback,
 )
 from astrai.trainer.train_context import TrainContext, TrainContextBuilder