加載Tensorflow模型並通過其他函數重用它

Question

我必須在初始化中加載模型，然後通過某個函數重用它。在我的架構函數中，通過一個模型運行預測多個數據集，我認爲每個數據集的硬盤驅動器負載模型不是一個好的解決方案。

我需要這樣的事情，即功能之間共享會話（或模型）：

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() 

有沒有做到這一點任何適當的方式？

Answer 1

您可以定義模型架構並將其權重作爲類加載。將其用於不同功能的各種數據集。

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

這是否可以幫助你嗎？