Read existing dataset to save time.

usnistgov · May 4, 2024 · e600009 · e600009
1 parent 575216b
commit e600009
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/alignn/lmdb_dataset.py b/alignn/lmdb_dataset.py
@@ -108,6 +108,7 @@ def get_torch_dataset(
     output_dir=".",
     tmp_name="dataset",
     map_size=1e12,
+    read_existing=True,
 ):
     """Get Torch Dataset with LMDB."""
     vals = np.array([ii[target] for ii in dataset])  # df[target].values
@@ -119,6 +120,13 @@ def get_torch_dataset(
     f.write(line)
     f.close()
     ids = []
+    if os.path.exists(tmp_name) and read_existing:
+        for idx, (d) in tqdm(enumerate(dataset), total=len(dataset)):
+            ids.append(d[id_tag])
+        dat = TorchLMDBDataset(lmdb_path=tmp_name, ids=ids)
+        print("Reading dataset", tmp_name)
+        return dat
+    ids = []
     env = lmdb.open(tmp_name, map_size=int(map_size))
     with env.begin(write=True) as txn:
         for idx, (d) in tqdm(enumerate(dataset), total=len(dataset)):