当前的答案进行随机分裂,其缺点是不能保证每个类别的样本数量保持平衡。当您希望每个类别的样本数量较少时,这尤其成问题。例如,MNIST有60,000个示例,即每位数6000个。假设您只希望训练集中的每位数字有30个示例。在这种情况下,随机拆分可能会在各班级之间产生不平衡(一个位数比其他位数更多的训练数据)。因此,您要确保每个数字恰好只有30个标签。这称为分层抽样 。
一种方法是使用Pytorch中的sampler接口,此处提供示例代码 。
做到这一点的另一种方法就是通过:)破解自己的方法。例如,以下是MNIST的简单实现,其中ds
是MNIST数据集, k
是每个类别所需的样本数。
def sampleFromClass(ds, k):
class_counts = {}
train_data = []
train_label = []
test_data = []
test_label = []
for data, label in ds:
c = label.item()
class_counts[c] = class_counts.get(c, 0) + 1
if class_counts[c] <= k:
train_data.append(data)
train_label.append(torch.unsqueeze(label, 0))
else:
test_data.append(data)
test_label.append(torch.unsqueeze(label, 0))
train_data = torch.cat(train_data)
for ll in train_label:
print(ll)
train_label = torch.cat(train_label)
test_data = torch.cat(test_data)
test_label = torch.cat(test_label)
return (TensorDataset(train_data, train_label),
TensorDataset(test_data, test_label))
您可以这样使用此功能:
def main():
train_ds = datasets.MNIST('../data', train=True, download=True,
transform=transforms.Compose([
transforms.ToTensor()
]))
train_ds, test_ds = sampleFromClass(train_ds, 3)
2
通过使用其他存储库中的引用,我可以设法做到这一点。但是,我想将此数据集拆分为训练和测试。
在课堂上我该怎么做?还是我需要做一个单独的课堂来做到这一点?