Hi,
I’ve used PyTorch for my project. I run it in my GTX Nvidia with 4 GB of memory. I have a customized data loader, like the following code:
class LoadDataset(data.Dataset):

    def __init__(self, root, type_data):
        self.filenames = []
        self.root = root
        self.transform = TRANSFORM_IMG[type_data]
        self.classes = None
        for (path, dirs, files) in os.walk(self.root):
            self.classes = dirs
            break
        
        for c in self.classes:
            c_dir = os.path.join(self.root, c)
            walk = os.walk(c_dir).__next__()
            for k in walk[2]:
                if k.endswith('.png'):
                    self.filenames.append(os.path.join(c_dir, k))
 
        self.len = len(self.filenames)

    def __len__(self):
        return self.len

    def __getitem__(self, index):
        # Select sample
        image = Image.open(self.filenames[index]).convert("RGB")
        image = self.transform(image)
        label = int(self.filenames[index].split("\\")[1])
        path = self.filenames[index]
        sample = {'image': image, 'path': path, 'label': label}

        return sample

And actually, the format of my data is also available for using ImageLoader method.
When I tried with the following image transformation:
TRANSFORM_IMG = {
    'train':
        transforms.Compose([
                transforms.Resize(224),
                transforms.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4),
                transforms.RandomAffine(degrees=0, translate=(0, 0.001)),
                transforms.ToTensor(),
                transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                     std=[0.229, 0.224, 0.225])]),
    'validation':
        transforms.Compose([
                transforms.Resize(224),
                transforms.ToTensor(),
                transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                     std=[0.229, 0.224, 0.225])])
    }

My customized loader gives me a dump about Out Of Memory when I run on my GPU. However, the same method is performed with ImageLoader gives the process run without any complaining and I can finish my training.
So, does the ImageLoader has some optimization for consuming resources rather than my customized method? If I’m to use my customized method, what should I do to avoid OOM?
Thank you in advance.