How to print data split into train- and test-data sets in PyTorch?

Yeongkyu Lee 0
Viewed 25

I have a data set loaded by pandas. My data set looks like:

hour score
0.5 10.0
1.2 8.0
1.8 14.0
2.4 26.0
2.6 22.0
3.2 30.0

I wrote the code shown below:

data = pd.read_csv("file_name")
X = data["hour"][:] ; Y = data["score"][:]
train_ratio = 0.8 ; test_ratio = 0.2
Ndata = len(X)
NumTrainData = int(Ndata*train_ratio)
NumTestData = int(Ndata*test_ratio)

train_set, test_set = torch.utils.data.random_split(data, [NumTrainData, NumTestData])

I want to check which data are stored in train_set and which in test_set. How can I check them?

pytorch
1 Answers

torch.utils.data.random_split returns torch.utils.data.Dataset object (official doc).

These object are iterators and can be browsed simply with a for loop:

def show_dataset(dataset: torch.utils.data.Dataset, num_elem: int = 5):
    for i, sample in enumerate(dataset):
        print(sample)
        if i == num_elem:
            break

Now you can display your top elements:

print('train dataset:')
show_dataset(train_set)
print('test dataset:')
show_dataset(test_set)
Valentin Goldité0
Related