SLR training script

Author: Lojze Žust

README.md

@@ -60,6 +60,25 @@ For additional details please refer to the [paper](https://arxiv.org/abs/2108.00
 
 ### SLR Training
 
+Use the utility script `tools/train_slr.sh` to train a model using the entire SLR pipeline. 
+
+```bash
+chmod +x tools/train_slr.sh
+tools/train_slr.sh
+```
+
+The script contains the following variables, that can be changed to achieve the desired results.
+
+- `MASTR_DIR`: Location of the dataset used for training.
+- `ARCHITECTURE`: Which architecture to use (use `python tools/train.py warmup --help` for more info).
+- `MODEL_NAME`: Name of the model. Used for saving logs and weights.
+- `BATCH_SIZE`: Batch size per gpu.
+- `WARMUP_EPOCHS`: Number of epochs for the warm-up phase.
+- `FINETUNE_EPOCHS`: Number of epochs for the fine-tuning phase.
+- `NUM_ITER`: Number of iterations of the SLR pseudo label estimation and fine-tuning.
+
+Individual steps of the SLR pipeline can also be executed separately, with the following python scripts.
+
 #### Step I: Feature warm-up
 
 Train an initial model on partial labels generated from weak annotations and IMU. Uses additional object-wise losses.
@@ -71,6 +90,7 @@ python tools/train.py warmup \
 --batch-size 4
 ```
 
+Use the `--help` switch for more details on all possible arguments and settings.
 
 #### Step II: Generate pseudo labels
 
@@ -85,6 +105,9 @@ python tools/generate_pseudo_labels.py \
 
 This creates the pseudo-labels and stores them into `output/pseudo_labels/wasr_slr_warmup_v0`.
 
+
+Use the `--help` switch for more details on all possible arguments and settings.
+
 #### Step III: Fine-tune model
 
 Fine-tune the initial model on the estimated pseudo-labels from the previous step. 
@@ -99,6 +122,9 @@ python tools/train.py finetune \
 --pretrained-weights output/logs/wasr_slr_warmup/version_0/checkpoints/last.ckpt \
 --mask-dir output/pseudo_labels/wasr_slr_warmup_v0
 ```
+
+Use the `--help` switch for more details on all possible arguments and settings.
+
 ### Inference
 
 #### General inference
@@ -109,19 +135,23 @@ Run inference using a trained model. `tools/general_inference.py` script is able
 export CUDA_VISIBLE_DEVICES=0,1
 python tools/general_inference.py \
 --architecture wasr_resnet101 \
---weights-file output/logs/wasr_slr/version_0/checkpoints/last.ckpt \
+--weights-file output/logs/wasr_slr_v2_it1/version_0/checkpoints/last.ckpt \
 --image-dir data/example_dir \
 --output-dir output/predictions/test_predictions
 ```
 
 Additionally, `--imu-dir` can be used to supply a directory with corresponding IMU horizon masks. The directory structure should match the one of image dir.
 
-**NOTE**: The IMU dir has to be provided for models architectures relying on IMU data (i.e. WaSR).
+**NOTE**: The IMU dir has to be provided for models architectures relying on IMU data (i.e. WaSR with IMU).
+
+Use the `--help` switch for more details on all possible arguments and settings.
 
 #### MODS inference
 
 `tools/mods_inference.py` can be used in a similar fashion to run inference on the MODS benchmark.
 
+Use the `--help` switch for more details on all possible arguments and settings.
+
 ## Pretrained models
 
 Currently available pretrained model weights. All models are trained on the MaSTr1325 dataset using SLR and weak annotations.

tools/train_slr.sh

@@ -0,0 +1,51 @@
+#!/bin/bash
+
+# Arguments:
+MASTR_DIR=data/mastr1325
+ARCHITECTURE=wasr_resnet101_imu
+MODEL_NAME=wasr_slr_v2
+BATCH_SIZE=3
+WARMUP_EPOCHS=25
+FINETUNE_EPOCHS=50
+NUM_ITER=2
+
+# 1. Warm-up model
+echo "1. Warm-up model"
+python tools/train.py warmup \
+    --architecture $ARCHITECTURE \
+    --model-name ${MODEL_NAME}_it0 \
+    --batch-size $BATCH_SIZE \
+    --epochs $WARMUP_EPOCHS
+
+for i in $(seq 1 $NUM_ITER)
+do  
+    let "PREV_I = $i - 1"
+    CUR_I=$i
+    PREV_VERSION=$(ls -t1 output/logs/${MODEL_NAME}_it${PREV_I} | head -n 1 | cut -d _ -f 2)
+
+    PREV_MODEL_FILENAME="${MODEL_NAME}_it${PREV_I}_v${PREV_VERSION}"
+    PREV_MODEL_WEIGHTS=output/logs/${MODEL_NAME}_it${PREV_I}/version_${PREV_VERSION}/checkpoints/last.ckpt
+    FILLED_MASKS_DIR=output/pseudo_labels/${PREV_MODEL_FILENAME}
+
+    echo "-------------------------"
+    echo "Fine-tuning, iteration $CUR_I"
+    echo "-------------------------"
+
+    # 2. Estimate pseudo labels
+    echo "2. Estimate pseudo labels"
+    python tools/generate_pseudo_labels.py \
+        --architecture $ARCHITECTURE \
+        --weights-file $PREV_MODEL_WEIGHTS \
+        --output-dir $FILLED_MASKS_DIR
+
+    # 3. Re-train model
+    echo "3. Retrain model"
+    python tools/train.py finetune \
+        --architecture $ARCHITECTURE \
+        --model-name ${MODEL_NAME}_it${CUR_I} \
+        --batch-size $BATCH_SIZE \
+        --pretrained-weights $PREV_MODEL_WEIGHTS \
+        --mask-dir $FILLED_MASKS_DIR \
+        --epochs $FINETUNE_EPOCHS
+
+done