Loading in a large tensor just to take a small crop takes too long

I’m trying to train a GAN and loading each batch takes about 3-4x as long as the actual training for each batch. Is this normal? It doesn’t seem to be the result of the transformations I apply. The data is stored as PyTorch tensors in .pt files.

Each tensor is of shape (3601, 3601), which is massive, but I only load each one in order to take a single random crop of size (256, 256), which then gets put into the batch. I strongly suspect the issue is that I am using torch.load() to load in these massive images even though I end up needing only like 0.5% of their data.

Is there any easy way to do this crop without loading in the entire image? Even if it requires me changing the way I have my data stored. If there is no easy way, would it be doable to write my own loading function that builds the cropped image directly from the.pt file containing the larger image?

Thanks in advance for any advice.

You could check, if np.memmap could work to load only parts of a larger file into memory, which might speed up your data loading.