diff --git a/CHANGELOG.md b/CHANGELOG.md index 1999e866..a81e4f50 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,8 @@ +## 1.6.1 + +### Enhancement +- Free intermediate arrays (`origin_img`, `img`, `ort_inputs`, `output`) and PIL pixel buffer at dead points during YoloX `image_processing()` to reduce peak memory during inference + ## 1.6.0 ### Fix diff --git a/benchmarks/__init__.py b/benchmarks/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/benchmarks/test_benchmark_yolox.py b/benchmarks/test_benchmark_yolox.py new file mode 100644 index 00000000..4074530c --- /dev/null +++ b/benchmarks/test_benchmark_yolox.py @@ -0,0 +1,63 @@ +"""Benchmark for YoloX image_processing() memory optimization. + +Uses a fake ONNX session to isolate the memory behavior of image_processing() +without requiring the real model weights. The fake session allocates a realistic +35 MiB workspace to simulate ONNX inference memory pressure. +""" + +import numpy as np +from PIL import Image as PILImage + +from unstructured_inference.models.yolox import UnstructuredYoloXModel + + +class _FakeInput: + def __init__(self) -> None: + self.name = "input" + + +class _FakeSession: + """Simulates an ONNX inference session with realistic memory allocation.""" + + def get_inputs(self): + return [_FakeInput()] + + def run(self, _names, _inputs): + workspace = np.empty((35 * 1024 * 1024,), dtype=np.uint8) # 35 MiB # noqa: F841 + # input_shape (1024,768), strides [8,16,32] → 128*96 + 64*48 + 32*24 = 16128 + return [np.random.randn(1, 16128, 16).astype(np.float32)] + + +def make_model() -> UnstructuredYoloXModel: + model = object.__new__(UnstructuredYoloXModel) + model.model = _FakeSession() + model.model_path = "yolox_fake" + model.layout_classes = { + 0: "Caption", + 1: "Footnote", + 2: "Formula", + 3: "List-item", + 4: "Page-footer", + 5: "Page-header", + 6: "Picture", + 7: "Section-header", + 8: "Table", + 9: "Text", + 10: "Title", + } + return model + + +# Letter-size page at 200 DPI — the default render resolution +def make_letter_200dpi() -> PILImage.Image: + return PILImage.fromarray(np.random.randint(0, 255, (2200, 1700, 3), dtype=np.uint8)) + + +def run_image_processing(): + model = make_model() + img = make_letter_200dpi() + return model.image_processing(img) + + +def test_benchmark_yolox_image_processing(benchmark): + benchmark(run_image_processing) diff --git a/pyproject.toml b/pyproject.toml index 68a5eb67..f50a2618 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -128,5 +128,8 @@ filterwarnings = [ "ignore::DeprecationWarning", ] +[tool.codeflash] +benchmarks-root = "benchmarks" + [tool.coverage.report] fail_under = 90 diff --git a/unstructured_inference/__version__.py b/unstructured_inference/__version__.py index e10d078b..14b8f278 100644 --- a/unstructured_inference/__version__.py +++ b/unstructured_inference/__version__.py @@ -1 +1 @@ -__version__ = "1.6.0" # pragma: no cover +__version__ = "1.6.1" # pragma: no cover diff --git a/unstructured_inference/models/yolox.py b/unstructured_inference/models/yolox.py index f467857d..cdf1f405 100644 --- a/unstructured_inference/models/yolox.py +++ b/unstructured_inference/models/yolox.py @@ -107,13 +107,17 @@ def image_processing( # TODO (benjamin): check other shapes for inference input_shape = (1024, 768) origin_img = np.array(image) + image.close() img, ratio = preprocess(origin_img, input_shape) + del origin_img # Free full-size image array before ONNX inference session = self.model ort_inputs = {session.get_inputs()[0].name: img[None, :, :, :]} output = session.run(None, ort_inputs) + del img, ort_inputs # Free preprocessed inputs after inference # TODO(benjamin): check for p6 predictions = demo_postprocess(output[0], input_shape, p6=False)[0] + del output boxes = predictions[:, :4] scores = predictions[:, 4:5] * predictions[:, 5:]