【python/丁寧解説】データセットをtrain/val/testにランダムで分ける方法

概要

データセットを準備した後にtrain,val,testにランダムに分けるのは意外と手間です。

scikit-learnのtrain_test_split()を使って少し工夫すると簡単にできます。

実装方法

まずは必要なライブラリのインポートとダミーデータを作成します。

ダミーデータは0から99の連番の数値となっています。

from sklearn.model_selection import train_test_split
data_list = list(range(0,100))

train_test_split()を2段階に使って実装します。

まずはtrain,val,testの割合を指定します。

train_rate = 0.75
val_rate = 0.15
test_rate = 0.10

まずは1段階目としてtrain(学習用)とval+test(評価とテスト用)に分けます。

trainの割合を基準として分けます。

train, val_test = train_test_split(data_list, train_size=train_rate, random_state=0)

次に2段階目としてval+testをvalとtestに分けます。

valとtestにおけるvalの割合は以下のようになります。

$$\frac{valの割合}{valの割合 + testの割合}$$

コードにすると以下のようになります。

val, test = train_test_split(val_test, train_size=val_rate / (val_rate + test_rate), random_state=0)

これでtrain,val,testに分けることができました。

念のため、ちゃんとデータがうまく振り分けられているか確認しました。

#学習、評価、テストを合体
data_list_merge = train + val + test
#データをソートする
data_list_merge.sort()
#元のリストと比較
data_list == data_list_merge

結果はtrueでした。OKですね。

実データでの確認

ダミーデータで確認しましたので、実データでも確認してみます。

今回はMNISTで試してみます。

まずはMNISTのデータを読み込みます。

tensorflowから読み込むとtrainとtestに分かれていますので一旦合体させて1つのデータとします。

#mnistデータを読み込む
from tensorflow.keras.datasets import mnist
(X_train, labels_train), (X_test, labels_test) = mnist.load_data()
#データを合体させる
X_all = np.concatenate([X_train, X_test])
labels_all = np.concatenate([labels_train, labels_test])

#分割の割合を設定
train_rate = 0.75
val_rate = 0.15
test_rate = 0.10

先ほどと同じように分けることができます。

ポイントとしては入力データとラベルデータを関数に与えているところです。

X_train, X_val_test, y_train, y_val_test = train_test_split(X_all, labels_all, train_size=train_rate, random_state=0)
X_val, X_test, y_val, y_test = train_test_split(X_val_test, y_val_test, train_size=val_rate / (val_rate + test_rate), random_state=0)

確認として入力データ=画像とラベルが合っているか確認してみます。

plt.imshow(X_train[1000])
print(y_train[1000])

出力は以下の通りです。7の画像に対して7が出力されているのでOKですね。

最後に

学習(train)と評価(val)に分ける記事は多いですが、train/val/testに分ける例があまり見かけなかったので記事させていただきました。

参考になればと思います。

コメント

タイトルとURLをコピーしました