sklearn使用StratifiedShuffleSplit完成train:test:dev = 6:2:2的数据集分割

import numpy as np
from sklearn.model_selection import StratifiedShuffleSplit

X = []
y = []
num = 0
with open('./bankV4.txt', 'r', encoding='utf8') as v4:
    old_lines = v4.readlines()
    print(len(old_lines))
    for i in range(len(old_lines)):
        if i > 0:
            #print(old_lines[i][0])
            y.append(int(old_lines[i][0]))
    #print(y)
            X.append(old_lines[i][1:])

split = StratifiedShuffleSplit(n_splits=1, test_size=0.4, random_state=0)
for train_index, test_index in split.split(X, y):
    print("TRAIN:", train_index, "\nTEST:", test_index)
    X_train, X_test = np.array(X)[train_index], np.array(X)[test_index]
    y_train, y_test = np.array(y)[train_index], np.array(y)[test_index]

X_train = X_train.tolist()
X_test = X_test.tolist()
#   X_dev = X_dev.tolist()
y_train = y_train.tolist()
y_test = y_test.tolist()
#    y_dev = y_dev.tolist()
train = []
for i in range(len(X_train)):
        train.append(str(y_train[i]) + str(X_train[i]))

split1 = StratifiedShuffleSplit(n_splits=1, test_size=0.5, random_state=0)
for test1_index, dev_index in split1.split(X_test, y_test):
    print("TRAIN:", test1_index, "\nTEST:", dev_index)
    X_test_test1, X_test_dev = np.array(X_test)[test1_index], np.array(X_test)[dev_index]
    y_test_test1, y_test_dev = np.array(y_test)[test1_index], np.array(y_test)[dev_index]

X_test_test1 = X_test_test1.tolist()
X_test_dev = X_test_dev.tolist()
y_test_test1 = y_test_test1.tolist()
y_test_dev = y_test_dev.tolist()

test = []
dev = []
for i in range(len(X_test_test1)):
        test.append(str(y_test_test1[i]) + str(X_test_test1[i]))
for i in range(len(X_test_dev)):
        dev.append(str(y_test_dev[i]) + str(X_test_dev[i]))

print("length:", len(train), len(test), len(dev))

#print(X_test)
#print(type(X_train), type(X_test), type(y_train), type(y_test))
with open('./train_bank.txt', 'w', encoding='utf8') as output_file1:
    output_file1.writelines(train)
with open('./test_bank.txt', 'w', encoding='utf8') as output_file2:
    output_file2.writelines(test)
with open('./dev_bank.txt', 'w', encoding='utf8') as output_file3:
    output_file3.writelines(dev)