Dataset loader with LMDB image caching and checkpoint resume. Supports automatic format detection, parallel image loading, adaptive batch sizing, and item-level checkpoint for resumable processing. Example:: loader = StandardDataLoader( data_root='/data/datasets',
| 17 | |
| 18 | @DATALOADER.register_module() |
| 19 | class StandardDataLoader: |
| 20 | """Dataset loader with LMDB image caching and checkpoint resume. |
| 21 | |
| 22 | Supports automatic format detection, parallel image loading, adaptive |
| 23 | batch sizing, and item-level checkpoint for resumable processing. |
| 24 | |
| 25 | Example:: |
| 26 | |
| 27 | loader = StandardDataLoader( |
| 28 | data_root='/data/datasets', |
| 29 | dataset={'file_path': 'train.jsonl'}, |
| 30 | batch_size=32, |
| 31 | logger=logger, |
| 32 | ) |
| 33 | for batch in loader: |
| 34 | process(batch) |
| 35 | """ |
| 36 | |
| 37 | def __init__( |
| 38 | self, |
| 39 | data_root, |
| 40 | dataset, |
| 41 | batch_size, |
| 42 | parallel_loading=True, |
| 43 | num_workers=256, |
| 44 | logger=None, |
| 45 | cache_dir="~/cache/images_lmdb_sharded", |
| 46 | lmdb_num_shards=32, |
| 47 | lmdb_map_size_per_shard=1024 * 1024 * 1024 * 1024, # 1TB per shard |
| 48 | lmdb_readonly=False, |
| 49 | lmdb_lock=False, |
| 50 | resize_image=True, |
| 51 | resize_image_size=1024, |
| 52 | use_image=True, |
| 53 | use_lmdb_cache=True, # Whether to cache images in LMDB; False = read from disk directly |
| 54 | adjust_batch_size=True, # Whether to adjust batch_size based on turns |
| 55 | checkpoint_manager=None, # Checkpoint manager for item-level resume |
| 56 | **kwargs, # Accept but ignore deprecated parameters like pre_load_img |
| 57 | ): |
| 58 | self.data_root = data_root |
| 59 | self.batch_size = batch_size |
| 60 | self.original_config = None |
| 61 | self.parallel_loading = parallel_loading |
| 62 | self.num_workers = num_workers |
| 63 | self.logger = logger |
| 64 | self.resize_image = resize_image |
| 65 | self.resize_image_size = resize_image_size |
| 66 | self.use_image = use_image |
| 67 | self.use_lmdb_cache = use_lmdb_cache |
| 68 | self.adjust_batch_size = adjust_batch_size |
| 69 | self.checkpoint_manager = checkpoint_manager |
| 70 | self.name = "" |
| 71 | |
| 72 | # Store LMDB config for lazy initialization |
| 73 | self._cache_dir = os.path.expanduser(cache_dir) if cache_dir else None |
| 74 | self._lmdb_num_shards = lmdb_num_shards |
| 75 | self._lmdb_map_size_per_shard = lmdb_map_size_per_shard |
| 76 | self._lmdb_readonly = lmdb_readonly |
no outgoing calls