[人工智能]sklearn划分训练集和测试集

from sklearn.model_selection import train_test_split
import os
 
pth = 'E:\\LONGBO\\crops_images\\train_data\\image\\'
files = [img for img in os.listdir(pth) if img.endswith('jpg')]
 
train, test = train_test_split(files, test_size=0.20, random_state=2020)
 
with open('valid.txt', 'w') as f:
    for line in test:
        f.write(pth+line+'\n')
 
with open('train.txt', 'w') as f:
    for line in train:
        f.write(pth+line+'\n')