创建一个 pyTorch 测试数据集(不带标签)

时间:2021-02-24 17:37:38

标签: python pytorch pytorch-dataloader

我为我的训练数据创建了一个 pyTorch 数据集,其中包含特征和标签,以便能够使用 this 教程使用 pyTorch DataLoader。这对我的训练数据很有效,但在加载测试 csv 文件时出现错误 (KeyError: "['label'] not found in axis"),该文件除了没有“标签”列之外是相同的。

如果有帮助,预期的输入 csv 文件是 csv 文件中的 MNIST 数据,其中包含 28*28 个特征列。

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'])

1 个答案:

答案 0 :(得分:0)

您应该能够使用这两种数据:

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'])