Creating a custom Dataloader to load .mat files

Hello, Am a beginner in deep-learning, Am trying to do image holographic image reconstruction and i need help on creating a DataLoader to take into a CNN .mat images. my images are divided into 3 folders ie training, testing and Validation. So i need a hand in creating an algorithm to take in these 3 categories of files.

Thank you

This might be the dataset module:

import os
import hdf5storage as h5
import torch
from import DataLoader, Dataset
import torchvision.transforms

class DataSet(Dataset):
	def __init__(self, images_folder,transform):
		assert os.path.exists(images_folder), 'Images folder does not exist'
		self.images = os.listdir(self.images_folder)
		self.transform = transform

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

	def __getitem__(self, index):
		img_filename = self.images[index]
		img = h5.loadmat(img_filename)['Insert_your_array_name']
		if self.transform:
			img = torch.from_numpy(img)	
		return img

if __name__=='__main__':

	image				[Training Images]
	image_test			[Testing Images]	
	image_validation	[Validation Images]

	dataset_train = DataSet(os.path.join(data_path, 'image'), transform)
	dataset_test = DataSet(os.path.join(data_path, 'image_test'), transform)
	dataset_val = DataSet(os.path.join(data_path, 'image_validation'), transform)

	train_loader = DataLoader(dataset_train, batch_size=batch_size, shuffle=True)
    test_loader = DataLoader(dataset_test, batch_size=batch_size, shuffle=True)
    val_loader= DataLoader(dataset_val, batch_size=batch_size, shuffle=True)

Does that solve your problem?

Thank you Hmrishav,

Your dataLoader works well for me