code: def save_model_with_external_data(
      basepath: str,
      model_location: str,
      initializer_location: str,
      torch_load_paths: Tuple[Union[str, io.BytesIO], ...],
      onnx_model: onnx.ModelProto,  # type: ignore[name-defined]
      rename_initializer: bool = False,
  ) -> None:
      """Load PyTorch tensors from files and add to "onnx_model" as external initializers.
  
      Output files:
          ONNX model file path:
          ONNX initializer folder: os.path.join(basepath, initializer_location)
  
      After running this function, you can do
          ort_sess = onnxruntime.InferenceSession(os.path.join(basepath, model_location))
      to execute the model.
  
      Arguments:
          basepath: Base path of the external data file (e.g., "/tmp/large-onnx-model").
          model_location: Relative location of the ONNX model file.
              E.g., "model.onnx" so that the model file is saved to
              "<model_location>/model.onnx".
          initializer_location: Relative location of the ONNX initializer folder.
              E.g., "initializers" so that the initializers are saved to
              "<model_location>/initializers".
              Note: When initializers are >2GB, must be the same as `model_location`.
          torch_load_paths: Files which containing serialized PyTorch tensors to be saved
              as ONNX initializers. They are loaded by torch.load.
          onnx_model: ONNX model to be saved with external initializers.
              If an input name matches a tensor loaded from "torch_load_paths",
              the tensor will be saved as that input's external initializer.
          rename_initializer: Replaces "." by "_" for all ONNX initializer names.
              Not needed by the official torch.onnx.dynamo_export. This is a hack
              for supporting `FXSymbolicTracer` tracer with fake tensor mode.
              In short, `FXSymbolicTracer` lifts FX parameters (self.linear_weight)
              as inputs (`def forward(self, linear_weight)`) and therefore, `.` cannot be used.
      """
      # FIXME: Avoid importing onnx into torch.onnx.
      import onnx
  
      onnx_model_with_initializers = onnx.ModelProto()  # type: ignore[attr-defined]
      onnx_model_with_initializers.CopyFrom(onnx_model)
      onnx_input_names = {
          input.name for input in onnx_model.graph.input
      }
  
      for path in torch_load_paths:
          state_dict = torch.load(path)
          for name, tensor in state_dict.items():
              if rename_initializer:
                  # Basically, "transformer.attention.self.query.weight" is mapped
                  # to "transformer_attention_self_query_weight" for mimicking the
                  # name-modifying code in FX-to-ONNX exporter.
                  # See function _replace_get_attr_with_placeholder for details.
                  name = name.replace(".", "_")
  
              # This block tries to match the onnx initializer name with torch parameter/buffer
              #  e.g. A pytorch buffer 'transformer.h.0.attn.bias' can be named 'h.0.attn.bias' in a ONNX initializer
              # For each PyTorch tensor name loaded by torch.load,
              #  1.  Search its best match in ONNX model. E.g., the match of
              #       "transformer_attention_weight" could be "attention_weight".
              #  2.  Set "tensor" as the initializer of the matched ONNX input.
              #      E.g., "tensor" is stored as the initializer of "attention_weight".
              # Step 1 is required because sometimes, tensor names are stored with prefix the dictionary
              # loaded by torch.load.
              if name in onnx_input_names:
                  # Same input name shouldn't be matched again
                  onnx_input_names.remove(name)
              else:
                  for onnx_input_name in onnx_input_names:
                      if onnx_input_name.endswith(
                          name
                      ) or name.endswith(onnx_input_name):
                          # Find a match. Change name to the matched ONNX input name, so that we
                          # create initializer with the right ONNX name.
                          name = onnx_input_name
                          onnx_input_names.remove(onnx_input_name)
                          break
  
              relative_tensor_file_path = os.path.join(
                  initializer_location, name
              )
              # Create one file per tensor.
              # tensor_proto.raw_data is stored to external file at
              # os.path.join(basepath, relative_tensor_file_path).
              tensor_proto = _create_tensor_proto_with_external_data(
                  tensor, name, relative_tensor_file_path, basepath
              )
              # Add the tensor_proto to the ONNX model as an initializer with external data.
              onnx_model_with_initializers.graph.initializer.append(
                  tensor_proto
              )
  
      # model_location should be a pure file name such as "file_name.onnx", not "folder/file_name.onnx".
      onnx.save(
          onnx_model_with_initializers,
          os.path.join(basepath, model_location),
      )
  
file_path: pytorch/torch/onnx/_internal/fx/serialization.py
imports: from __future__ import annotations
  import io
  import os
  from typing import Tuple, TYPE_CHECKING, Union
  import torch
  from torch.onnx import _type_utils
  from torch.onnx._internal import _beartype
label_desc: Use image node_id 2 for calling a global function with example usage: save_model_with_external_data(basepath, model_location, initializer_location, torch_load_paths, onnx_model, rename_initializer) without return types