GPU not fully used, how to optimize the code

You could try to profile the data loading and check if it might be slowing down your code using the ImageNet example.

If the data loading time is not approaching zero, you might want to take a look at this post, which discusses common issues and provides more information.

If the data loading is not an issue, you might need to increase the batch size to increase the workload on the GPU.

Also, setting torch.backends.cudnn.benchmark = True at the beginning of your script might help to select faster kernels, if your input shapes are not highly variable.

1 Like