train, valid, test split with random sampling

Stratified Sampling in Pandas - GeeksforGeeks

 

10개의 비디오에서 이미지를 추출하여 car, plane, tree로 labeling된 아래와 같은 dataframe이 있을때,

• 비디오 번호 별 sampling
• Label의 비율을 유지하는 proportionate sampling

을 정리한다.

 

 

1. 비디오 번호 별 sampling

1-1. 파이썬 random 라이브러리를 이용한 sampling

import random

list_video = df['video'].unique()
print(list_video)

출력 결과

random.seed(42)
random.shuffle(list_video)

list_train = list_video[:int(len(list_video)*0.8)]
list_valid = list_video[int(len(list_video)*0.8):int(len(list_video)*0.9)]
list_test = list_video[int(len(list_video)*0.9):]

print(list_train)
print(list_valid)
print(list_test)

출력 결과

 

1-2. 비디오 이름 별로 train, valid, test dataframe 필터링하기

train = df[df.video.isin(list_train)].copy()
valid = df[df.video.isin(list_valid)].copy()
test = df[df.video.isin(list_test)].copy()

print(train['label'].value_counts())
print(valid['label'].value_counts())
print(test['label'].value_counts())

출력 결과

왼쪽부터 train, valid, test

 

2. Label의 비율을 유지하는 proportionate sampling (train:valid:test = 80:10:10)

2-1. valid와 test를 함께 추출

test_valid = df.groupby('label', group_keys = False).apply(lambda x:x.sample(frac=0.2, random_state = 42))

 

2-2. 함께 추출된 dataframe에서 valid와 test를 나눔

(valid를 먼저 sampling 한 뒤, valid에 포함되지 않은 index를 추출하여 test로 정의)

valid = test_valid.groupby('label', group_keys = False).apply(lambda x:x.sample(frac=0.5, random_state = 42))

test = test_valid[~test_valid.index.isin(valid.index)].copy()

 

2-3. valid와 test가 아닌 나머지를 train으로 추출

train = df[~df.index.isin(test_valid.index)].copy()

 

2-4. 추출 결과 확인

print(train['label'].value_counts())
print(valid['label'].value_counts())
print(test['label'].value_counts())

출력 결과

왼쪽부터 train, valid, test

 

끝