dataset.py

import linecache
import subprocess

import torch
from torch.utils.data import Dataset, DataLoader


class EEGDataset(Dataset):
    def __init__(self, filename):
        self.filename = filename
        self.total_size = int(subprocess.check_output("wc -l " + filename, shell=True).split()[0])

    def __getitem__(self, idx):
        line = linecache.getline(self.filename, idx + 1)
        str_x, str_y = line.split("\t")
        str_x = str_x.split("")
        y = int(str_y)
        x = torch.tensor([float(x) for x in str_x],
                           dtype=torch.float)
        x = x.view(2500, 2, 3)
        return x, y

    def __len__(self):
        return self.total_size


if __name__ == "__main__":
    dataset = EEGDataset("./data/eeg.txt")
    dataloader = DataLoader(dataset, batch_size=32, shuffle=False)
    for batch in dataloader:
        x, y = batch
        print(x.size(), y.size())