Update config example parameters (#198)

Update config example for batch creation
openclimatefix · May 17, 2024 · 130846c · 130846c
1 parent f9831cd
commit 130846c
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 18 deletions.
diff --git a/README.md b/README.md
@@ -152,10 +152,15 @@ Where `FULL-PATH-TO-REPO` represent the whole path to the PVNet repo on your loc
 
 ### Running the batch creation script
 
-Run the save_batches.py script to create batches with the following example arguments as:
+Run the save_batches.py script to create batches if setting parameters in the datamodule config (`streamed_batches.yaml` in this example):
 
 ```
-python scripts/save_batches.py datamodule=streamed_batches +batch_output_dir="./output" +num_train_batches=10 +num_val_batches=5
+python scripts/save_batches.py
+```
+or with the following example arguments to override config:
+
+```
+python scripts/save_batches.py datamodule=streamed_batches datamodule.batch_output_dir="./output" datamodule.num_train_batches=10 datamodule.num_val_batches=5
 ```
 
 In this function the datamodule argument looks for a config under `PVNet/configs/datamodule`. The examples here are either to use "premade_batches" or "streamed_batches".

diff --git a/configs.example/datamodule/streamed_batches.yaml b/configs.example/datamodule/streamed_batches.yaml
@@ -6,6 +6,11 @@ configuration: "PLACEHOLDER.yaml"
 num_workers: 20
 prefetch_factor: 2
 batch_size: 8
+batch_output_dir: "PLACEHOLDER"
+num_train_batches: 2
+num_val_batches: 1
+
+
 train_period:
   - null
   - "2022-05-07"

diff --git a/scripts/save_batches.py b/scripts/save_batches.py
@@ -5,15 +5,20 @@
 the same config file currently set to train the model.
 
 use:
+```
+python save_batches.py
+```
+if setting all values in the datamodule config file, or
+
 ```
 python save_batches.py \
-    +batch_output_dir="/mnt/disks/bigbatches/batches_v0" \
+    datamodule.batch_output_dir="/mnt/disks/bigbatches/batches_v0" \
     datamodule.batch_size=2 \
     datamodule.num_workers=2 \
-    +num_train_batches=0 \
-    +num_val_batches=2
+    datamodule.num_train_batches=0 \
+    datamodule.num_val_batches=2
 ```
-
+if wanting to override these values for example
 """
 # This is needed to get multiprocessing/multiple workers to behave
 try:
@@ -110,12 +115,14 @@ def main(config: DictConfig):
     print_config(config, resolve=False)
 
     # Set up directory
-    os.makedirs(config.batch_output_dir, exist_ok=False)
+    os.makedirs(config_dm.batch_output_dir, exist_ok=False)
 
-    with open(f"{config.batch_output_dir}/datamodule.yaml", "w") as f:
-        f.write(OmegaConf.to_yaml(config.datamodule))
+    with open(f"{config_dm.batch_output_dir}/datamodule.yaml", "w") as f:
+        f.write(OmegaConf.to_yaml(config_dm))
 
-    shutil.copyfile(config_dm.configuration, f"{config.batch_output_dir}/data_configuration.yaml")
+    shutil.copyfile(
+        config_dm.configuration, f"{config_dm.batch_output_dir}/data_configuration.yaml"
+    )
 
     dataloader_kwargs = dict(
         shuffle=False,
@@ -132,8 +139,8 @@ def main(config: DictConfig):
         persistent_workers=False,
     )
 
-    if config.num_val_batches > 0:
-        os.mkdir(f"{config.batch_output_dir}/val")
+    if config_dm.num_val_batches > 0:
+        os.mkdir(f"{config_dm.batch_output_dir}/val")
         print("----- Saving val batches -----")
 
         val_batch_pipe = _get_datapipe(
@@ -145,14 +152,14 @@ def main(config: DictConfig):
 
         _save_batches_with_dataloader(
             batch_pipe=val_batch_pipe,
-            batch_dir=f"{config.batch_output_dir}/val",
-            num_batches=config.num_val_batches,
+            batch_dir=f"{config_dm.batch_output_dir}/val",
+            num_batches=config_dm.num_val_batches,
             dataloader_kwargs=dataloader_kwargs,
             output_format="torch" if config.renewable == "pv" else "netcdf",
         )
 
-    if config.num_train_batches > 0:
-        os.mkdir(f"{config.batch_output_dir}/train")
+    if config_dm.num_train_batches > 0:
+        os.mkdir(f"{config_dm.batch_output_dir}/train")
         print("----- Saving train batches -----")
 
         train_batch_pipe = _get_datapipe(
@@ -164,8 +171,8 @@ def main(config: DictConfig):
 
         _save_batches_with_dataloader(
             batch_pipe=train_batch_pipe,
-            batch_dir=f"{config.batch_output_dir}/train",
-            num_batches=config.num_train_batches,
+            batch_dir=f"{config_dm.batch_output_dir}/train",
+            num_batches=config_dm.num_train_batches,
             dataloader_kwargs=dataloader_kwargs,
             output_format="torch" if config.renewable == "pv" else "netcdf",
         )