在机器学习多分类任务中有时候需要针对类别进行分层采样，比如说类别不均衡的数据，这时候随机采样会造成训练集、验证集、测试集中不同类别的数据比例不一样，这是会在一定程度上影响分类器的性能的，这时候就需要进行分层采样保证训练集、验证集、测试集中每一个类别的数据比例差不多持平。
下面python代码。
# 将数据按照类别进行分层划分 def save_file_stratified(filename, ssdfile_dir, categories): """ 将文件分流到3个文件中 filename: 原数据地址，一个csv文件 文件内容格式: 类别\t内容 """ f_train = open('../data/usefuldata-711depart/train.txt', 'w', encoding='utf-8') f_val = open('../data/usefuldata-711depart/val.txt', 'w', encoding='utf-8') f_test = open('../data/usefuldata-711depart/test.txt', 'w', encoding='utf-8') # f_class = open('../data/usefuldata-37depart/class.txt', 'w', encoding='utf-8') dict_ssdqw = {} for ssdfile in os.listdir(ssdfile_dir): ssdfile_name = os.path.join(ssdfile_dir, ssdfile) f = open(ssdfile_name, 'r', encoding='utf-8') content_qw = '' content = f.readline() # 以下部分，因为统计整个案件基本情况他有换行，所以将多行处理在一行里面 while content: content_qw += content content_qw = content_qw.replace('\n', '') content = f.readline() ssdfile_key = str(ssdfile).replace('.txt','') dict_ssdqw[ssdfile_key] = content_qw # doc_count代表每一类数据总共有多少个 doc_count_0 = 0 doc_count_1 = 0 doc_count_2 = 0 doc_count_3 = 0 doc_count_4 = 0 doc_count_5 = 0 doc_count_6 = 0 doc_count_7 = 0 doc_count_8 = 0 doc_count_9 = 0 doc_count_10 = 0 doc_count_11 = 0 doc_count_12 = 0 temp_file = open(filename, 'r', encoding='utf-8') line = temp_file.readline() while line: line_content = line.split(',') name = line_content[0] if name in dict_ssdqw: label = line_content[1] if label == categories[0]: doc_count_0 += 1 elif label == categories[1]: doc_count_1 += 1 elif label == categories[2]: doc_count_2 += 1 elif label == categories[3]: doc_count_3 += 1 elif label == categories[4]: doc_count_4 += 1 elif label == categories[5]: doc_count_5 += 1 elif label == categories[6]: doc_count_6 += 1 elif label == categories[7]: doc_count_7 += 1 elif label == categories[8]: doc_count_8 += 1 elif label == categories[9]: doc_count_9 += 1 elif label == categories[10]: doc_count_10 += 1 elif label == categories[11]: doc_count_11 += 1 elif label == categories[12]: doc_count_12 += 1 line = temp_file.readline() temp_file.close() # 总数量 doc_count = doc_count_0 + doc_count_1 + doc_count_2 + doc_count_3 +\ doc_count_4 + doc_count_5 + doc_count_6 + doc_count_7 +\ doc_count_8 + doc_count_9 + doc_count_10 + doc_count_11 + doc_count_12 class_set = set() tag_train_0 = doc_count_0 * 70 / 100 tag_train_1 = doc_count_1 * 70 / 100 tag_train_2 = doc_count_2 * 70 / 100 tag_train_3 = doc_count_3 * 70 / 100 tag_train_4 = doc_count_4 * 70 / 100 tag_train_5 = doc_count_5 * 70 / 100 tag_train_6 = doc_count_6 * 70 / 100 tag_train_7 = doc_count_7 * 70 / 100 tag_train_8 = doc_count_8 * 70 / 100 tag_train_9 = doc_count_9 * 70 / 100 tag_train_10 = doc_count_10 * 70 / 100 tag_train_11= doc_count_11 * 70 / 100 tag_train_12 = doc_count_12 * 70 / 100 tag_val_0 = doc_count_0 * 85 / 100 tag_val_1 = doc_count_1 * 85 / 100 tag_val_2 = doc_count_2 * 85 / 100 tag_val_3 = doc_count_3 * 85 / 100 tag_val_4 = doc_count_4 * 85 / 100 tag_val_5 = doc_count_5 * 85 / 100 tag_val_6 = doc_count_6 * 85 / 100 tag_val_7 = doc_count_7 * 85 / 100 tag_val_8 = doc_count_8 * 85 / 100 tag_val_9 = doc_count_9 * 85 / 100 tag_val_10 = doc_count_10 * 85 / 100 tag_val_11 = doc_count_11 * 85 / 100 tag_val_12 = doc_count_12 * 85 / 100 # tag_test = doc_count * 70 / 100 tag_0 = 0 tag_1 = 0 tag_2 = 0 tag_3 = 0 tag_4 = 0 tag_5 = 0 tag_6 = 0 tag_7 = 0 tag_8 = 0 tag_9 = 0 tag_10 = 0 tag_11 = 0 tag_12 = 0 # 有些文书行业标记是空！！我想看看有多少条？ blank_tag = 0 # 标记一下，每个类别有多少个训练集、验证集、测试集？ train_class_tag = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] val_class_tag = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] test_class_tag = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] # csvfile = open(filename, 'r', encoding='utf-8') txtfile = open(filename, 'r', encoding='utf-8') process_line = txtfile.readline() while process_line: line_content = process_line.split(',') name = line_content[0] if name in dict_ssdqw: content = dict_ssdqw[name] label = line_content[1] # if label != '' and label != '其他行业': if label != '': class_set.add(label) # 对每一类进行分层采样 if label == categories[0]: tag_0 += 1 if tag_0 < tag_train_0: f_train.write(label + '\t' + content + '\n') train_class_tag[0] += 1 elif tag_0 < tag_val_0: f_val.write(label + '\t' + content + '\n') val_class_tag[0] += 1 else: f_test.write(label + '\t' + content + '\n') test_class_tag[0] += 1 elif label == categories[1]: tag_1 += 1 if tag_1 < tag_train_1: f_train.write(label + '\t' + content + '\n') train_class_tag[1] += 1 elif tag_1 < tag_val_1: f_val.write(label + '\t' + content + '\n') val_class_tag[1] += 1 else: f_test.write(label + '\t' + content + '\n') test_class_tag[1] += 1 elif label == categories[2]: tag_2 += 1 if tag_2 < tag_train_2: f_train.write(label + '\t' + content + '\n') train_class_tag[2] += 1 elif tag_2 < tag_val_2: f_val.write(label + '\t' + content + '\n') val_class_tag[2] += 1 else: f_test.write(label + '\t' + content + '\n') test_class_tag[2] += 1 elif label == categories[3]: tag_3 += 1 if tag_3 < tag_train_3: f_train.write(label + '\t' + content + '\n') train_class_tag[3] += 1 elif tag_3 < tag_val_3: f_val.write(label + '\t' + content + '\n') val_class_tag[3] += 1 else: f_test.write(label + '\t' + content + '\n') test_class_tag[3] += 1 elif label == categories[4]: tag_4 += 1 if tag_4 < tag_train_4: f_train.write(label + '\t' + content + '\n') train_class_tag[4] += 1 elif tag_4 < tag_val_4: f_val.write(label + '\t' + content + '\n') val_class_tag[4] += 1 else: f_test.write(label + '\t' + content + '\n') test_class_tag[4] += 1 elif label == categories[5]: tag_5 += 1 if tag_5 < tag_train_5: f_train.write(label + '\t' + content + '\n') train_class_tag[5] += 1 elif tag_5 < tag_val_5: f_val.write(label + '\t' + content + '\n') val_class_tag[5] += 1 else: f_test.write(label + '\t' + content + '\n') test_class_tag[5] += 1 elif label == categories[6]: tag_6 += 1 if tag_6 < tag_train_6: f_train.write(label + '\t' + content + '\n') train_class_tag[6] += 1 elif tag_6 < tag_val_6: f_val.write(label + '\t' + content + '\n') val_class_tag[6] += 1 else: f_test.write(label + '\t' + content + '\n') test_class_tag[6] += 1 elif label == categories[7]: tag_7 += 1 if tag_7 < tag_train_7: f_train.write(label + '\t' + content + '\n') train_class_tag[7] += 1 elif tag_7 < tag_val_7: f_val.write(label + '\t' + content + '\n') val_class_tag[7] += 1 else: f_test.write(label + '\t' + content + '\n') test_class_tag[7] += 1 elif label == categories[8]: tag_8 += 1 if tag_8 < tag_train_8: f_train.write(label + '\t' + content + '\n') train_class_tag[8] += 1 elif tag_8 < tag_val_8: f_val.write(label + '\t' + content + '\n') val_class_tag[8] += 1 else: f_test.write(label + '\t' + content + '\n') test_class_tag[8] += 1 elif label == categories[9]: tag_9 += 1 if tag_9 < tag_train_9: f_train.write(label + '\t' + content + '\n') train_class_tag[9] += 1 elif tag_9 < tag_val_9: f_val.write(label + '\t' + content + '\n') val_class_tag[9] += 1 else: f_test.write(label + '\t' + content + '\n') test_class_tag[9] += 1 elif label == categories[10]: tag_10 += 1 if tag_10 < tag_train_10: f_train.write(label + '\t' + content + '\n') train_class_tag[10] += 1 elif tag_10 < tag_val_10: f_val.write(label + '\t' + content + '\n') val_class_tag[10] += 1 else: f_test.write(label + '\t' + content + '\n') test_class_tag[10] += 1 elif label == categories[11]: tag_11 += 1 if tag_11 < tag_train_11: f_train.write(label + '\t' + content + '\n') train_class_tag[11] += 1 elif tag_11 < tag_val_11: f_val.write(label + '\t' + content + '\n') val_class_tag[11] += 1 else: f_test.write(label + '\t' + content + '\n') test_class_tag[11] += 1 elif label == categories[12]: tag_12 += 1 if tag_12 < tag_train_12: f_train.write(label + '\t' + content + '\n') train_class_tag[12] += 1 elif tag_12 < tag_val_12: f_val.write(label + '\t' + content + '\n') val_class_tag[12] += 1 else: f_test.write(label + '\t' + content + '\n') test_class_tag[12] += 1 else: blank_tag += 1 process_line = txtfile.readline() txtfile.close() print("有" + str(blank_tag) + "个文书的行业标记为空！") print("train:") print(train_class_tag) train_tag_total =0 for i_total in train_class_tag: train_tag_total += i_total train_class_tag_distribute = [] for i in train_class_tag: train_class_tag_distribute.append((i / train_tag_total) * 100) print("分布:") print(train_class_tag_distribute) print("val:") print(val_class_tag) val_tag_total = 0 for i_total in val_class_tag: val_tag_total += i_total val_class_tag_distribute = [] for i in val_class_tag: val_class_tag_distribute.append((i / val_tag_total) * 100) print("分布:") print(val_class_tag_distribute) print("test:") print(test_class_tag) test_tag_total = 0 for i_total in test_class_tag: test_tag_total += i_total test_class_tag_distribute = [] for i in test_class_tag: test_class_tag_distribute.append((i / test_tag_total) * 100) print("分布:") print(test_class_tag_distribute) f_train.close() f_test.close() f_val.close() if __name__ == '__main__': categories = [ "class1", "class2", "class3", "class4", "class5", "class6", "class7", "class8", "class9", "class10", "class11", "class12", "class13" ] save_file_stratified('../data/qwdata/shuffle-try3/classified_table_ms.txt', '../data/qwdata/ms-ygscplusssdqw',categories)
后面可以看到类别划分
这里要注意的一点是：这是我早期写的文章，需要注意的一点是，我们通常在训练集和验证集上做分层采样即可，测试集最好保持原样不要动。