mirror of
https://github.com/vale981/ray
synced 2025-03-06 10:31:39 -05:00
61 lines
1.6 KiB
Python
61 lines
1.6 KiB
Python
import numpy as np
|
|
import ray
|
|
import ray.rllib.agents.ppo as ppo
|
|
import onnxruntime
|
|
import os
|
|
import shutil
|
|
|
|
# Configure our PPO trainer
|
|
config = ppo.DEFAULT_CONFIG.copy()
|
|
config["num_gpus"] = 0
|
|
config["num_workers"] = 1
|
|
config["framework"] = "tf"
|
|
|
|
outdir = "export_tf"
|
|
if os.path.exists(outdir):
|
|
shutil.rmtree(outdir)
|
|
|
|
np.random.seed(1234)
|
|
|
|
# We will run inference with this test batch
|
|
test_data = {
|
|
"obs": np.random.uniform(0, 1., size=(10, 4)).astype(np.float32),
|
|
}
|
|
|
|
# Start Ray and initialize a PPO trainer
|
|
ray.init()
|
|
trainer = ppo.PPOTrainer(config=config, env="CartPole-v0")
|
|
|
|
# You could train the model here
|
|
# trainer.train()
|
|
|
|
# Let's run inference on the tensorflow model
|
|
policy = trainer.get_policy()
|
|
result_tf, _ = policy.model(test_data)
|
|
|
|
# Evaluate tensor to fetch numpy array
|
|
with policy._sess.as_default():
|
|
result_tf = result_tf.eval()
|
|
|
|
# This line will export the model to ONNX
|
|
res = trainer.export_policy_model(outdir, onnx=11)
|
|
|
|
# Import ONNX model
|
|
exported_model_file = os.path.join(outdir, "saved_model.onnx")
|
|
|
|
# Start an inference session for the ONNX model
|
|
session = onnxruntime.InferenceSession(exported_model_file, None)
|
|
|
|
# Pass the same test batch to the ONNX model (rename to match tensor names)
|
|
onnx_test_data = {f"default_policy/{k}:0": v for k, v in test_data.items()}
|
|
|
|
result_onnx = session.run(["default_policy/model/fc_out/BiasAdd:0"],
|
|
onnx_test_data)
|
|
|
|
# These results should be equal!
|
|
print("TENSORFLOW", result_tf)
|
|
print("ONNX", result_onnx)
|
|
|
|
assert np.allclose(result_tf, result_onnx), \
|
|
"Model outputs are NOT equal. FAILED"
|
|
print("Model outputs are equal. PASSED")
|