Added cache_size parameter to the post-processing training config

Added `cache_size` parameter to the post-processing training config
OML-Team · Jun 17, 2024 · 8bf558a · 8bf558a
1 parent 0e3776f
commit 8bf558a
Show file tree

Hide file tree

Showing 4 changed files with 5 additions and 1 deletion.
diff --git a/oml/lightning/pipelines/train.py b/oml/lightning/pipelines/train.py
@@ -32,7 +32,7 @@ def get_retrieval_loaders(cfg: TCfg) -> Tuple[DataLoader, DataLoader, ILabeledDa
         transforms_train=get_transforms_by_cfg(cfg["transforms_train"]),
         transforms_val=get_transforms_by_cfg(cfg["transforms_val"]),
         dataframe_name=cfg["dataframe_name"],
-        cache_size=cfg["cache_size"],
+        cache_size=cfg.get("cache_size", 0),
         verbose=cfg.get("show_dataset_warnings", True),
     )
 

diff --git a/oml/lightning/pipelines/train_postprocessor.py b/oml/lightning/pipelines/train_postprocessor.py
@@ -88,13 +88,15 @@ def get_loaders_with_embeddings(
 
     train_dataset = ImageLabeledDataset(
         dataset_root=cfg["dataset_root"],
+        cache_size=cfg.get("cache_size", 0),
         df=train_extraction.df,
         transform=get_transforms_by_cfg(cfg["transforms_train"]),
         extra_data={EMBEDDINGS_KEY: emb_train},
     )
 
     valid_dataset = ImageQueryGalleryLabeledDataset(
         dataset_root=cfg["dataset_root"],
+        cache_size=cfg.get("cache_size", 0),
         df=val_extraction.df,
         transform=transforms_extraction,
         extra_data={EMBEDDINGS_KEY: emb_val},

diff --git a/pipelines/postprocessing/pairwise_postprocessing/postprocessor_train.yaml b/pipelines/postprocessing/pairwise_postprocessing/postprocessor_train.yaml
@@ -12,6 +12,7 @@ dataframe_name: df.csv
 logs_folder: ${now:%Y-%m-%d_%H-%M-%S}_${postfix}
 
 num_workers: 20
+cache_size: 0
 max_epochs: 1000
 valid_period: 10
 

diff --git a/tests/test_runs/test_pipelines/configs/train_postprocessor.yaml b/tests/test_runs/test_pipelines/configs/train_postprocessor.yaml
@@ -11,6 +11,7 @@ dataframe_name: df_with_category.csv
 logs_root: logs
 
 num_workers: 0
+cache_size: 50
 max_epochs: 2
 valid_period: 1