拆分訓練/測試數據

Question

我需要將數據集划分為訓練和驗證數據集。 我試着做80-20，但它不符合我的期望。

train_dataset, test_dataset = train_test_split(df, test_size = 0.2, random_state = 0)

我有一個包含 25 個城市的城市變量，每個城市都有幾個觀察值（在不同的日期）。 我想要的是對每個城市及其觀察結果進行 80-20 划分。 我不知道這個方法是否有名字，但我不知道如何 go 關於它。

謝謝你。

Answer 1

這樣的事情怎么樣？

import numpy as np
from pandas import DataFrame

def split_data(df: DataFrame, ratio: int): 
    length  = len(df)
    indices = list(range(length))
    np.random.shuffle(indices)

    train_indices = indices[:int(length * ratio)]
    test_indices  = indices[round(length * ratio):]

    train_set = df.iloc[train_indices]
    test_set  = df.iloc[test_indices]
    
    return (train_set, test_set)


train, test = split_data(df, 0.80)