如何将tf.data.Dataset拆分成x_train、y_train、x_test、y_test

import tensorflow as tf
from sklearn.model_selection import train_test_split

dataset = tf.data.Dataset.from_tensor_slices((features, labels))

train_dataset, test_dataset = train_test_split(dataset, test_size=0.2, random_state=42)

x_train = train_dataset.map(lambda x, y: x)
y_train = train_dataset.map(lambda x, y: y)
x_test = test_dataset.map(lambda x, y: x)
y_test = test_dataset.map(lambda x, y: y)