ModelZoo.load_model: use memory_limit instead of memory_growth

- 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
This commit is contained in:
Robert Sachunsky 2026-05-21 02:43:34 +02:00
parent 94a5e9da14
commit bf7ec0233d

View file

@ -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'))