[ziw-liu]

@mattersoflight and I had a discussion about potential data utilization problem when only sampling one patch from a stack at each training step. While with longer training schedule, all ROIs will still get sampled, this does waste a lot of I/O bandwidth, since most of the data loaded from disk is not used in the batch. This problem becomes more prominent when the stack depth increases, where even local NVME pools cannot keep the GPU fed with data.

This PR implements a new data module argument train_patches_per_stack, allowing sampling multiple patches from a single stack, effective only in training. The batch size has to be divisible by this number for collation to work. When I/O-bound, this can greatly speed up data-loading. Below is a test of single-precision batch tensor with size (32, 2, 5, 2048, 2048), worker number equals number of batches loaded (all 5 batches loaded in parallel, sharing I/O bandwidth). Data is stored on a Lustre server in the local infiniband network, and cache is warmed up by running the test several times before recording the result. Increasing train_patches_per_stack results in decreased total time. However the decrease is sub-linear, as the bottleneck shifts to compute.

train_patches_per_stack=1: 15869 function calls (15786 primitive calls) in 89.330 seconds
train_patches_per_stack=2: 13055 function calls (12972 primitive calls) in 49.898 seconds
train_patches_per_stack=4: 14321 function calls (14238 primitive calls) in 33.688 seconds
train_patches_per_stack=8: 14110 function calls (14027 primitive calls) in 24.769 seconds

However the train_patches_per_stack cannot be arbitrarily large. Sampling too many patches from the same stack reduces the variance of the batch, effectively reducing the batch size (in terms of statistical power). Also the chance of sampling background will increase. In practice this should be evaluated with batch size to avoid having too few stacks for a batch.

Below is the first 16 samples of the same batch when train_patches_per_stack=4 showing some 'batch effect' for samples from the same stack.

[ziw-liu]

@mattersoflight I misinterpreted the training speed because as the epochs have train_patches_per_stack-times more batches. When taking that into account, this PR does make actual training faster.