You can pass in a list of file paths into scikit-learn's train_test_split
.
Here is an example:
from pathlib import Path
from sklearn.model_selection import train_test_split
file_locations = [Path("images1.png"), Path("images2.png")]
X_train, X_test = train_test_split(file_locations)
The choice between loading data all-at-once vs batch-by-batch is context dependent. If you have small amount of data relative to RAM, then loading in the data into memory is preferable because it will be faster and easier to handle. If you do not have enough RAM available, then taking the pointer-to-the-data approach will enable training to happen. The training in this scenario might be slower and require custom code to load the data at the moment it is needed.