How to select the checkpoint for predicting?

[NiekPennings]

When using the predict function, I would like to specify which checkpoint to use. For example, using checkpoint_best.pth instead of checkpoint_final.pth (or vice versa, I don't know which one is exactly used). I looked at the code but couldn't figure out where the checkpoint is loaded.

The reason I want to do this is that I saved the checkpoints at every 50-epoch interval, and I want to gather the results using the predict function with the different checkpoints. The function provides a very nice summary.json that I want to use for the different versions.

Thanks!

[vmiller987]

@NiekPennings

You should be able to by specifying it with the -chk flag. It defaults to the final_checkpoint as well.

I made a nnUNetv2 CLI Catalog awhile ago. Here's the section on prediction. Hope this helps!

## Command: `nnUNetv2_predict`

### Arugments
- `-i`
  - **Type**: `str`
  - **Required**: Yes
  - **Help**: Input folder. Ensure the correct channel numbering for your files (e.g., `_0000`). File endings must match the training dataset.

- `-o`
  - **Type**: `str`
  - **Required**: Yes
  - **Help**: Output folder. If it does not exist, it will be created. Predicted segmentations will have the same name as their source images.

- `-m`
  - **Type**: `str`
  - **Required**: Yes
  - **Help**: Folder containing the trained model. Must have subfolders `fold_X` for the different folds you trained.

- `-f`
  - **Type**: `list`
  - **Required**: No
  - **Default**: `(0, 1, 2, 3, 4)`
  - **Help**: Specify the folds of the trained model to be used for prediction. Default: `(0, 1, 2, 3, 4)`.

- `-step_size`
  - **Type**: `float`
  - **Required**: No
  - **Default**: `0.5`
  - **Help**: Step size for sliding window prediction. Larger values are faster but less accurate. Cannot exceed 1. Default: `0.5`.

- `--disable_tta`
  - **Type**: `flag`
  - **Required**: No
  - **Help**: Disable test-time augmentation (mirroring). Faster but less accurate inference. Not recommended.

- `--verbose`
  - **Type**: `flag`
  - **Required**: No
  - **Help**: Enable verbose output.

- `--save_probabilities`
  - **Type**: `flag`
  - **Required**: No
  - **Help**: Export predicted class probabilities. Required for ensembling multiple configurations.

- `--continue_prediction`
  - **Type**: `flag`
  - **Required**: No
  - **Help**: Continue an aborted previous prediction (will not overwrite existing files).

- `-chk`
  - **Type**: `str`
  - **Required**: No
  - **Default**: `checkpoint_final.pth`
  - **Help**: Name of the checkpoint to use. Default: `checkpoint_final.pth`.

- `-npp`
  - **Type**: `int`
  - **Required**: No
  - **Default**: `3`
  - **Help**: Number of processes used for preprocessing. Beware of out-of-RAM issues. Default: `3`.

- `-nps`
  - **Type**: `int`
  - **Required**: No
  - **Default**: `3`
  - **Help**: Number of processes used for segmentation export. Beware of out-of-RAM issues. Default: `3`.

- `-prev_stage_predictions`
  - **Type**: `str`
  - **Required**: No
  - **Default**: `None`
  - **Help**: Folder containing the predictions of the previous stage. Required for cascaded models.

- `-device`
  - **Type**: `str`
  - **Required**: No
  - **Default**: `cuda`
  - **Help**: Set the device for inference. Options: `cuda` (GPU), `cpu` (CPU), `mps` (Apple M1/M2). Do not use this to set GPU ID.

- `--disable_progress_bar`
  - **Type**: `flag`
  - **Required**: No
  - **Help**: Disable the progress bar. Recommended for non-interactive HPC environments.

[NiekPennings]

Ah you’re right, I must’ve missed it. Thanks!