CHAP2.0 CHAP-FT Fine-Tuning

CHAP2.0 fine-tuning uses the split HDF5 files generated by create_dataset_split.py and trains the model with main_finetune.py.

If you only want to apply an existing checkpoint to daily HDF5 files, you do not need to fine-tune a model. Use Generating Predictions instead.

Prerequisite

Before running fine-tuning, create CHAP2.0 split HDF5 files:

10s_train.h5
10s_val.h5
10s_test_complete.h5

See Creating Dataset Splits for the conversion from daily preprocessed HDF5 files to these split files.

Fine-Tuning Command

From inside the CHAP2/ directory, invoke main_finetune.py as follows:

torchrun --nproc_per_node=<num_gpus> -m main_finetune \
  --data_path <split_data_dir> \
  --model CHAP \
  --checkpoint ../MSSE-2021/pre-trained-models-pt/CHAP_ALL_ADULTS.pth \
  --output_dir <checkpoint_output_dir> \
  --remark <experiment_name> \
  --blr 1e-3 \
  --epochs 10 \
  --warmup_epochs 2 \
  --batch_size 16 \
  --weight_decay 1e-3 \
  --use_data_aug 1

For single-GPU or CPU runs, python -m main_finetune can be used instead of torchrun.

Complete usage details of the most commonly used arguments are as follows:

usage: main_finetune.py --data_path DATA_PATH
                        [--model MODEL]
                        [--checkpoint CHECKPOINT]
                        [--output_dir OUTPUT_DIR]
                        [--log_dir LOG_DIR]
                        [--remark REMARK]
                        [--epochs EPOCHS]
                        [--batch_size BATCH_SIZE]
                        [--accum_iter ACCUM_ITER]
                        [--blr BLR]
                        [--lr LR]
                        [--weight_decay WEIGHT_DECAY]
                        [--warmup_epochs WARMUP_EPOCHS]
                        [--pos_weight POS_WEIGHT]
                        [--use_focal_loss]
                        [--subset_ratio SUBSET_RATIO]
                        [--use_data_aug USE_DATA_AUG]
                        [--eval EVAL]
                        [--make_prediction]
                        [--prediction_dir PREDICTION_DIR]
                        [--device DEVICE]
                        [--num_workers NUM_WORKERS]
                        [--pin_mem | --no_pin_mem]

required arguments:
  --data_path DATA_PATH
      Path to the split HDF5 directory containing 10s_train.h5,
      10s_val.h5, and 10s_test_complete.h5.

optional arguments:
  --model MODEL
      Model architecture to train. The main CHAP2.0 workflow uses CHAP.
      Default: CHAP.
  --checkpoint CHECKPOINT
      Checkpoint used to initialize the model. For CHAP-FT, this is usually
      an MSSE/CHAP1 pretrained checkpoint such as CHAP_ALL_ADULTS.pth.
      Default: None.
  --output_dir OUTPUT_DIR
      Directory where checkpoints and training outputs are saved.
      Default: ./output.
  --log_dir LOG_DIR
      Directory for W&B logs. If unset, W&B logging is disabled.
      Default: None.
  --remark REMARK
      Experiment name used in output and log directory names.
      Default: Debug.
  --epochs EPOCHS
      Number of fine-tuning epochs. Default: 20.
  --batch_size BATCH_SIZE
      Batch size per GPU. Default: 64.
  --accum_iter ACCUM_ITER
      Gradient accumulation steps. Effective batch size is
      batch_size * accum_iter * number of GPUs. Default: 1.
  --blr BLR
      Base learning rate. The absolute learning rate is scaled by total
      batch size when --lr is not provided. Default: 5e-4.
  --lr LR
      Absolute learning rate. If provided, overrides --blr scaling.
      Default: None.
  --weight_decay WEIGHT_DECAY
      Weight decay. Default: 5e-2.
  --warmup_epochs WARMUP_EPOCHS
      Number of learning-rate warmup epochs. Default: 2.
  --pos_weight POS_WEIGHT
      Positive-class weight for BCEWithLogitsLoss. Default: 1.0.
  --use_focal_loss
      Use focal loss instead of BCEWithLogitsLoss.
  --subset_ratio SUBSET_RATIO
      Fraction of the dataset to use. Default: 1.0.
  --use_data_aug USE_DATA_AUG
      Whether to use data augmentation. Default: 1.
  --eval EVAL
      Run evaluation only using the checkpoint path provided here.
      Default: None.
  --make_prediction
      During split-based evaluation, also write prediction outputs.
  --prediction_dir PREDICTION_DIR
      Directory for predictions produced during split-based evaluation.
  --device DEVICE
      Device used for training or evaluation. Default: cuda.
  --num_workers NUM_WORKERS
      Number of DataLoader workers. Default: 4.
  --pin_mem / --no_pin_mem
      Enable or disable pinned CPU memory in the DataLoader.
      Default: --pin_mem.

Model Initialization

CHAP-FT initializes from an MSSE/CHAP1 pretrained checkpoint, usually CHAP_ALL_ADULTS.pth, and fine-tunes that checkpoint on a target CHAP2.0 dataset.

If --checkpoint is not provided, the script trains a CHAP model from scratch. That is supported by the code, but it is not the main CHAP-FT workflow.

Class Imbalance

--pos_weight controls the positive-class weight used by BCEWithLogitsLoss. This can be useful when the training labels are imbalanced between sitting and not-sitting windows.

The helper script pre_process/find_label_dist.py can be used to estimate label distribution and pos_weight from training subjects. It is an optional utility and is most useful when working from daily HDF5 files and a subject split dictionary.

W&B Logging

W&B logging is disabled by default. To enable it, pass --log_dir <log_dir> and log in before running fine-tuning:

wandb login

Alternatively, set WANDB_API_KEY in your shell environment before running main_finetune.py.

Example Scripts

CHAP2.0 includes example fine-tuning scripts:

script/chap_ft_sol.sh       # SOL/PASOS CHAP fine-tuning example
script/chap_ft_iwatch.sh    # iWatch CHAP fine-tuning example
script/chap_scratch_sol.sh  # SOL/PASOS training-from-scratch example

These scripts contain lab-specific paths and should be edited before use.

Evaluation and Prediction

main_finetune.py can also run evaluation on split HDF5 data:

python -m main_finetune \
  --data_path <split_data_dir> \
  --model CHAP \
  --eval <checkpoint_path>

To write prediction files during split-based evaluation, add:

--make_prediction --prediction_dir <prediction_output_dir>

This is different from the prediction-only workflow. For prediction-only use on daily HDF5 files, use make_predictions.py as described in Generating Predictions.