Albumentations blocks CPU by spawning multiple threads

[Levin-Kobelke]

🐛 Bug/Unexpected behavior

First of all, I am not sure if this is considered a bug, but since it caused trouble for me and I could not find any easy help, I want to leave a hint for others that run into the same problem.

Problem: When using albumentations as transforms in a data loader, each data loading worker spawns additional threads. When training on a multi GPU Node this causes the shared CPU to be clogged by the albumentation processes and takes up to much CPU for only using part of the GPU.

To Reproduce

Steps to reproduce the behavior:

1. Build a dataloader with albumentation transforms if self.transform is not None: res = self.transform(image=image)
2. Use multiple workers to load the data onto the GPU
3. Observe each worker spawn multiple threads to perform augmentations (e.g. $htop ->Tree view)

Expected behavior

Some control parameter or hint of this behavior.

Environment

• Albumentations version (1.1.0):
• Python version (e.g., 3.10.4):
• OS (Ubuntu):
• How you installed albumentations ( pip):
• pythorch lightning 1.6.0
• torch1.11.0+cu113

Solution

Set cv2.setNumThreads(0) to fix the number of threads spawned

[MeteorsHub]

I observed the same problem. It consumed most of my cpu usage and the training is very slow. The cv2 trick worked for me

[zakajd]

+1 here. Same problem fixed with cv2.setNumThreads(0). I didn't notice such behaviour when using Albumentations 6 month ago (~June-July). Did something changes this then?

[Dipet]

This is known issue with OpenCV: https://github.com/albumentations-team/albumentations#comments I can not imagine common fix.