PyTorch CSV

pytorch_csv.py

from torch.utils.data import Dataset, DataLoader

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


class CSVDataset(Dataset):
    def __init__(self, x, y):
        self.data = torch.tensor(x, dtype=torch.float, device=device)
        self.target = torch.tensor(y, dtype=torch.long, device=device)

    def __getitem__(self, index):
        return self.data[index], self.target[index]

    def __len__(self):
        return self.data.shape[0]


x = (pd.read_csv('./data/AI/x', sep='\t', header=None)).values
y = (pd.read_csv('./data/AI/y', sep='\t', header=None)[0]).values

X_train, X_test, y_train, y_test = train_test_split(x, y, test_size=0.2, random_state=0)

train_dataset = CSVDataset(X_train, y_train)
test_dataset = CSVDataset(X_test, y_test)