关于 Tensorflow 使用 Dataset API 占用内存高的问题

初学深度学习，看的是《 Hands on Machine Learning with Scikit-Learn and Tensorflow 》这本书，书中用自己定义的shuffle_batch函数实现将数据分批输入神经网络的功能，数据集用的是 MNIST。

书上的函数定义如下：

def shuffle_batch(X, y, batch_size):
    rnd_idx = np.random.permutation(len(X))
    n_batches = len(X) // batch_size
    for batch_idx in np.array_split(rnd_idx, n_batches):
        X_batch, y_batch = X[batch_idx], y[batch_idx]
        yield X_batch, y_batch

楼主上网搜索一下发现用 Dataset API 和它的shuffle、batch和repeat函数可能可以更加优雅地实现分批输入的功能，于是就写了下面的代码：

train_data = tf.data.Dataset.from_tensor_slices((X_train, y_train))
train_data = train_data.shuffle(m)
train_data = train_data.batch(batch_size)
train_data = train_data.repeat()
td_iter = train_data.make_one_shot_iterator()
features, labels = td_iter.get_next()

with tf.Session() as sess:
    sess.run(init)
    for epoch in range(n_epochs):
        for iteration in range(n_batchs):
            X_batch, y_batch = sess.run([features, labels])
            sess.run(training_op, feed_dict={X:X_batch, y:y_batch})
        acc_train = accuracy.eval(feed_dict={X:X_batch, y:y_batch})
        acc_test = accuracy.eval(feed_dict={X:X_test, y:y_test})
        print(epoch, "Train accuracy:", acc_train, "Test accuracy:", acc_test)
    save_path = saver.save(sess, './my_model')

但是我发现这段代码虽然也能训练出类似精度的模型，但是在打印出第一个 epoch 的输出前，内存占用极高，而且要等好久才会有第一个输出（后面的输出就花费正常时间）。

如果是按照书上的代码来训练（不使用 Dataset API ），内存几乎没有任何波动。但是我觉得就算是用了 Dataset API，MNIST 这个数据集也不大吧？要占用这么多内存么？

同样的内存占用情况也发生在下面的代码：

with tf.Session() as sess:
    sess.run(init)
    sess.run([features, labels])

我觉得是不是我代码哪里写错了？因为刚接触这个 API，是模仿人家的写法写的，希望大家解答下疑惑哈