Train
To train a pose regression model from scratch on a single patient or a set of preregistered subjects, use xvr train:
- The
--volpathargument 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
wandbto log experiments. To use this feature, set theWANDB_API_KEYenvironment variable by adding the following line to your.zshrcor.bashrcfile:
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)