split custom pytorch dataset

1import torch
2import numpy as np
3from torchvision import datasets
4from torchvision import transforms
5from torch.utils.data.sampler import SubsetRandomSampler
6
7class CustomDatasetFromCSV(Dataset):
8    def __init__(self, csv_path, transform=None):
9        ....
10
11
12dataset = CustomDatasetFromCSV(my_path)
13batch_size = 16
14validation_split = .2
15shuffle_dataset = True
16random_seed= 42
17
18# Creating data indices for training and validation splits:
19dataset_size = len(dataset)
20indices = list(range(dataset_size))
21split = int(np.floor(validation_split * dataset_size))
22if shuffle_dataset :
23    np.random.seed(random_seed)
24    np.random.shuffle(indices)
25train_indices, val_indices = indices[split:], indices[:split]
26
27# Creating PT data samplers and loaders:
28train_sampler = SubsetRandomSampler(train_indices)
29valid_sampler = SubsetRandomSampler(val_indices)
30
31train_loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, 
32                                           sampler=train_sampler)
33validation_loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size,
34                                                sampler=valid_sampler)
35
36# Usage Example:
37num_epochs = 10
38for epoch in range(num_epochs):
39    # Train:   
40    for batch_index, (faces, labels) in enumerate(train_loader):
41        # ...
42