Object Detection에서는 Train/Valid set을 어떻게 나눌까? -Stratified Group KFold (feat. sklearn)

CV를 나누기 어려운 Object Detection의 특징

Object Detection은 이미지 내에서 Object의 위치를 찾고(Localization), 어떤 종류인지 분류(Classification)하는 태스크입니다. 위의 슬라이드처럼 한 이미지에 한 개의 Object가 있을 수도, 여러 개의 Object가 있을 수도 있습니다. 하지만 후자와 같은 이미지가 당연히 세상에는 더 많겠죠?

한 이미지에 여러 개의 Object가 있다는 말은 곧 1개의 데이터에 여러 개의 label이 붙는다라고도 해석해볼 수 있습니다. 이런 경우에는 일반적으로 사용되던 CV 방법들을 적용하기가 어려워집니다. Random split 혹은 이미지 단위로 Train set과 Validation set을 나누게 될 경우, 특정한 Object가 Validation set에 아예 포함되지 않거나, 지나치게 많게 혹은 적게 포함되는 경우가 발생할 수 있습니다. 이에 대한 해결책으로는 Bouding box의 라벨을 고려해 split 하는 방법이 있을 수 있습니다. 즉, Train set과 Validation set의 bbox 라벨 분포를 맞춰주는 것이죠. 하지만 이 경우에는 Train set과 Validation set에 같은 이미지가 포함되는 문제가 발생할 수 있습니다. 이미 학습 때 봤던 이미지이기 때문에, 당연히 처음 보는 이미지에 비해 정답을 잘 맞힐 수밖에 없고, Validation 성적이 실제보다 좋게 나올 가능성이 높아지게 됩니다.따라서 이런 경우에는 Train set과 Validation set에 중복되는 이미지가 없으면서, 두 세트의 bbox 라벨 분포가 어느정도 유사하도록 맞춰줘야 합니다.

 

Stratified Group K Fold

Stratified Group K Fold는 Stratified K Fold와 Group K Fold가 합쳐진 말입니다. 말 그대로, 타겟 데이터의 분포를 유지하며 Train/Valid set을 나누는 Stratified K Fold와 동일한 그룹(= 같은 이미지)에서 나온 데이터가 Train set과 Valid set에 동시에 포함되지 않도록 하는 Group K Fold를 결합한 방법입니다.

https://scikit-learn.org/stable/auto_examples/model_selection/plot_cv_indices.html#sphx-glr-auto-examples-model-selection-plot-cv-indices-py

 

Sklearn - Stratified Group K Fold

벌써 코딩을 어떻게 해야하나... 머리가 지끈하지 않나요? ㅎㅎ 그런데 최근에 sklearn에 Stratified Group K Fold를 처리할 수 있는 모듈이 추가 되었답니다. 🎉

필요한 파라미터는 X, y, groups로  간단합니다. X에는 학습 데이터를, y에 label을, groups에는 어떤 그룹(= 우리의 경우 이미지)에서 나왔는지 라벨 데이터를 넣어주면 됩니다.

제가 가지고 있던 데이터셋은 COCO format의 json 파일로 annotation이 주어져서 아래와 같이 코드를 짰어요. 

import json
import numpy as np
from sklearn.model_selection import StratifiedGroupKFold

# load json
annotation = {dataset file 경로}

with open(annotation) as f:
    data = json.load(f)
    
var = [(ann['image_id'], ann['category_id']) for ann in data['annotations']]

X = np.ones((len(data['annotations']),1))
y = np.array([v[1] for v in var])
groups = np.array([v[0] for v in var])

cv = StratifiedGroupKFold(n_splits=5, shuffle=True, random_state=411)

for train_idx, val_idx in cv.split(X, y, groups):
    print("TRAIN:", groups[train_idx])
    print("      ", y[train_idx])
    print(" TEST:", groups[val_idx])
    print("      ", y[val_idx])

위 이미지와 코드에서 알 수 있듯이 우리는 train, valid index를 얻게 됩니다. 주의해야 할 점은 이 index는 annotation에 대한 index라는 점입니다. (Group index 아니에요!)

그럼 실제로 데이터가 잘 나누어졌는지 확인해보고 글을 마치도록 하겠습니다. (아래 코드는 캐글 노트북을 참고하였습니다. )

# check distribution
def get_distribution(y):
    y_distr = Counter(y)
    y_vals_sum = sum(y_distr.values())
    
    return [f'{y_distr[i]/y_vals_sum:.2%}'  for i in range(np.max(y) +1)]
    
distrs = [get_distribution(y)]
index = ['training set']

for fold_ind, (train_idx, val_idx) in enumerate(cv.split(X,y, groups)):
    train_y, val_y = y[train_idx], y[val_idx]
    # train_gr, val_gr = groups[train_idx], groups[val_idx]
    
    assert len(set(train_gr) & set(val_gr)) == 0
    
    distrs.append(get_distribution(train_y))
    distrs.append(get_distribution(val_y))
    
    index.append(f'train - fold{fold_ind}')
    index.append(f'val - fold{fold_ind}')
                 
categories = [d['name'] for d in data['categories']]

pd.DataFrame(distrs, index=index, columns = [categories[i] for i in range(np.max(y) + 1)])

 

위 코드를 실행해본 결과, 다음과 같이 각 Fold가 Training set의 데이터 분포와 유사한 분포를 가지고 있는 것을 확인할 수 있었습니다.