feat: Implement unified memory management for ONNX model inference
- Introduced a new memory management policy for ONNX Runtime to release temporary GPU memory after inference completion. - Established a shared inference entry point for all ONNX models, differentiating between long-lived sessions and temporary memory arenas. - Updated CUDA sessions to utilize a shrinkable memory arena, reclaiming GPU memory post-inference. - Ensured consistent memory release semantics across BEN2, DA3, MoGe-2, and Upscale models while reusing existing sessions. - Added lifecycle tests to validate temporary memory release and session reuse. - Modified model adapters to integrate the new memory management features, ensuring backward compatibility with existing inference behavior. - Enhanced unit tests to cover new memory release scenarios and verify correct behavior across different execution providers.
This commit is contained in:
@@ -504,8 +504,8 @@ class ServerJobsTest(unittest.TestCase):
|
||||
fake_torch.inference_mode = nullcontext
|
||||
fake_ben2 = ModuleType("server.models.onnx_ben2")
|
||||
fake_ben2.create_session = lambda _model_dir, _device: object()
|
||||
fake_ben2.infer = lambda _model, image: (
|
||||
inference_options.append((image, {})) or FakeImage()
|
||||
fake_ben2.infer = lambda _model, image, **options: (
|
||||
inference_options.append((image, options)) or FakeImage()
|
||||
)
|
||||
fake_pil = ModuleType("PIL")
|
||||
fake_pil.Image = SimpleNamespace(
|
||||
@@ -537,7 +537,7 @@ class ServerJobsTest(unittest.TestCase):
|
||||
|
||||
self.assertEqual(
|
||||
inference_options[0][1],
|
||||
{},
|
||||
{"release_memory": True},
|
||||
)
|
||||
self.assertEqual(
|
||||
save_options[0],
|
||||
@@ -556,7 +556,7 @@ class ServerJobsTest(unittest.TestCase):
|
||||
model_result = Image.new("RGBA", (2, 1), (10, 20, 30, 128))
|
||||
received = []
|
||||
|
||||
def infer(_model, image):
|
||||
def infer(_model, image, **_options):
|
||||
received.append(np.asarray(image.convert("RGB"), dtype=np.uint8))
|
||||
return model_result
|
||||
|
||||
@@ -644,8 +644,8 @@ class ServerJobsTest(unittest.TestCase):
|
||||
fake_torch.inference_mode = nullcontext
|
||||
fake_ben2 = ModuleType("server.models.onnx_ben2")
|
||||
fake_ben2.create_session = lambda _model_dir, _device: object()
|
||||
fake_ben2.infer = lambda _model, image: (
|
||||
inference_options.append((image, {})) or FakeImage()
|
||||
fake_ben2.infer = lambda _model, image, **options: (
|
||||
inference_options.append((image, options)) or FakeImage()
|
||||
)
|
||||
fake_pil = ModuleType("PIL")
|
||||
fake_pil.Image = SimpleNamespace(
|
||||
@@ -689,8 +689,10 @@ class ServerJobsTest(unittest.TestCase):
|
||||
)
|
||||
|
||||
self.assertEqual(len(inference_options), 2)
|
||||
for _image, options in inference_options:
|
||||
self.assertEqual(options, {})
|
||||
self.assertEqual(
|
||||
[options["release_memory"] for _image, options in inference_options],
|
||||
[False, True],
|
||||
)
|
||||
|
||||
def test_da3_debug_saves_only_direct_api_prediction(self):
|
||||
prediction = SimpleNamespace(
|
||||
@@ -754,6 +756,52 @@ class ServerJobsTest(unittest.TestCase):
|
||||
self.assertFalse(any(context.directory.glob("*.png")))
|
||||
self.assertEqual(calls[0][0], [str(input_path)])
|
||||
self.assertEqual(calls[0][1]["process_res"], 518)
|
||||
self.assertTrue(calls[0][1]["release_memory"])
|
||||
|
||||
def test_da3_per_frame_releases_only_after_final_frame(self):
|
||||
releases = []
|
||||
|
||||
def prediction():
|
||||
return SimpleNamespace(
|
||||
depth=np.ones((1, 2, 2), dtype=np.float32),
|
||||
conf=np.ones((1, 2, 2), dtype=np.float32),
|
||||
extrinsics=np.eye(4, dtype=np.float32)[None, :3],
|
||||
intrinsics=np.eye(3, dtype=np.float32)[None],
|
||||
processed_images=np.zeros((1, 2, 2, 3), dtype=np.uint8),
|
||||
)
|
||||
|
||||
fake_da3 = ModuleType("server.models.onnx_da3")
|
||||
|
||||
def infer(_session, _images, **options):
|
||||
releases.append(options["release_memory"])
|
||||
return prediction()
|
||||
|
||||
fake_da3.infer = infer
|
||||
fake_runtime = ModuleType("server.models.onnx_runtime")
|
||||
fake_runtime.provider_label = lambda _session: "CPU"
|
||||
parameters = {
|
||||
"model": "DA3",
|
||||
"device": "cpu",
|
||||
"process_res": 56,
|
||||
"process_res_method": "upper_bound_resize",
|
||||
"ref_view_strategy": "middle",
|
||||
"inference_mode": "per_frame",
|
||||
}
|
||||
|
||||
with patch.dict(
|
||||
sys.modules,
|
||||
{
|
||||
"server.models.onnx_da3": fake_da3,
|
||||
"server.models.onnx_runtime": fake_runtime,
|
||||
},
|
||||
):
|
||||
self.da3_model.infer(
|
||||
FakeModels(),
|
||||
parameters,
|
||||
[Path("a.png"), Path("b.png")],
|
||||
)
|
||||
|
||||
self.assertEqual(releases, [False, True])
|
||||
|
||||
def test_moge2_debug_saves_each_direct_api_dictionary(self):
|
||||
prediction = {
|
||||
@@ -786,12 +834,42 @@ class ServerJobsTest(unittest.TestCase):
|
||||
},
|
||||
)
|
||||
self.assertEqual(infer_one.call_count, 2)
|
||||
self.assertEqual(
|
||||
[call.kwargs["release_memory"] for call in infer_one.call_args_list],
|
||||
[False, True],
|
||||
)
|
||||
archives = sorted(context.directory.glob("prediction.*.npz"))
|
||||
self.assertEqual(len(archives), 2)
|
||||
with np.load(archives[0]) as result:
|
||||
self.assertEqual(set(result.files), set(prediction))
|
||||
self.assertFalse(any(context.directory.glob("*.glb")))
|
||||
|
||||
def test_moge2_batch_releases_only_after_final_image(self):
|
||||
prediction = {
|
||||
"points": np.ones((2, 2, 3), dtype=np.float32),
|
||||
"depth": np.ones((2, 2), dtype=np.float32),
|
||||
"normal": np.ones((2, 2, 3), dtype=np.float32),
|
||||
"mask": np.ones((2, 2), dtype=bool),
|
||||
"intrinsics": np.eye(3, dtype=np.float32),
|
||||
}
|
||||
with patch.object(
|
||||
self.moge2_model,
|
||||
"infer_one",
|
||||
return_value=prediction,
|
||||
) as infer_one:
|
||||
frames = self.moge2_model.infer(
|
||||
FakeModels(),
|
||||
{"model": "MOGE2_VITS_NORMAL", "device": "cpu"},
|
||||
[Path("a.png"), Path("b.png")],
|
||||
include_points=True,
|
||||
)
|
||||
|
||||
self.assertEqual(len(frames), 2)
|
||||
self.assertEqual(
|
||||
[call.kwargs["release_memory"] for call in infer_one.call_args_list],
|
||||
[False, True],
|
||||
)
|
||||
|
||||
def test_ben2_debug_saves_direct_api_foreground(self):
|
||||
class Foreground:
|
||||
def save(self, path):
|
||||
@@ -817,6 +895,9 @@ class ServerJobsTest(unittest.TestCase):
|
||||
},
|
||||
)
|
||||
infer_foreground.assert_called_once()
|
||||
self.assertTrue(
|
||||
infer_foreground.call_args.kwargs["release_memory"]
|
||||
)
|
||||
self.assertEqual(
|
||||
(context.directory / "foreground.png").read_bytes(),
|
||||
b"foreground",
|
||||
@@ -848,6 +929,10 @@ class ServerJobsTest(unittest.TestCase):
|
||||
)
|
||||
|
||||
self.assertEqual(infer_one.call_count, 2)
|
||||
self.assertEqual(
|
||||
[call.kwargs["release_memory"] for call in infer_one.call_args_list],
|
||||
[False, True],
|
||||
)
|
||||
predictions = sorted(context.directory.glob("prediction.*.png"))
|
||||
self.assertEqual(len(predictions), 2)
|
||||
with Image.open(predictions[0]) as prediction:
|
||||
@@ -880,6 +965,10 @@ class ServerJobsTest(unittest.TestCase):
|
||||
)
|
||||
|
||||
self.assertEqual(infer_one.call_count, 2)
|
||||
self.assertEqual(
|
||||
[call.kwargs["release_memory"] for call in infer_one.call_args_list],
|
||||
[False, True],
|
||||
)
|
||||
self.assertEqual(result, {"upscale": "."})
|
||||
colors = sorted(context.directory.glob("upscale.*.png"))
|
||||
self.assertEqual(
|
||||
|
||||
Reference in New Issue
Block a user