qing
73b72d6889
- train/val 归一化统计量改为仅用训练集计算并落盘复用，修复 val loss 不可信 - predict.py 统计文件路径改用 os.path 处理，修复跨平台与无分隔符路径问题 - dataset.py jpg/png 混合时统一按文件名排序，保证时间顺序 - 新增 --seed 全局随机种子（含 DataLoader generator），训练可复现 - 新增 --max-grad-norm 梯度裁剪（默认 1.0），防 GRU 梯度爆炸 - get_episode_count 与数据集目录过滤条件对齐，避免无效目录导致索引错位 - 空验证集返回 inf 并保底保存最终权重，避免伪造 best 模型 - 清理 export_onnx.py 重复分支死代码
42 lines
1.3 KiB
Python
42 lines
1.3 KiB
Python
import argparse
|
|
import torch
|
|
from config import DataConfig, ModelConfig
|
|
from model import ActionPredictor
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Export actpred model to ONNX")
|
|
parser.add_argument("--model", required=True, help="checkpoint path (best_inference.pt)")
|
|
parser.add_argument("--output", default="model.onnx", help="output ONNX file path")
|
|
args = parser.parse_args()
|
|
|
|
checkpoint = torch.load(args.model, map_location="cpu", weights_only=True)
|
|
cfg_dict = checkpoint["config"]
|
|
data_cfg = DataConfig(**cfg_dict["data"])
|
|
model_cfg = ModelConfig(**cfg_dict["model"])
|
|
|
|
model = ActionPredictor(model_cfg, data_cfg)
|
|
model.load_state_dict(checkpoint["model_state_dict"])
|
|
model.eval()
|
|
|
|
dummy_image = torch.randn(1, 3, data_cfg.image_size, data_cfg.image_size)
|
|
dummy_history = torch.randn(1, data_cfg.history_len, data_cfg.action_dim)
|
|
|
|
torch.onnx.export(
|
|
model, (dummy_image, dummy_history),
|
|
args.output,
|
|
input_names=["image", "action_history"],
|
|
output_names=["predicted_action"],
|
|
dynamic_axes={
|
|
"image": {0: "batch"},
|
|
"action_history": {0: "batch"},
|
|
"predicted_action": {0: "batch"},
|
|
},
|
|
opset_version=17,
|
|
)
|
|
print(f"ONNX model exported to: {args.output}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|