code: def save_onnx_graph(
      input_nodes, output_nodes, checkpoint_save_dir: str
  ) -> None:
      """
      Given the input nodes and output nodes of the TF graph, save it as an onnx graph
      This requires the TF graph and the weights checkpoint to be stored in the experiment directory.
      It then freezes the graph (merging the graph and weights checkpoint), and converts it to ONNX.
  
      :param input_nodes: A list of input nodes for the TF graph
      :param output_nodes: A list of output nodes for the TF graph
      :param checkpoint_save_dir: The directory to save the ONNX graph to
      :return: None
      """
      import tf2onnx  # just to verify that tf2onnx is installed
  
      # freeze graph
      frozen_graph_path = os.path.join(
          checkpoint_save_dir, "frozen_graph.pb"
      )
      freeze_graph_command = [
          "python -m tensorflow.python.tools.freeze_graph",
          "--input_graph={}".format(
              os.path.join(checkpoint_save_dir, "graphdef.pb")
          ),
          "--input_binary=true",
          "--output_node_names='{}'".format(
              ",".join([o.split(":")[0] for o in output_nodes])
          ),
          "--input_checkpoint={}".format(
              tf.train.latest_checkpoint(checkpoint_save_dir)
          ),
          "--output_graph={}".format(frozen_graph_path),
      ]
      start_shell_command_and_wait(" ".join(freeze_graph_command))
  
      # convert graph to onnx
      onnx_graph_path = os.path.join(checkpoint_save_dir, "model.onnx")
      convert_to_onnx_command = [
          "python -m tf2onnx.convert",
          "--input {}".format(frozen_graph_path),
          "--inputs '{}'".format(",".join(input_nodes)),
          "--outputs '{}'".format(",".join(output_nodes)),
          "--output {}".format(onnx_graph_path),
          "--verbose",
      ]
      start_shell_command_and_wait(" ".join(convert_to_onnx_command))
  
imports: import os
  import time
  from typing import Any, List, Tuple, Dict
  import numpy
  import tensorflow
  from rl_coach.architectures.architecture import Architecture
  from rl_coach.architectures.tensorflow_components.savers import GlobalVariableSaver
  from rl_coach.base_parameters import AgentParameters, DistributedTaskParameters
  from rl_coach.core_types import GradientClippingMethod
  from rl_coach.saver import SaverCollection
  from rl_coach.spaces import SpacesDefinition
  from rl_coach.utils import force_list, squeeze_list, start_shell_command_and_wait
label_desc: Use image node_id 3 for calling a global function with example usage: save_onnx_graph(input_nodes, output_nodes, checkpoint_save_dir) without return types