title: Split dataset based on file names in pytorch Dataset
body: <p>Is there a way to divide the dataset into training and testing based on the filenames. I have a folder containing two folders: input and output. Input folder has the images and output are the labels for that image. The file names in the input folder are something like <code>input01_train.png</code> and <code>input01_test.png</code> like shown below.</p>
  
  <pre><code>                          Dataset
                            /     \
                       Input       Output
                        |             |
             input01_train.png   output01_train.png
                      .                 .
                      .                 .
             input01_test.png    output01_test.png
  </code></pre>
  
  <p>The code I have only divides the dataset into inputs and labels not test and train.</p>
  
  <pre><code>class CancerDataset(Dataset):
    def __init__(self, dataset_folder):#,label_folder):
      self.dataset_folder = torchvision.datasets.ImageFolder(dataset_folder ,transform = transforms.Compose([transforms.Resize(512),transforms.ToTensor()]))
      self.label_folder = torchvision.datasets.ImageFolder(dataset_folder ,transform = transforms.Compose([transforms.Resize(512),transforms.ToTensor()]))
  
    def __getitem__(self,index):
      img = self.dataset_folder[index]
      label = self.label_folder[index]
      return img,label
  
    def __len__(self):
      return len(self.dataset_folder)
  
  trainset = CancerDataset(dataset_folder = '/content/drive/My Drive/cancer_data/')
  trainsetloader = DataLoader(trainset,batch_size = 1, shuffle = True,num_workers = 0,pin_memory = True)
  </code></pre>
  
  <p>I would like to be able to divide the train and test set by their names if that is possible .</p>