From bf7ec0233df245ff14b18472fdaa2cb2bda51a1e Mon Sep 17 00:00:00 2001 From: Robert Sachunsky Date: Thu, 21 May 2026 02:43:34 +0200 Subject: [PATCH] =?UTF-8?q?ModelZoo.load=5Fmodel:=20use=20`memory=5Flimit`?= =?UTF-8?q?=20instead=20of=20`memory=5Fgrowth`=E2=80=A6?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - growth strategy is more flexible, but uses much more VRAM - limit strategy needs to be calibrated to models (currently fixed), and batch size, but needs much less VRAM and is faster --- src/eynollah/model_zoo/model_zoo.py | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/src/eynollah/model_zoo/model_zoo.py b/src/eynollah/model_zoo/model_zoo.py index b97911a..c63a58d 100644 --- a/src/eynollah/model_zoo/model_zoo.py +++ b/src/eynollah/model_zoo/model_zoo.py @@ -169,7 +169,23 @@ class EynollahModelZoo: gpus = gpus[:1] # TF will always use first allowable tf.config.set_visible_devices(gpus, 'GPU') for device in gpus: - tf.config.experimental.set_memory_growth(device, True) + # tf.config.experimental.set_memory_growth(device, True) + # dynamic growth never frees memory (to avoid fragmentation), + # so the VRAM requirements end up much larger than feasible + # (for small GPUs); so try hard (calibrated) limits instead: + tf.config.set_logical_device_configuration( + device, + [tf.config.LogicalDeviceConfiguration(memory_limit={ + "binarization": 868, # due to bs 5 + "enhancement": 980, # due to bs 3 + "col_classifier": 210, + "page": 618, + "textline": 1680, # 954 for bs 1 + "region_1_2": 1580, + "region_fl_np": 1756, + "table": 1818, + "reading_order": 632, + }[model_category])]) vendor_name = ( tf.config.experimental.get_device_details(device) .get('device_name', 'unknown'))