1 回答

TA贡献2041条经验 获得超4个赞
我终于设法处理了 NoneType 问题。保留问题中定义的函数。
class MNIST01(MNIST):
def __getitem__(self, idx):
features, target = super(MNIST01, self).__getitem__(idx)
if target.item() <= 1:
return features, target
我们现在需要为我们的数据加载器定义一个自定义整理函数 collate_fn,它处理样本列表以形成一个批次。在这个函数中,我们可以应用过滤器来处理None值并忽略它们。
from torch.utils.data.dataloader import default_collate
def filter_collate(batch):
batch = list(filter(lambda x: x is not None, batch))
return default_collate(batch)
然后我们只需要将这个函数传递给DataLoader:
from torch.utils.data import DataLoader
train_loader = DataLoader(train_dataset, collate_fn=filter_collate, **kwargs)
test_loader = DataLoader(test_dataset, collate_fn=filter_collate, **kwargs)
版本 2
比第一个容易多了,避免了访问数据时的一些问题。只需从类的实例中直接过滤train_data和train_label属性(以及对应于测试集)MNIST。
train_dataset.train_data = train_dataset.train_data[train_dataset.train_labels <= 1]
train_dataset.train_labels = train_dataset.train_labels[train_dataset.train_labels <= 1]
添加回答
举报