qing
63cfdd2066
- 第 1 轮安全审计修复: F-1~F-10（路径校验、预训练架构修正等） - 第 2 轮复审修复 11 项: 3 项上轮修复不完整（R01 checkpoint 键校验/ R02 fused_hidden_dim 值域/R03 pretrained=False）+ 8 项新问题 （actions.npy 形状校验、帧数警告、CLI 值域校验、stats 形状与路径、 onnxscript 依赖、归一化语义、解析去重） - 单元测试由 17 扩展至 51 用例（50 通过 + 1 环境性跳过） - 审计文档两轮合并至 actpred_CODE_AUDIT_20260808.md，更新 INDEX
191 lines
8.0 KiB
Python
191 lines
8.0 KiB
Python
import argparse
|
||
from dataclasses import dataclass, asdict, field
|
||
from typing import Optional
|
||
|
||
|
||
@dataclass
|
||
class DataConfig:
|
||
image_size: int = 64
|
||
action_dim: int = 18
|
||
history_len: int = 5
|
||
data_dir: str = "./data"
|
||
normalize_actions: bool = False
|
||
action_stats_path: Optional[str] = None
|
||
|
||
|
||
@dataclass
|
||
class ModelConfig:
|
||
cnn_backbone: str = "simple"
|
||
cnn_feature_dim: int = 128
|
||
action_hidden_dim: int = 64
|
||
fused_hidden_dim: int = 128
|
||
dropout: float = 0.1
|
||
history_encoder: str = "mlp"
|
||
|
||
|
||
@dataclass
|
||
class TrainConfig:
|
||
batch_size: int = 32
|
||
epochs: int = 100
|
||
lr: float = 1e-3
|
||
weight_decay: float = 1e-4
|
||
val_split: float = 0.15
|
||
checkpoint_dir: str = "./checkpoints"
|
||
patience: int = 15
|
||
num_workers: int = 2
|
||
seed: int = 42
|
||
max_grad_norm: float = 1.0
|
||
|
||
|
||
@dataclass
|
||
class Config:
|
||
data: DataConfig = field(default_factory=DataConfig)
|
||
model: ModelConfig = field(default_factory=ModelConfig)
|
||
train: TrainConfig = field(default_factory=TrainConfig)
|
||
|
||
def to_dict(self):
|
||
return asdict(self)
|
||
|
||
@classmethod
|
||
def from_dict(cls, d):
|
||
return cls(
|
||
data=DataConfig(**d.get("data", {})),
|
||
model=ModelConfig(**d.get("model", {})),
|
||
train=TrainConfig(**d.get("train", {})),
|
||
)
|
||
|
||
|
||
def build_argparser(description: str = "") -> argparse.ArgumentParser:
|
||
parser = argparse.ArgumentParser(description=description)
|
||
parser.add_argument("--image-size", type=int, default=None)
|
||
parser.add_argument("--action-dim", type=int, default=None)
|
||
parser.add_argument("--history-len", type=int, default=None)
|
||
parser.add_argument("--data-dir", type=str, default=None)
|
||
parser.add_argument("--normalize-actions", action=argparse.BooleanOptionalAction, default=None)
|
||
parser.add_argument("--action-stats-path", type=str, default=None)
|
||
parser.add_argument("--backbone", type=str, default=None, choices=["simple", "mobilenet"])
|
||
parser.add_argument("--cnn-feature-dim", type=int, default=None)
|
||
parser.add_argument("--action-hidden-dim", type=int, default=None)
|
||
parser.add_argument("--fused-hidden-dim", type=int, default=None)
|
||
parser.add_argument("--dropout", type=float, default=None)
|
||
parser.add_argument("--history-encoder", type=str, default=None, choices=["mlp", "gru"])
|
||
parser.add_argument("--batch-size", type=int, default=None)
|
||
parser.add_argument("--epochs", type=int, default=None)
|
||
parser.add_argument("--lr", type=float, default=None)
|
||
parser.add_argument("--weight-decay", type=float, default=None)
|
||
parser.add_argument("--val-split", type=float, default=None)
|
||
parser.add_argument("--checkpoint-dir", type=str, default=None)
|
||
parser.add_argument("--patience", type=int, default=None)
|
||
parser.add_argument("--num-workers", type=int, default=None)
|
||
parser.add_argument("--seed", type=int, default=None)
|
||
parser.add_argument("--max-grad-norm", type=float, default=None)
|
||
parser.add_argument("--init-from", type=str, default=None,
|
||
help="预训练 checkpoint 路径，作为微调起点（strict=False 加载兼容权重）")
|
||
parser.add_argument("--save-config", type=str, default=None, help="保存配置到 JSON 文件")
|
||
return parser
|
||
|
||
|
||
def validate_config_values(data_cfg: DataConfig, model_cfg: ModelConfig,
|
||
train_cfg: TrainConfig = None):
|
||
"""校验从 checkpoint 恢复或 CLI 传入的配置值域合理性
|
||
|
||
防止恶意/损坏 checkpoint 注入异常值（如超大维度导致 OOM），
|
||
也防止训练 CLI 越界参数（如 --history-len 0 / --lr -1）静默进入训练。
|
||
"""
|
||
checks = [
|
||
("image_size", data_cfg.image_size, 1, 1024, int),
|
||
("history_len", data_cfg.history_len, 1, 1024, int),
|
||
("action_dim", data_cfg.action_dim, 1, 1024, int),
|
||
("cnn_feature_dim", model_cfg.cnn_feature_dim, 1, 4096, int),
|
||
("action_hidden_dim", model_cfg.action_hidden_dim, 1, 4096, int),
|
||
("fused_hidden_dim", model_cfg.fused_hidden_dim, 1, 4096, int),
|
||
("dropout", model_cfg.dropout, 0.0, 1.0, float),
|
||
]
|
||
for name, val, lo, hi, typ in checks:
|
||
if typ is int:
|
||
ok = isinstance(val, int) and not isinstance(val, bool) and lo <= val <= hi
|
||
else:
|
||
ok = isinstance(val, (int, float)) and not isinstance(val, bool) and lo <= val <= hi
|
||
if not ok:
|
||
raise ValueError(
|
||
f"checkpoint config {name}={val} 超出合理范围 [{lo}, {hi}]")
|
||
|
||
if train_cfg is not None:
|
||
t_checks = [
|
||
("batch_size", train_cfg.batch_size, 1, 65536, int),
|
||
("epochs", train_cfg.epochs, 1, 100000, int),
|
||
("lr", train_cfg.lr, 1e-9, 10.0, float),
|
||
("weight_decay", train_cfg.weight_decay, 0.0, 1.0, float),
|
||
("val_split", train_cfg.val_split, 0.0, 1.0, float),
|
||
("patience", train_cfg.patience, 0, 100000, int),
|
||
("num_workers", train_cfg.num_workers, 0, 256, int),
|
||
("max_grad_norm", train_cfg.max_grad_norm, 0.0, 1e6, float),
|
||
]
|
||
for name, val, lo, hi, typ in t_checks:
|
||
if typ is int:
|
||
ok = isinstance(val, int) and not isinstance(val, bool) and lo <= val <= hi
|
||
else:
|
||
ok = isinstance(val, (int, float)) and not isinstance(val, bool) and lo <= val <= hi
|
||
if not ok:
|
||
raise ValueError(
|
||
f"配置 {name}={val} 超出合理范围 [{lo}, {hi}]")
|
||
|
||
|
||
def parse_checkpoint_config(cfg_dict: dict) -> Config:
|
||
"""从 checkpoint 恢复 Config，校验键集合与值域
|
||
|
||
- 顶层缺 data/model 键 → ValueError（带明确信息）
|
||
- data/model 段含未知键 → ValueError（列出未知键，避免裸 TypeError）
|
||
- 值域非法 → ValueError（复用 validate_config_values）
|
||
"""
|
||
if not isinstance(cfg_dict, dict):
|
||
raise ValueError(f"checkpoint config 应为 dict，实际为 {type(cfg_dict).__name__}")
|
||
for section in ("data", "model"):
|
||
if section not in cfg_dict:
|
||
raise ValueError(f"checkpoint config 缺少 '{section}' 段")
|
||
section_dict = cfg_dict[section]
|
||
if not isinstance(section_dict, dict):
|
||
raise ValueError(f"checkpoint config['{section}'] 应为 dict，实际为 {type(section_dict).__name__}")
|
||
known = set(DataConfig.__dataclass_fields__) if section == "data" \
|
||
else set(ModelConfig.__dataclass_fields__)
|
||
unknown = set(section_dict) - known
|
||
if unknown:
|
||
raise ValueError(
|
||
f"checkpoint config['{section}'] 含未知键: {sorted(unknown)}")
|
||
cfg = Config.from_dict(cfg_dict)
|
||
validate_config_values(cfg.data, cfg.model, cfg.train)
|
||
return cfg
|
||
|
||
|
||
def apply_overrides(cfg: Config, args: argparse.Namespace) -> Config:
|
||
mapping = {
|
||
"image_size": ("data", "image_size"),
|
||
"action_dim": ("data", "action_dim"),
|
||
"history_len": ("data", "history_len"),
|
||
"data_dir": ("data", "data_dir"),
|
||
"normalize_actions": ("data", "normalize_actions"),
|
||
"action_stats_path": ("data", "action_stats_path"),
|
||
"backbone": ("model", "cnn_backbone"),
|
||
"cnn_feature_dim": ("model", "cnn_feature_dim"),
|
||
"action_hidden_dim": ("model", "action_hidden_dim"),
|
||
"fused_hidden_dim": ("model", "fused_hidden_dim"),
|
||
"dropout": ("model", "dropout"),
|
||
"history_encoder": ("model", "history_encoder"),
|
||
"batch_size": ("train", "batch_size"),
|
||
"epochs": ("train", "epochs"),
|
||
"lr": ("train", "lr"),
|
||
"weight_decay": ("train", "weight_decay"),
|
||
"val_split": ("train", "val_split"),
|
||
"checkpoint_dir": ("train", "checkpoint_dir"),
|
||
"patience": ("train", "patience"),
|
||
"num_workers": ("train", "num_workers"),
|
||
"seed": ("train", "seed"),
|
||
"max_grad_norm": ("train", "max_grad_norm"),
|
||
}
|
||
args_dict = vars(args)
|
||
for arg_name, (section, field) in mapping.items():
|
||
val = args_dict.get(arg_name)
|
||
if val is not None:
|
||
setattr(getattr(cfg, section), field, val)
|
||
return cfg
|