Splitting ImageFolder into data and labels?

Hello all, I’m new here and sorry if this question has been asked before, I have tried searching for a long time now and cannot seem to find the solution.

When I load a standard dataset from pytorch such as CIFAR, I can access the images and labels separately with dataset.data and dataset.targets. When I create a dataset with ImageFolder however, I only get back one list that does not contain .data or .label fields. I know that I can access the images and labels separately by doing dataset[i][0] or [1] to get respectively the ith image or label. But this list structure is bothersome because a) I cannot take the mean or std of my training set easily and b) I cannot pass the images and labels separately to stratifiedshufflesplit to create a validation set. (I know that I can use random_split but that does not guarantee equal distribution of classes in the validation set).

The way that I was using was to do:

images = []
labels = []
for i in range(len(my_data)):
images.append(my_data[i][0])
labels.append(my_data[i][1])

Which would give me separate lists of just the data or labels that I could then turn into a tensor with tensor.stack() for the data and a numpy array for the labels, both much nicer to work with.

This method works fine but uses up a lot of memory and is very slow. When trying to run it on Kaggle, 90% of the time my notebook restarts due to trying to use too much memory.

Can someone help me divide up my imagefolder into data and labels in a neater/more efficient way? Or am I not supposed to try and do that? Is there something conceptually wrong with what I am attempting here?

I also considered editing ImageFolder itself to have .data and .targets fields but I don’t know how I would go about doing that.

Could you please explain about your folder structure? data, labels, …

It goes like this:

/root/train/[class names]/[images]
/root/test/[only one class]/[testing images]

Since my test data is not labeled, I’m only reading the train folder into ImageFolder and trying to split that into a training and validation set

I usually use a custom dataset class for this type of problems and in the __init__ method read the name of images and do sorting and other stuff if needed and then in the __getitem__ method, load the image and do the processing and return the image and its label. You can have a lot of freedom and can control everything in the class. You can check this tutorial for more details.