How to use pillow_heif in Dataloader (Pytorch)

[qilong-zhang]

Hi @bigcat88 , I want to load ".heic" images in Dataloader. So my code is

Class Dataset(torch.utils.data.Dataset):
    def __init__(self, path=None, mode='train'):
          pass
    def __len__(self):
          pass
    def __getitem__(self, index):
          register_heif_opener()
          pass

But this code seems to cause a memory leak. Is there any other better solution?

[bigcat88]

Can you provide fuller example?

1. register_heif_opener can be called only once somewhere when program starts
2. very strange that register_heif_opener causes a memory leak. What version of Pillow are you using?

[qilong-zhang]

Can you provide fuller example?

1. register_heif_opener can be called only once somewhere when program starts
2. very strange that register_heif_opener causes a memory leak. What version of Pillow are you using?

from multiprocessing import shared_memory
import glob
from pillow_heif import register_heif_opener
import torch

class Dataset_Test(torch.utils.data.Dataset):                                    
    def __init__(self, path=None):
        self.transform = ToTensor()
        self.common_path = glob.glob(os.path.join(path, '*', '*'))
        self.common_path = shared_memory.ShareableList(self.common_path)

    def __len__(self):
        return len(self.common_path)

    def __getitem__(self, index):        
        register_heif_opener()                              
        name = self.common_path[index]
        image = Image.open(name).convert('RGB')
        
        return self.transform(image)

Here is my complete code, the description just given about the memory leak may not be accurate. My main problem at the moment is that I find that if I call register_heif_opener() every time inside "_getitem_", it will load slower, is there any way to call it only once?

Besides, If I use the above code, after several iterations, the running code will terminate, and report:

Traceback (most recent call last):
  File "train.py", line 289, in <module>
    run_demo(main, n_gpus)
  File "train.py", line 204, in run_demo
    mp.spawn(demo_fn,
  File "/root/mambaforge3/envs/torch-gpu/lib/python3.8/site-packages/torch/multiprocessing/spawn.py", line 240, in spawn
    return start_processes(fn, args, nprocs, join, daemon, start_method='spawn')
  File "/root/mambaforge3/envs/torch-gpu/lib/python3.8/site-packages/torch/multiprocessing/spawn.py", line 198, in start_processes
    while not context.join():
  File "/root/mambaforge3/envs/torch-gpu/lib/python3.8/site-packages/torch/multiprocessing/spawn.py", line 140, in join
    raise ProcessExitedException(
torch.multiprocessing.spawn.ProcessExitedException: process 2 terminated with signal SIGKILL

[bigcat88]

You can call it before the definition of Dataset_Test class for example.

from multiprocessing import shared_memory
import glob
from pillow_heif import register_heif_opener
import torch

register_heif_opener() 

class Dataset_Test(torch.utils.data.Dataset):  

if it works, you can simplify it even more with one import line and without calling register_heif_opener:

from multiprocessing import shared_memory
import glob
import pillow_heif.HeifImagePlugin  # noqa
import torch

[qilong-zhang]

Thanks! My code aims to train a classfication model with pytorch DDP (4 GPUs). Therefore, it is not easy to provide a minimal example, as this also requires large-scale data sets. But one thing is for sure, if I don't consider heic format images, I won't have these problems.

Interestingly, it seems that the memory will no longer continue to rise slightly when the model continues to be trained after finishing one epoch training.

[qilong-zhang]

Thanks! My code aims to train a classfication model with pytorch DDP (4 GPUs). Therefore, it is not easy to provide a minimal example, as this also requires large-scale data sets.

Interestingly, it seems that the memory will no longer continue to rise slightly when the model continues to be trained after finishing one epoch training.

Not right, the memory still rise silightly..

[bigcat88]

Without an example in format:

for i in range(1000):
    image = Image.open("path_to_image").convert('RGB')
    image = self.transform(image)

i probably can not help.
Are you sure that this problem is not present without heic images?(saw that you mention that there is no problem when you do not open HEIF images)

P.S. (not related to current problem): in next Pillow(9.5.0) they will add support for multiply calls of register_heif_opener
python-pillow/Pillow#6916

[qilong-zhang]

I retest my code and I find that if I don't train the model, the memory doesn't grow; but if I train the model, the memory grows regardless of whether it's 1 GPU or multiple GPUs. Fortunately, memory grows relatively slowly, so for now I don't have SIGKILL as long as I have enough memory.

[qilong-zhang]

Without an example in format:

for i in range(1000):
    image = Image.open("path_to_image").convert('RGB')
    image = self.transform(image)

i probably can not help. Are you sure that this problem is not present without heic images?(saw that you mention that there is no problem when you do not open HEIF images)

P.S. (not related to current problem): in next Pillow(9.5.0) they will add support for multiply calls of register_heif_opener python-pillow/Pillow#6916

I checked my code again. I found that the memory growth was a problem that was already present in my code. If I consider .heic format images, the memory grows relatively faster, thus causing SIGKILL. So for now, it looks like the problem may not be fundamentally caused by .heic.