multi-columns: stratify split, 여러 개의 클래스 stratify로 데이터셋 분리하기

skmultilearn 이용

from skmultilearn.model_selection import iterative_train_test_split

%%time
X = straw['image'].to_numpy().reshape(-1,1)
Y = straw[['disease', 'area', 'grow']].to_numpy()
X_train, y_train, X_test, y_test = iterative_train_test_split(X,Y,test_size=0.20)

---

y_train = y_train.reshape(3,-1)
y_df = pd.DataFrame({'disease' : y_train[0,:],
                    'area':y_train[1,:],
                    'grow':y_train[2,:]})
for col in y_df.columns :
    df = y_df.groupby(col).size()
    display(df)

###### 출력결과 ######

disease
0    5731
1    2838
2    2896
3    4009
4    1358
5     360
dtype: int64

area
0    5731
1    2862
2    2483
3    6042
4      72
5       2
dtype: int64

grow
0    4281
1    1014
2    4614
3    3740
4     611
5    1483
7     601
8     848
dtype: int64

---

y_test = y_test.reshape(3, -1)
y_df = pd.DataFrame({'disease' : y_test[0,:],
                    'area':y_test[1,:],
                    'grow':y_test[2,:]})
for col in y_df.columns :
    df = y_df.groupby(col).size()
    display(df)

###### 출력결과 ######

disease
0    1433
1    1433
2       1
3       5
4    1113
5     313
dtype: int64

area
0    1433
1     283
2    1154
3    1143
4     274
5      11
dtype: int64

grow
0    1070
1      53
2    1416
3    1093
4     300
5       4
7     354
8       8
dtype: int64