使用python实现对样本的分层均衡抽样

前言

• 文章来源：CSDN@LawsonAbs
• 代码见我Github

---

1. 需求

在深度学习中，我们时常会碰到样本不均衡的情况，这种情况下，如果我们的样本没有给出dev集合，那么就需要手动分割train得到一份dev，但是dev的分割最好不要随机实现，而是要按照类别的个数分类取样。众所周知，sklearn中实现了常用的机器学习方法，但是 sklearn 中没有相应的函数实现此需求，【只有一个叫做 StratifiedKFold 的类，它只能做到尽可能按照类别取数，但还是有遗漏】
所以我们自己实现一份分层抽样算法，实现如下。

2.代码

def split_data_by_class(x,y,rate,seed=22):
    '''
    x,y 表示训练数据的输入和标签
    如果random = True, 则随机分割~
    rate表示dev所占数据的比例，如果不足1条，则按1条处理

    returns:
        train_idx,dev_idx
    '''    
    # 先shuffle一下再说，保持相同的shuffle  
    random.seed(seed)
    random.shuffle(y)
    random.seed(seed)
    random.shuffle(x)

    cont_id = {} # 每个类别都放到一个list中    
    # 保持每个类别划分相同
    for i in range(len(y)):
        y_idx = y[i] # y_idx是个分类值
        if y_idx not in cont_id.keys():
            cont_id[y_idx] = []
        cont_id[y_idx].append(i) # 将该类别的下标放到其中
    
    train_idx,dev_idx = [],[] # 最后返回的值，是一个下标
    # 保持按照类别均匀抽取，一共有35个类
    for item in cont_id.items():
        key,value = item # key 是类别信息，value是该类别的所有（在y中）下标
        if len(value) >= 2: # 如果当前类别的样本数大于2，才分成dev
            mid = int(len(value)*rate)
            if mid == 0: # 如果只有0，那么就得将其改成1
                mid = 1
            for i in range(mid): # 按照个数取前 rate * len(value) 个
                dev_idx.append(value[i])         
            for i in range(mid,len(value)): # 从上次的开始，作为train 集合
                train_idx.append(value[i])
        else:            
             train_idx.extend(value) # 将整个值都作为train
    
    # 得到值后再次shuffle，避免输出连续相同的值
    random.seed(seed)
    random.shuffle(train_idx)
    random.seed(seed)
    random.shuffle(dev_idx)
    
    # 需要判断一下最后返回的dev_idx 的长度是否符合要求，如果不合要求，则要随机再抽取一点儿到其中
    while len(dev_idx) < len(train_idx) * rate:
        idx = random.choice(train_idx)        
        dev_idx.append(idx) # 放到dev中
        train_idx.remove(idx) # 删除这个idx

    return train_idx,dev_idx # 返回train/dev 数据所在的下标