[onnxruntime.training.artifacts] convert the'. h5 'file to onnx for training

[laborchen]

Use 'generate_artifacts' to convert the'. h5 'file to onnx for training
But the training effect in Qt is completely inferior to using Python's'. h5 '
I suspect it's a code conversion issue with 'onnxrruntime. training. artifacts' that caused the fusion of operators in training_madel.onnx, which can only predict but cannot be trained

import os
import onnx
import tf2onnx
import numpy as np
import tensorflow as tf
from tensorflow.keras.models import load_model, save_model
from tensorflow.keras.layers import Activation
from tensorflow.keras.models import clone_model
from tensorflow.keras.activations import relu
import onnxruntime.training.artifacts as artifacts

def export_initial_models(keras_model_path, output_dir):
"""
Args:
keras_model_path: Keras模型文件路径 (.h5)
output_dir: 输出目录
"""
# 创建输出目录
baseline_dir = os.path.join(output_dir, 'baseline')
current_dir = os.path.join(output_dir, 'current')
os.makedirs(baseline_dir, exist_ok=True)
os.makedirs(current_dir, exist_ok=True)
# 加载Keras模型
print("Loading Keras model...")
keras_model = load_model(keras_model_path)
# ==========================================================
# 导出 ONNX + ORT Training Artifacts
# ==========================================================
try:
print("开始导出 ONNX...")
# ONNX输出路径
onnx_path = os.path.join(baseline_dir, 'inference_model.onnx')
spec = (
tf.TensorSpec(
(1, 1, 5, 1000),
tf.float32,
name="input"
),
)
model_proto, _ = tf2onnx.convert.from_keras(
keras_model,
input_signature=spec,
opset=17
)
with open(onnx_path, "wb") as f:
f.write(model_proto.SerializeToString())
print("✓ inference_model.onnx 导出成功")
# ======================================================
# 生成 ORT Training Artifacts
# ======================================================
print("开始生成 ORT Training Artifacts...")
onnx_model = onnx.load(onnx_path)

    # 自动打印参数名
    print("=== 模型参数 ===")
    trainable_params = []
    for init in onnx_model.graph.initializer:
        name = init.name
        print(name)
        # 只训练真正权重
        if any(k in name.lower() for k in [
            "weight",
            "bias",
            "kernel",
            "gamma",
            "beta",
            "matmul",
            "conv"
        ]):
            trainable_params.append(name)
    print("\n=== 可训练参数 ===")
    for p in trainable_params:
        print(p)

    # 生成训练图
    artifacts.generate_artifacts(
        onnx_model,
        requires_grad=trainable_params,
        frozen_params=[],
        loss=artifacts.LossType.CrossEntropyLoss,
        optimizer=artifacts.OptimType.AdamW,
        artifact_directory=baseline_dir
    )
    print("✓ ORT Training Artifacts 生成成功")
except Exception as e:
    print("ORT Artifacts 导出失败:")
    print(e)

# 5. 复制到current目录
import shutil
for item in os.listdir(baseline_dir):
    src = os.path.join(baseline_dir, item)
    dst = os.path.join(current_dir, item)
    if os.path.isfile(src):
        shutil.copy2(src, dst)
    elif os.path.isdir(src):
        if os.path.exists(dst):
            shutil.rmtree(dst)
        shutil.copytree(src, dst)
print(f"\n✅ Initial models exported to {output_dir}")
print("\nGenerated files:")
for item in os.listdir(baseline_dir):
    item_path = os.path.join(baseline_dir, item)
    if os.path.isfile(item_path):
        size = os.path.getsize(item_path) / 1024
        print(f"  - {item} ({size:.2f} KB)")
    else:
        print(f"  - {item}/ (directory)")

if name == "main":
keras_model_path = "./Model/CrossSubject/Net.h5"
export_initial_models(keras_model_path, "./models")