Monitor data loading time with lightning

[arthursw]

I would like to monitor the data loading and processing ; and compare the time it takes to prepare the queue on the CPU to the time it takes to process a batch on the GPU.

The lightning Trainer can be initialized with a DeviceStatsMonitor callback, and gives information like:

-------------------------------------------------------------------------------------------------------------------------------------------------------
|  Action                                                 |  Mean duration (s)    |  Num calls        |  Total time (s)       |  Percentage %         |
-------------------------------------------------------------------------------------------------------------------------------------------------------
|  Total                                                  |  -                    |  14374            |  1374.9               |  100 %                |
-------------------------------------------------------------------------------------------------------------------------------------------------------
|
|  run_training_epoch                                     |  272.48               |  5                |  1362.4               |  99.089               |
|  run_training_batch                                     |  3.4777               |  310              |  1078.1               |  78.41                |
[...]
|  [_TrainingEpochLoop].train_dataloader_next             |  0.44467              |  310              |  137.85               |  10.026               |

I guess 0.44467 seconds is the time it took for the DataLoader to get the next batch which was already prepared by the tio.Queue? So this does not reflect the total queue processing time?
How can I measure it? (and is it possible with lightning?)

[arthursw]

I just created a test script to profile processing speed ; it creates a queue of patches and consumes them:

import argparse
import datetime
import multiprocessing
from pathlib import Path
import torch
import torchio as tio
import lovely_tensors as lt
lt.monkey_patch()

seed = 42 
num_workers = multiprocessing.cpu_count() // 2

parser = argparse.ArgumentParser(prog=__file__, formatter_class=argparse.ArgumentDefaultsHelpFormatter, description="""Test data augmentation.""")
parser.add_argument('-d', '--dataset', required=True, help='Path to the test dataset folder.')
args = parser.parse_args()

dataset_dir = Path(args.dataset)
subject_paths = list((dataset_dir).iterdir())
subjects = []
for subject_path in subject_paths:
	subject = tio.Subject(
		flair=tio.ScalarImage(subject_path / 't2.nii.gz'),
		segmentation=tio.LabelMap(subject_path / 'seg.nii.gz'),
	)
	subjects.append(subject)
dataset = tio.SubjectsDataset(subjects)

training_transform = tio.Compose([
	tio.EnsureShapeMultiple(8, method='pad'),
	tio.RandomMotion(p=0.3),
	tio.RandomGamma(p=0.3),
	tio.RandomBiasField(p=0.3),
	tio.RandomBlur(std=(0.5, 1.), p=0.3),
	tio.RandomAnisotropy(downsampling=(1.1, 2.0), p=0.3),
	tio.RandomFlip(axes=(0, 1, 2), flip_probability=0.5, p=1.0),
	tio.OneOf({
		tio.RandomAffine(scales=0.2, degrees=(15, 15, 15), translation=0, center='image'): 0.8,
		tio.RandomElasticDeformation(num_control_points=(4, 6, 7), max_displacement=(3, 10, 24), locked_borders=0): 0.2,
	}, p=1.0),
	tio.ZNormalization(masking_method=tio.ZNormalization.mean),
	tio.RandomNoise(std=(0, 0.2), p=0.3),
	tio.OneHot(num_classes=2),
])

training_subjects = subjects

training_batch_size = 2

patch_size = (48, 128, 496)
samples_per_volume = 1
max_queue_length = 13

training_set = tio.SubjectsDataset(training_subjects, transform=training_transform)

sampler = tio.data.LabelSampler(patch_size, 'segmentation', label_probabilities={0: 0.3, 1: 0.6})

patches_training_set = tio.Queue(
	subjects_dataset=training_set,
	max_length=max_queue_length,
	samples_per_volume=samples_per_volume,
	sampler=sampler,
	num_workers=num_workers,
	shuffle_subjects=True,
	shuffle_patches=True,
	verbose=True
)

training_loader_patches = torch.utils.data.DataLoader(patches_training_set, batch_size=training_batch_size)

print('start: ', datetime.datetime.now())
for n in range(10):
	for batch_index, batch in enumerate(training_loader_patches):
		print('   batch ', batch_index, n, datetime.datetime.now())
		
		one_batch_flairs = batch['flair'][tio.DATA]
		one_batch_segmentations = batch['segmentation'][tio.DATA]
		print('   path: ', batch['flair']['path'], one_batch_flairs.max(), one_batch_segmentations.max())

I could see the time needed to fill the queue and the patch list.

[romainVala]

Hi
thanks for sharing, did you get different result than with DeviceStatsMonitor ?

[arthursw]

I didn't use both techniques together so I did not compare. I assumed ligthning is not aware that transforms are applied on the CPU while the training happens on the GPU at the same time, but I don't know.