(dataset, map_fn, cache_dir, cache_file_prefix='', new_fingerprint_args=None, regenerate_cache=False, caching_batch_size=1)
| 82 | |
| 83 | |
| 84 | def _map_and_cache(dataset, map_fn, cache_dir, cache_file_prefix='', new_fingerprint_args=None, regenerate_cache=False, caching_batch_size=1): |
| 85 | new_fingerprint_args = [] if new_fingerprint_args is None else new_fingerprint_args |
| 86 | new_fingerprint_args.append(dataset._fingerprint) |
| 87 | new_fingerprint = Hasher.hash(new_fingerprint_args) |
| 88 | if cache_file_prefix: |
| 89 | cache_dir = cache_dir / cache_file_prefix.strip('_') |
| 90 | |
| 91 | cache = Cache(cache_dir, new_fingerprint, shard_size_gb=10) |
| 92 | |
| 93 | if map_fn is None: |
| 94 | # loading directly from cache without mapping |
| 95 | assert new_fingerprint == cache.fingerprint |
| 96 | return cache |
| 97 | |
| 98 | if regenerate_cache: |
| 99 | cache.clear() |
| 100 | |
| 101 | # Cache has either been cleared if fingerprint didn't match, or has some (maybe 0) existing items in it. |
| 102 | |
| 103 | # Skip existing items |
| 104 | cache_size = len(cache) |
| 105 | dataset_size = len(dataset) |
| 106 | assert cache_size <= dataset_size |
| 107 | if cache_size == dataset_size: |
| 108 | return cache |
| 109 | dataset = dataset.select(range(cache_size, dataset_size), keep_in_memory=True) |
| 110 | |
| 111 | # Let each worker process know its rank |
| 112 | manager = mp.Manager() |
| 113 | id_queue = manager.Queue() |
| 114 | |
| 115 | def init(queue): |
| 116 | global rank |
| 117 | rank = queue.get() |
| 118 | |
| 119 | for i in range(NUM_PROC): |
| 120 | id_queue.put(i) |
| 121 | |
| 122 | pool = mp.Pool(NUM_PROC, init, (id_queue,)) |
| 123 | |
| 124 | def wrapper(example): |
| 125 | global rank |
| 126 | return map_fn(example, rank) |
| 127 | |
| 128 | # Tensor slices reference the entire memory of the original tensor, and everything would be pickled and stored |
| 129 | # in cache, so we do this. |
| 130 | def recursive_clone_tensors(obj): |
| 131 | if torch.is_tensor(obj): |
| 132 | return obj.clone() |
| 133 | elif isinstance(obj, dict): |
| 134 | for k, v in obj.items(): |
| 135 | obj[k] = recursive_clone_tensors(v) |
| 136 | return obj |
| 137 | elif isinstance(obj, (list, tuple)): |
| 138 | return [recursive_clone_tensors(x) for x in obj] |
| 139 | else: |
| 140 | return obj |
| 141 |
no test coverage detected