Helper functions for large ONNX models support

[dbogunowicz]

Implementing onnx helper functions that will be shared across our repositories (deepsparse, sparseml, etc.) to reliably save and validate not only the models which fit into a protobuf (as we did up until now) but also the larger ones (where the large data information is stored in the additional accompanying file).

Did not include the helper function for loading onnx models, since this is not needed if we work under the assumption that the onnx file and the auxiliary data file are in the same directory, see: https://github.com/onnx/onnx/blob/main/docs/ExternalData.md

Related PRs:
neuralmagic/deepsparse#995 (deepsparse)
neuralmagic/sparseml#1540 (sparseml)

Manual testing (only the big model pathway):

import shutil
from sparsezoo.utils.onnx import save_onnx, validate_onnx
import os
import onnx
import glob

### Validate the new functions for >2Gb models ###

## Prepare the model and test directory ##
model_path = "/network/damian/big_models/codegen-2B-multi/decoder_model.onnx"
model_onnx = onnx.load(model_path)
test_dir = "test_dir"

## Test validate_onnx ##
# should pass for a big model, input is path
validate_onnx(model = model_path)
print("Test 1 passed")

# should raise warning for a big model, input is model
validate_onnx(model = model_onnx)
print("Test 2 passed")

## Test save_onnx ##
# should produce two files
os.mkdir(test_dir)
save_onnx(model_onnx, os.path.join(test_dir, "test.onnx"))
files = glob.glob(os.path.join(test_dir, "*"))
print(f"Produced files {files}")
shutil.rmtree(test_dir)
print("Test 3 passed")

#should produce two files as well
os.mkdir(test_dir)
save_onnx(model_onnx, os.path.join(test_dir, "test.onnx"), external_data_file="something")
files = glob.glob(os.path.join(test_dir, "*"))
print(f"Produced files {files}")
shutil.rmtree(test_dir)
print("Test 4 passed")

Out:

Test 1 passed
Test 2 passed
Attempting to validate an in-memory ONNX model with size > 2000000000 bytes.`validate_onnx` skipped, as large ONNX models cannot be validated in-memory. To validate this model, save it to disk and call `validate_onnx` on the file path.
Produced files ['test_dir/test.onnx', 'test_dir/963eac78-e441-11ed-8c90-d52fb1102de6']
Test 3 passed
Produced files ['test_dir/test.onnx', 'test_dir/something']
Test 4 passed

[KSGulin]

Looks good! Left a couple comments

[KSGulin]

Approved 🙂

[bfineran]

LGTM pending comments