小编Kim*_*men的帖子

我是 pytorch 的新手，正在尝试实现一个前馈神经网络来对 mnist 数据集进行分类。我在尝试使用交叉验证时遇到了一些问题。我的数据具有以下形状 x_train：： torch.Size([45000, 784])和 y_train：torch.Size([45000])

我尝试使用 sklearn 中的 KFold。

kfold =KFold(n_splits=10)

这是我的训练方法的第一部分，我将数据分成几部分：

for  train_index, test_index in kfold.split(x_train, y_train): 
        x_train_fold = x_train[train_index]
        x_test_fold = x_test[test_index]
        y_train_fold = y_train[train_index]
        y_test_fold = y_test[test_index]
        print(x_train_fold.shape)
        for epoch in range(epochs):
         ...

y_train_fold变量的索引是正确的，它只是: [ 0 1 2 ... 4497 4498 4499]，但它不是 for x_train_fold，而是[ 4500 4501 4502 ... 44997 44998 44999]。测试折叠也是如此。

对于第一次迭代，我希望变量x_train_fold是前 4500 张图片，换句话说，具有 shape torch.Size([4500, 784])，但它具有 shapetorch.Size([40500, 784])

关于如何做到这一点的任何提示？