Skip to content

Train

To train a pose regression model from scratch on a single patient or a set of preregistered subjects, use xvr train:

  • The --volpath argument should point to a directory containing CT volumes for training.
  • If the directory contains a single CT scan, the resulting model be patient-specific.
  • If the directory contains multiple CTs, it's beneficial to preregister them to a common reference frame (e.g., using Greedy). This will improve the accuracy of the model, but this isn't strictly necessary.
  • We use wandb to log experiments. To use this feature, set the WANDB_API_KEY environment variable by adding the following line to your .zshrc or .bashrc file:

    export WANDB_API_KEY=your_api_key
    

xvr train

$ xvr train --help
Usage: xvr train [OPTIONS]

Train a pose regression model.

Data:
* --volpath -v TEXT              CT or directory of CTs for pretraining [required]
* --outpath -o TEXT              Directory in which to save model weights [required]
  --maskpath -m TEXT             Optional labelmaps corresponding to the CTs
  --patch-size TEXT              Optional random crop size e.g. 'h,w,d'; if None,
                                 return entire volume
  --sample-weights TEXT          Probability for sampling each volume in volpath
  --num-workers INTEGER          Number of subprocesses to use in the dataloader
                                 [default: 4]
  --pin-memory --no-pin-memory   Copy volumes into CUDA pinned memory before
                                 returning [default: False]

Sampling:
* --r1 <FLOAT FLOAT>...          Range for primary angle (in degrees) [required]
* --r2 <FLOAT FLOAT>...          Range for secondary angle (in degrees) [required]
* --r3 <FLOAT FLOAT>...          Range for tertiary angle (in degrees) [required]
* --tx <FLOAT FLOAT>...          Range for x-offset (in millimeters) [required]
* --ty <FLOAT FLOAT>...          Range for y-offset (in millimeters) [required]
* --tz <FLOAT FLOAT>...          Range for z-offset (in millimeters) [required]
  --batch-size INTEGER           Number of DRRs per batch [default: 116]
  --img-threshold FLOAT          Minimum fraction of foreground pixels to keep a DRR
                                 [default: 0.1]
  --mask-threshold FLOAT         Minimum fraction of mask pixels to keep a DRR
                                 [default: 0.05]
  --n-samples INTEGER            Number of points sampled along each ray when
                                 rendering DRRs [default: 500]
  --geodesic-only                Skip re-rendering from predicted poses; supervise
    --no-geodesic-only           pose directly with geodesic loss only (fast)
                                 [default: False]

Renderer:
* --sdd FLOAT                    Source-to-detector distance (in millimeters)
                                 [required]
* --height INTEGER               DRR height (in pixels) [required]
* --delx FLOAT                   DRR pixel size (in millimeters / pixel) [required]
  --orientation TEXT             Orientation of CT volumes [choices: AP, PA]
                                 [default: AP]
  --reverse-x-axis               Obey radiologic convention (e.g., heart on right)
    --no-reverse-x-axis          [default: False]

Model:
  --model-name TEXT              Name of model to instantiate from the timm library
                                 [default: resnet18]
  --norm-layer TEXT              Normalization layer [default: groupnorm]
  --pretrained --no-pretrained   Load pretrained ImageNet-1k weights [default: 
                                 False]
  --parameterization TEXT        Parameterization of SO(3) for regression [default: 
                                 quaternion_adjugate]
  --convention TEXT              If parameterization='euler_angles', specify order
                                 [default: ZXY]
  --unit-conversion-factor FLOAT Scale factor for translation prediction (e.g., from
                                 m to mm) [default: 1000.0]
  --p-augmentation FLOAT         Base probability of image augmentations during
                                 training [default: 0.333]

Optimizer:
  --lr FLOAT                     Maximum learning rate [default: 0.0002]
  --weight-ncc FLOAT             Weight on mNCC loss term [default: 1.0]
  --weight-geo FLOAT             Weight on geodesic loss term [default: 0.01]
  --weight-dice FLOAT            Weight on Dice loss term [default: 1.0]
  --weight-haus FLOAT            Weight on Hausdorff loss term [default: 0.1]
  --n-total-itrs INTEGER         Number of iterations for training the model
                                 [default: 1000000]
  --n-warmup-itrs INTEGER        Number of iterations for warming up the learning
                                 rate [default: 1000]
  --n-grad-accum-itrs INTEGER    Number of iterations for gradient accumulation
                                 [default: 4]
  --n-save-every-itrs INTEGER    Number of iterations before saving a new model
                                 checkpoint [default: 1000]
  --disable-scheduler            Turn off cosine learning rate scheduler [default: 
    --no-disable-scheduler       False]

Checkpoint:
  --ckptpath TEXT                Checkpoint of a pretrained pose regressor
  --reuse-optimizer              Initialize the previous optimizer's state [default:
    --no-reuse-optimizer         False]
  --warp TEXT                    SimpleITK transform to warp input CT to
                                 checkpoint's reference frame
  --invert --no-invert           Whether to invert the warp or not [default: False]

Logging:
  --project TEXT                 WandB project name [default: xvr]
  --group TEXT                   WandB run group
  --name TEXT                    WandB run name
  --id TEXT                      WandB run ID (useful when restarting from a
                                 checkpoint)