Create a pyTorch testing Dataset (without labels)

user_3042868 0
Viewed 1591

I have created a pyTorch dataset for my training data which consists of features and a label to be able to utilize the pyTorch DataLoader using this tutorial. This works well for my training data, but I get an error (KeyError: "['label'] not found in axis") when loading the testing csv file, which is identical other than there being no "label" column.

If it helps, the intended input csv file is MNIST data in csv file which has 28*28 feature columns.

import torch

class mnist(torch.utils.data.Dataset):
    
    def __init__(self, csv_file):
        self.train = pd.read_csv(csv_file)
        self.train_x = self.train.drop("label", axis=1)
    
    def __len__(self):
        return len(self.train)
    
    def __getitem__(self, idx):
        if torch.is_tensor(idx):
            idx = idx.tolist()
        if isinstance(idx, list):
            idx_len = len(idx)
        else:
            idx_len = 1
        
        X = np.asarray(self.train_x.iloc[idx], dtype=np.float32)
        X = np.reshape(X, (1,28,28))
        y = np.asarray(self.train.iloc[idx]['label'])
        
        sample = {'X': X, 'y':y}
        
        return torch.from_numpy(sample['X']), torch.from_numpy(sample['y'])
pytorch-dataloader pytorch python
1 Answers

You should be able to use both data with this:

import torch

class mnist(torch.utils.data.Dataset):
    
    def __init__(self, csv_file):
        self.train = pd.read_csv(csv_file)

        self.training = "label" in self.train.columns
        self.train_x = self.train if not self.training else self.train.drop("label", axis=1)
    
    def __len__(self):
        return len(self.train)
    
    def __getitem__(self, idx):
        ...
        
        X = np.asarray(self.train_x.iloc[idx], dtype=np.float32)
        X = np.reshape(X, (1,28,28))
        if not self.training:
            return torch.from_numpy(X])

        y = np.asarray(self.train.iloc[idx]['label'])

        sample = {'X': X, 'y':y}
        return torch.from_numpy(sample['X']), torch.from_numpy(sample['y'])
user_80475350
Related