Spaces:
Sleeping
Sleeping
File size: 7,452 Bytes
07cb7d3 | 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 | """
有监督 Keras 模型任务的固定流水线。
这个文件放分类、分割这类“数据源 + 模型构建器 + 训练规则”任务共用的主流程。
任务入口只负责提供数据源、模型构建器和训练规则，具体训练、导出和样例测试由这里统一调度。
"""
from dataclasses import dataclass
from pathlib import Path
from deep_learning.data.spec import SupervisedDataSource
from deep_learning.env.keras import enable_mixed_precision
from deep_learning.env.logger import log
from deep_learning.env.resolve import resolve_saved
from deep_learning.pipeline.specs import Pipeline
from deep_learning.pipeline.specs.configs import CheckpointLoadRules, TrainingRule
from deep_learning.pipeline.services.logging_config import log_config
from deep_learning.models.spec import SupervisedModelBuilder
from deep_learning.pipeline.context import InferenceArtifactState, TrainingArtifactState
from deep_learning.pipeline.env.const import ENV
from deep_learning.pipeline.services import CheckpointService, build_runtime
from deep_learning.pipeline.services.training_callbacks import build_common_callbacks
from deep_learning.pipeline.stages.export import SaveKerasModelStage
@dataclass
class SupervisedModelPipeline(Pipeline):
"""
有监督 Keras 任务主入口。
这个类把训练、导出和样例测试固定成显式动作，分类和分割任务只需要替换数据源、
模型构建器。
"""
name: str
data_source: SupervisedDataSource
model_builder: SupervisedModelBuilder
training_rule: TrainingRule
checkpoint_load_rules: CheckpointLoadRules = None
task_dir: Path = None
def __post_init__(self):
runtime = build_runtime(self.name, self.task_dir)
self.task_dir = runtime.task_dir
self._runtime_env = runtime
if self.checkpoint_load_rules is None:
super().__setattr__("checkpoint_load_rules", CheckpointLoadRules())
def __setattr__(self, name, value):
if "_runtime_env" in self.__dict__ and not name.startswith("_"):
public_fields = {"data_source", "model_builder"}
if name not in public_fields:
raise AttributeError(f"cannot assign to field '{name}'")
super().__setattr__(name, value)
def train(self):
with log():
enable_mixed_precision()
with log():
self.log_config()
with log():
self._runtime_env.checkpoint_dir.mkdir(parents=True, exist_ok=True)
with log():
training_state = self._load_or_build_training_state(
checkpoint_rule=self.checkpoint_load_rules.resolve_train_rule(
default_dirs=[self._runtime_env.checkpoint_dir]
),
checkpoint_must=False
)
with log("构建训练计划"):
self.model_builder.compile_training_model(training_state.training_artifact.model)
train_ds, validation_ds = self.data_source.training_ds()
callbacks = build_common_callbacks(
runtime=self._runtime_env,
checkpoint_filename="model_epoch_{epoch:03d}.weights.h5",
save_weights_only=True
)
if hasattr(training_state.training_artifact.model, "summary"):
training_state.training_artifact.model.summary()
with log("开始训练", "训练结束"):
return training_state.training_artifact.model.fit(
train_ds,
validation_data=validation_ds,
epochs=self.training_rule.epochs,
steps_per_epoch=self.training_rule.steps_per_epoch,
callbacks=callbacks
).history
def export_model(self) -> Path:
self.log_config()
self._runtime_env.checkpoint_dir.mkdir(parents=True, exist_ok=True)
training_state = self._load_or_build_training_state(
checkpoint_rule=self.checkpoint_load_rules.resolve_export_rule(
default_dirs=[self._runtime_env.checkpoint_dir]
),
checkpoint_must=True
)
inference_artifact = self.model_builder.convert_to_inference_artifact(
training_artifact=training_state.training_artifact
)
inference_state = InferenceArtifactState(
inference_artifact=inference_artifact,
inference_resource=None
)
export_stage = SaveKerasModelStage(resolve_saved(f"models/{self.name}"))
return export_stage.run(
self._runtime_env,
inference_state,
training_state.checkpoint_epoch
)
def test_examples(self, checkpoint: str = "test") -> None:
self.log_config()
if checkpoint == "test":
checkpoint_rule = self.checkpoint_load_rules.resolve_test_rule(
default_dirs=[resolve_saved(f"models/{self.name}")]
)
if checkpoint_rule["suffix"] is None:
checkpoint_rule["suffix"] = ".keras"
checkpoint_service = CheckpointService()
checkpoint_path, _ = checkpoint_service.resolve_existing_checkpoint(
checkpoint_rule=checkpoint_rule,
not_found_message="未找到任何导出模型文件"
)
artifact = self.model_builder.load_inference_artifact(checkpoint_path)
checkpoint_service.log_loaded_checkpoint(checkpoint_path)
self.data_source.test_examples(artifact.model)
return
training_state = self._load_or_build_training_state(
checkpoint_rule=self.checkpoint_load_rules.resolve_train_rule(
default_dirs=[self._runtime_env.checkpoint_dir]
),
checkpoint_must=True
)
inference_artifact = self.model_builder.convert_to_inference_artifact(
training_state.training_artifact
)
self.data_source.test_examples(
inference_artifact.model
)
def log_config(self):
self._runtime_env.log_dir.mkdir(parents=True, exist_ok=True)
config_path = self._runtime_env.log_dir / "config.txt"
output = log_config(self, config_path, header=f"ENV[{ENV}]")
print(output)
print(f"配置已保存到: {config_path}")
def _load_or_build_training_state(
self,
checkpoint_rule: dict,
checkpoint_must: bool
) -> TrainingArtifactState:
checkpoint_service = CheckpointService()
checkpoint_path, checkpoint_epoch = checkpoint_service.resolve_checkpoint(**checkpoint_rule)
if checkpoint_path is not None:
checkpoint_service.log_loading_checkpoint(checkpoint_path, checkpoint_epoch)
training_artifact = self.model_builder.build_training_artifact()
training_artifact.model.load_weights(str(checkpoint_path))
checkpoint_service.log_loaded_checkpoint(checkpoint_path)
return TrainingArtifactState(
training_artifact=training_artifact,
checkpoint_epoch=checkpoint_epoch
)
if checkpoint_must:
raise ValueError(f"目录 {self._runtime_env.checkpoint_dir} 中未找到检查点文件")
print("未找到检查点，使用新模型")
return TrainingArtifactState(
training_artifact=self.model_builder.build_training_artifact()
)
|