Skip to main content
`Ctrl`+`K`

# Using ONNX Model Inference and Saving Input Data in NPY Format

This topic explains how to perform inference with an ONNX model using floating-point inputs and save the input data in .npy format. This approach facilitates data storage and reuse and can serve as a calibration dataset during model quantization, provided that the data adequately reflects the typical distribution of the model inputs.

Through an example, we demonstrate how to define a simple dataset class (InputDataset), perform inference using an ONNX model, and save input data in .npy format to support subsequent model quantization.

## Detailed Code

    import onnxruntime as ort
    import numpy as np
    import os
    from torch.utils.data import Dataset, DataLoader

    # A simple dataset with two inputs (`input1`, `input2`) and random tensors.
    # Users can customize data generation to match their model's needs.
    class InputDataset(Dataset):
        def __init__(self, num_samples):
            super(InputDataset, self).__init__()
            self.num_samples = num_samples
            self.input1 = [np.random.rand(3, 224, 224).astype(np.float32) for _ in range(num_samples)]
            self.input2 = [np.random.rand(10).astype(np.float32) for _ in range(num_samples)]
            self.labels = [np.random.randint(0, 2) for _ in range(num_samples)]

        def __len__(self):
            return self.num_samples

        def __getitem__(self, idx):
            return {
                "input1": self.input1[idx],
                "input2": self.input2[idx],
                "label": self.labels[idx]
            }

    dataset = InputDataset(num_samples=10)
    data_loader = DataLoader(dataset, batch_size=1, shuffle=False, num_workers=4, pin_memory=True)

    onnx_model_path = "path/to/your/float_model.onnx"
    session = ort.InferenceSession(onnx_model_path)
    input_names = [inp.name for inp in session.get_inputs()]
    output_names = [out.name for out in session.get_outputs()]

    enable_data_caching = True
    calibration_cache_dir = "calibration_data/"

    if enable_data_caching:
        for name in input_names:
            input_folder_path = os.path.join(calibration_cache_dir, name)
            os.makedirs(input_folder_path, exist_ok=True)

    for batch_idx, batch in enumerate(data_loader):
        input_feed = {}

        for name in input_names:
            input_data = batch[name].numpy()
            input_feed[name] = input_data

            # If `enable_data_caching` is True, save input data as `.npy` files by input name for each batch.
            if enable_data_caching:
                file_path = os.path.join(calibration_cache_dir, name, f"calib_{batch_idx+1:06d}.npy")
                np.save(file_path, input_data)
                print(f"Saved input data for {name} to {file_path}")

        outputs = session.run(output_names, input_feed)

        predictions = np.argmax(outputs[0], axis=1)

        print(f"Predictions for batch {batch_idx}: {predictions}")

The input data saved during ONNX inference can serve as a calibration dataset for model quantization. For instructions on how to use the saved NPY data as a calibration dataset, refer to Calibration Data Path for AMD Quark Quantizer. The output data format saved during inference is as follows:

### For Single-Input Models

    calibration_data/
      calib_000001.npy
      calib_000002.npy
      calib_000003.npy
      calib_000004.npy
      calib_000005.npy
      ...

### For Multi-Input Models

    calibration_data/
      input1_name/
        calib_000001.npy
        calib_000002.npy
        calib_000003.npy
        calib_000004.npy
        calib_000005.npy
        ...
      input2_name/
        calib_000001.npy
        calib_000002.npy
        calib_000003.npy
        calib_000004.npy
        calib_000005.npy
        ...