saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-30 20:18:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Fix training artifacts for 2GB+ models and MSELoss (#22414)

Browse source
This commit is contained in:
Justin Beavers 2024-10-15 17:47:16 -06:00 committed by GitHub
parent 6407d81b35
commit a5e85a950c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 6 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
orttraining/orttraining/python/training/onnxblock/blocks.py
View file

@ -54,8 +54,10 @@ class Block(ABC):
output = self.build(*args, **kwargs)
if accessor._GLOBAL_ACCESSOR.has_path:
# `save` will destructively access any external data
copied_model = copy.deepcopy(accessor._GLOBAL_ACCESSOR.model)
onnx.save(
accessor._GLOBAL_ACCESSOR.model,
copied_model,
self.temp_onnx_file_path,
save_as_external_data=True,
all_tensors_to_one_file=True,

5
orttraining/orttraining/test/python/orttraining_test_ort_apis_onnxblock.py
View file

@ -1159,7 +1159,8 @@ def test_generate_artifacts_external_data_one_file():
assert os.path.exists(os.path.join(temp_dir, "checkpoint"))
def test_generate_artifacts_external_data_separate_files():
@pytest.mark.parametrize("loss", [loss_t for loss_t in artifacts.LossType])
def test_generate_artifacts_external_data_separate_files(loss):
with tempfile.TemporaryDirectory() as temp_dir:
_, simple_net = _get_models("cpu", 32, 28, 10, 10)
@ -1176,7 +1177,7 @@ def test_generate_artifacts_external_data_separate_files():
artifacts.generate_artifacts(
os.path.join(temp_dir, "simple_net.onnx"),
requires_grad=requires_grad_params,
loss=artifacts.LossType.CrossEntropyLoss,
loss=loss,
optimizer=artifacts.OptimType.AdamW,
artifact_directory=temp_dir,
)