train_test_split(X, y, test_size=0.2, stratify=y)

参数
stratify=y : 按照数据集中y的比例分配给train和test，使得train和test中各类别数据的比例与原数据集的比例一致。
举例：
原数据集中有100条数据，A类有80条，B类有20条，且前80条全为A类。
那么我们在不指定stratify参数的情况下，切分出前80条的数据就全部为A类。
所以我们要设置stratify=y来使得切分出的测试集与训练集中包含的类别为4:1。
通常在数据集的分类分布不平衡的情况下会用到stratify。