加载Tensorflow模型并通过其他功能重用它

时间:2017-03-01 16:17:45

标签: tensorflow

我必须在初始化中加载模型,然后通过某个函数重用它。在我的架构功能中,运行一个模型预测几个数据集,我认为每个数据集的硬盘驱动器负载模型不是一个好的解决方案。

我需要这样的东西,即在函数之间共享会话(或模型):

def __init__(self):
    self.graph = tf.Graph()
    with self.graph.as_default():
        self.sess = tf.Session()
        with sess.as_default():
           saver.restore(tf.get_default_session(), path_to_checkpoint)

def some_func():
    with self.graph.as_default():
        with self.sess.as_default():
            self.sess.run()

有没有正确的方法呢?

1 个答案:

答案 0 :(得分:0)

您可以定义模型体系结构并将其权重加载为类。将它用于不同功能的各种数据集。

class vgg16:
def __init__(self, imgs, weights=None, sess=None):
    self.imgs = imgs
    self.convlayers()
    self.fc_layers()
    self.probs = tf.nn.softmax(self.fc3l)
    if weights is not None and sess is not None:
        self.load_weights(weights, sess)

def convlayers(self):
    self.parameters = []
      ::
      ::

if __name__ == '__main__':
    sess = tf.Session()
    imgs = tf.placeholder(tf.float32, [None, 224, 224, 3])
    vgg = vgg16(imgs, 'vgg16_weights.npz', sess)

    # run inference here on multiple datasets

点击此处:https://www.cs.toronto.edu/~frossard/vgg16/vgg16.py

这对你有帮助吗?