train_test_split分层抽样失效问题求助

各位前辈好，我最近在做数据划分的时候碰到了分层抽样的问题，折腾了半天没搞定，想请大家指点一下。

先把我的代码贴出来：

import pandas as pd 
from sklearn.model_selection import train_test_split  

# 构建数据集，1000条数据，前500条标签是1，后500条是0
ids = list(range(1000)) 
label = 500 * [1.0] + 500 * [0.0] 
df = pd.DataFrame({"id": ids, "label": label})  

# 划分比例：训练集80%，验证集10%，测试集10%
train_p = 0.8 
val_p = 0.1 
test_p = 0.1  

# 手动计算各集样本数
n_train = int(len(df) * train_p) 
n_val = int(len(df) * val_p) 
n_test = len(df) - n_train - n_val  

print("* Step 3")

我本来想让训练、验证、测试集里的标签分布都和原数据集一样（1:1），但如果直接按顺序取前800条当训练集的话，训练集里会有500个1和300个0，完全不符合分层的要求。后来我试着用train_test_split来做，但不知道是不是参数没设置对，分层效果还是不对。请问有没有正确实现这种分层抽样的方法呀？

备注：内容来源于stack exchange，提问作者Antonio Sesto