如何创建将被重用的tensorflow操作

在输入/训练模型之前,我有一些数据要处理。对于这个例子,我想做一个max pool 2d。我用tensorflow写了一个简短的函数。

import tensorflow
import tensorflow.nn as nn

def _tfMaxPool(arr, pool=(4,4), sess=None):
    op = nn.max_pool(arr, (1, 1, pool[0], 1), (1, 1, pool[0], 1 ), padding="VALID")
    op = nn.max_pool(op, (1, 1, 1, pool[1]), (1, 1, 1, pool[1]), padding="VALID")
    if sess is None:
        sess = tensorflow.Session();

    return sess.run(op)

问题是这样每次都会在我的图中添加节点,这似乎会扰乱我的会话。另一种方法是创建模型。

import keras

seq = keras.Sequential([ 
            keras.layers.InputLayer((1, 512, 512)), 
            keras.layers.MaxPool2D((4, 4), (4, 4), data_format="channels_first")
                     ])
def _tfMaxPool2(arr, pool=(4,4), sess=None):
    swapped = arr.swapaxes(0,1)
    return seq.predict(swapped).swapaxes(0,1)

这个模型和我想要的几乎一模一样,但我认为我遗漏了一些基本的东西。