Parameters that effect memory and compute usage

[laynr]

I am running out of memory running the following code. Would anyone be willing to provide insight into what parameters effect compute resources the most when using torchio? Specifically wondering about 'queue_length', 'samples_per_volume', 'num_workers', and 'batch_size'.

    # dataset contains  500 subjects
    patch_size = (128, 128, 64)
    queue_length = 300
    samples_per_volume = 10
    sampler = tio.data.UniformSampler(patch_size)
    patches_queue = tio.Queue(
        dataset,
        queue_length,
        samples_per_volume,
        sampler,
        num_workers=4,
    )
    patches_loader = DataLoader(patches_queue, batch_size=16)    

    num_epochs = 2
    model = torch.nn.Identity()
    #model = torch.hub.load('fepegar/highresnet', 'highres3dnet', pretrained=True)
    for epoch_index in range(num_epochs):
        for patches_batch in patches_loader:
            inputs = patches_batch['image'][tio.DATA]
            targets = patches_batch['label']
            logits = model(inputs)
    
    torch.save(model.state_dict(), f'./saved_model')

[fepegar]

Hi, @laynr. It might also be related to the way you build your dataset. First, what is the output of patches_queue.get_max_memory_pretty() (docs)? How much RAM do you have?

[romainVala]

I agree this is not a torchio issue, but still, it will be good, to have an estimation (within torchio) of what GPU and CPU memory will be. no ?
(I did not play with tio.Queue.get_max_memory_pretty(), may it is already here, it is only cpu, and takes account of num_worker ?)

[fepegar]

I'm not sure I understand your question, @romainVala.

[romainVala]

get_max_memory, is only computing the CPU mem needed for the queue
but this is only half of the story : One should also compute the memory needed by each worker.
(which then depend on the total volume size, and the number of patch per volume)

(Note that , GPU memory is very different since it is only dependent on the input size, and the model used)

[fepegar]

The memory needed by each worker is much harder to compute because it depends on how much memory is used internally by each transform.

[romainVala]

yes you are right I forgot the transforms
may be an estimate without transform (min memory) could be of interest, to understand (and help to choose) num_worker or sample_per_volume

And then for each transform this should linked to the volume input size ... and one could start with a rough estimate (2 or 3 * input_size) depending on the transform.

Not sure if it is possible, but it is important. I get often issues with cpu memory on my GPU cluster, and I have no other choice, than try and test, to adapt those parameters