Multiple image folder

I am working on a classification problem. I split my dataset into sub-folders according to class labels. Then, I used ImageFolder and DataLoader from Pytorch to load the data. I have attached my directory structure below.


My dataset does not have a train and test folder. Now, I am trying to do splitting in a way that all four sub-folders get split into train and test folders.

Most of the solutions I came across have train and test folders pre-hand with respective .csv files. And use train_test split from sklearn.

Should I split data before splitting the entire data set it into image sub-folders based on class label?

It generally makes sense to split across class labels proportionally as splitting blindly across the entire dataset risks omitting smaller classes from either the training or testing set entirely.