title: Scikit learn train_test_split into Pytorch Dataloader
body: <p>I have a dataset for binary classification with PNGs titled as in the attachment below, where the first 0 or 1 in the title determines its class. They're in a folder called &quot;annotation_class&quot;, and I have a small script to separate these:</p>
  <pre class="lang-py prettyprint-override"><code>import cv2,glob
  import numpy as np
  from sklearn.model_selection import train_test_split
  
  filelist = glob.glob('annotation_class'+'/*.png')
  size_row, size_col = 256, 256
  X,y = [],[]
  
  for name in filelist:
          img = cv2.imread(name) 
          img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) 
          img = cv2.resize(img,(size_row, size_col))
          X.append(img)
          y.append(int(name.split('\\')[-1].split('_')[1]))
  
  
  x_train, x_test, y_train, y_test= train_test_split(X, y, test_size=0.2, train_size=0.8, random_state=4)
  </code></pre>
  <p>The returns are all lists. I'm using Pytorch for this project and would like to make a custom Dataset to use Dataloader, but I'm not sure how best to include these after I've used <code>train_test_split</code>. Should I scrap that altogether and use something else? I'd like to end up with two DataLoader's for training and testing.</p>
  <p><a href="https://i.stack.imgur.com/mnFtE.png" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/mnFtE.png" alt="sample titles" /></a></p>