csv logger

mwalmsley · Nov 21, 2023 · a038550 · a038550
1 parent 0212fcd
commit a038550
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/only_for_me/narval/train.sh b/only_for_me/narval/train.sh
@@ -21,7 +21,7 @@ export NCCL_BLOCKING_WAIT=1  #Set this environment variable if you wish to use t
 
 REPO_DIR=/project/def-bovy/walml/zoobot
 srun $PYTHON $REPO_DIR/only_for_me/narval/train.py \
-    --save-dir $REPO_DIR/only_for_me/narval/debug_models \
+    --save-dir $REPO_DIR/only_for_me/narval/debug_models_v2 \
     --batch-size 128 \
     --gpus 4 \
     --color --mixed-precision

diff --git a/zoobot/pytorch/training/train_with_pytorch_lightning.py b/zoobot/pytorch/training/train_with_pytorch_lightning.py
@@ -7,6 +7,7 @@
 from pytorch_lightning.strategies.ddp import DDPStrategy
 from pytorch_lightning.callbacks import ModelCheckpoint
 from pytorch_lightning.callbacks.early_stopping import EarlyStopping
+from pytorch_lightning.loggers import CSVLogger
 
 from galaxy_datasets.pytorch.galaxy_datamodule import GalaxyDataModule
 
@@ -190,8 +191,8 @@ def train_default_zoobot_from_scratch(
             'framework': 'pytorch'
         })
     else:
-        logging.warning('No wandb_logger passed. Disabling logging')
-        wandb_logger = False
+        logging.warning('No wandb_logger passed. Using CSV logging only')
+        wandb_logger = CSVLogger(save_dir=save_dir)
 
     # work out what dataset the user has passed
     single_catalog = catalog is not None