python分割train和test

用sklearn库中的train_test_split方法

from  sklearn.model_selection  import  train_test_split
train, test = train_test_split(data, random_state=2021, train_size=0.8)

自己用numpy写

import numpy as np

# 从 0~n 中随机选取 x 个数字
def getRandomIdx(n, x):
    return np.random.choice(np.arange(n), size=x, replace=False)

# 调用方法进行分割，获取train、test的index
total_cnt = data.shape[0]
train_idx = np.array(getRandomIdx(total_cnt, int(total_cnt * 0.8)))
test_idx = np.delete(np.arange(total_cnt), train_idx)

# 得到训练集和测试集
data_train = data[train_idx]
data_test = data[test_idx]