Fixed memory leak

main.py

@@ -49,6 +49,7 @@
 def train(save_dir, **kwargs):
   """"""
 
+  kwargs['config_file'] = kwargs.pop('config_file', '')
   load = kwargs.pop('load')
   try:
     if not load and os.path.isdir(save_dir):
@@ -58,7 +59,6 @@ def train(save_dir, **kwargs):
   except KeyboardInterrupt:
     print()
     sys.exit(0)
-
   network = Network(**kwargs)
   network.train(load=load)
   return

parser/bucket.py

@@ -119,6 +119,10 @@ def from_dataset(cls, dataset, bkt_idx, *args, **kwargs):
     return bucket
 
   #=============================================================
+  def reset_placeholders(self):
+    self.embed_model.reset_placeholders()
+    return
+  #=============================================================
   @property
   def tokens(self):
     return self._tokens

parser/multibucket.py

@@ -58,6 +58,13 @@ def __call__(self, vocab, keep_prob=None, moving_params=None):
         embeddings.append(bucket(vocab, keep_prob=keep_prob, moving_params=moving_params))
     return tf.nn.embedding_lookup(tf.concat(embeddings, axis=0), self.placeholder)
 
+  #=============================================================
+  def reset_placeholders(self):
+    self.placeholder = None
+    for bucket in self:
+      bucket.reset_placeholders()
+    return
+
   #=============================================================
   def generate_placeholder(self):
     """"""

parser/network.py

@@ -19,6 +19,7 @@
 from __future__ import division
 from __future__ import print_function
 
+import gc
 import os
 import time
 import codecs
@@ -85,12 +86,21 @@ def add_file_vocabs(self, conll_files):
         vocab.index_tokens()
     return
 
+  #=============================================================
+  def setup_vocabs(self):
+    """"""
+
+    for vocab in self.vocabs:
+      vocab.setup()
+    return 
+
   #=============================================================
   def train(self, load=False):
     """"""
 
     # prep the configurables
     self.add_file_vocabs(self.parse_files)
+    self.setup_vocabs()
     trainset = Trainset.from_configurable(self, self.vocabs, nlp_model=self.nlp_model)
     with tf.variable_scope(self.name.title()):
       train_tensors = trainset()
@@ -226,32 +236,28 @@ def parse(self, input_files, output_dir=None, output_file=None):
       raise ValueError('Cannot provide a value for --output_file when parsing multiple files')
     self.add_file_vocabs(input_files)
 
-    # load the model and prep the parse set
-    trainset = Trainset.from_configurable(self, self.vocabs, nlp_model=self.nlp_model)
-    with tf.variable_scope(self.name.title()):
-      train_tensors = trainset()
-    train_outputs = [train_tensors[train_key] for train_key in trainset.train_keys]
+    for input_file in input_files:
+      with tf.Graph().as_default():
+        config_proto = tf.ConfigProto()
+        if self.per_process_gpu_memory_fraction == -1:
+          config_proto.gpu_options.allow_growth = True
+        else:
+          config_proto.gpu_options.per_process_gpu_memory_fraction = self.per_process_gpu_memory_fraction
+        with tf.Session(config=config_proto) as sess:
+          # load the model and prep the parse set
+          self.setup_vocabs()
+          trainset = Trainset.from_configurable(self, self.vocabs, nlp_model=self.nlp_model)
+          with tf.variable_scope(self.name.title()):
+            train_tensors = trainset()
+          train_outputs = [train_tensors[train_key] for train_key in trainset.train_keys]
 
-    saver = tf.train.Saver(self.save_vars, max_to_keep=1)
-    config_proto = tf.ConfigProto()
-    if self.per_process_gpu_memory_fraction == -1:
-      config_proto.gpu_options.allow_growth = True
-    else:
-      config_proto.gpu_options.per_process_gpu_memory_fraction = self.per_process_gpu_memory_fraction
-    with tf.Session(config=config_proto) as sess:
-      for var in self.non_save_vars:
-        sess.run(var.initializer)
-      saver.restore(sess, tf.train.latest_checkpoint(self.save_dir))
-
-      # Iterate through files and batches
-      # (Hacky workaround to hopefully avoid memory issues)
-      default_graph = tf.get_default_graph()
-      default_graphdef = default_graph.as_graph_def()
-      for input_file in input_files:
-        with tf.Graph() as graph:
-          graph.import_graph_def(default_graphdef)
+          saver = tf.train.Saver(self.save_vars, max_to_keep=1)
+          for var in self.non_save_vars:
+            sess.run(var.initializer)
+          saver.restore(sess, tf.train.latest_checkpoint(self.save_dir))
 
-          parseset = Parseset.from_configurable(trainset, self.vocabs, parse_files=input_file, nlp_model=self.nlp_model)
+          # Iterate through files and batches
+          parseset = Parseset.from_configurable(self, self.vocabs, parse_files=input_file, nlp_model=self.nlp_model)
           with tf.variable_scope(self.name.title(), reuse=True):
             parse_tensors = parseset(moving_params=self.optimizer)
           parse_outputs = [parse_tensors[parse_key] for parse_key in parseset.parse_keys]
@@ -273,7 +279,9 @@ def parse(self, input_files, output_dir=None, output_file=None):
             probs.append(sess.run(parse_outputs, feed_dict=feed_dict))
             sents.append(tokens)
           parseset.write_probs(sents, output_path, probs)
-          del parseset
+      del trainset
+      del parseset
+      print('Finished one')
     if self.verbose:
       print(ctext('Parsing {0} file(s) took {1} seconds'.format(len(input_files), time.time()-start_time), 'bright_green'))
     return

parser/neural/models/embeds/base_embed.py

@@ -38,6 +38,11 @@ def __init__(self, *args, **kwargs):
     self.placeholder = None
     return
 
+  #=============================================================
+  def reset_placeholders(self):
+    self.placeholder = None
+    return
+
   #=============================================================
   def __call__(self, vocab, keep_prob=None, moving_params=None):
     """"""

parser/neural/models/nn.py

@@ -64,6 +64,8 @@ def embed_concat(self, vocabs, vocabs_to_merge=[['words', 'lemmas'], ['tags', 'x
         if merge_dict[vocab.name] == vocab.name:
           drop_mask = tf.expand_dims(linalg.random_mask(vocab.embed_keep_prob, tf.shape(placeholder)), 2)
           drop_masks.append(drop_mask)
+      for placeholder in placeholders:
+        print(placeholder.graph)
       total_masks = tf.add_n(drop_masks)
       scale_mask = len(drop_masks) / tf.maximum(total_masks, 1.)
     embed_dict = {}

parser/scripts/count_nonprojective.py

@@ -89,6 +89,7 @@ def __str__(self):
 
   args = parser.parse_args()
   for filename in args.files:
+    lang = re.search('([-\w]*)-ud', filename).group(1)
     nonproj = []
     with open(filename) as f:
       buff = []
@@ -101,4 +102,4 @@ def __str__(self):
           tree = DepTree(buff)
           nonproj.extend(tree.count_nonprojective())
           buff = []
-    print(filename, np.mean(nonproj)*100)
+    print(lang, np.mean(nonproj)*100)

parser/vocabs/base_vocab.py

@@ -90,6 +90,13 @@ def __call__(self, placeholder=None, moving_params=None):
     embeddings = self.embeddings if moving_params is None else moving_params.average(self.embeddings)
     return tf.nn.embedding_lookup(embeddings, placeholder)
 
+  #=============================================================
+  def setup(self):
+    """"""
+
+    self.placeholder = None
+    return
+
   #=============================================================
   def set_feed_dict(self, data, feed_dict):
     """"""
@@ -135,14 +142,14 @@ def counts(self):
   @property
   def embeddings(self):
     return self._embeddings
-  @embeddings.setter
-  def embeddings(self, matrix):
-    if matrix.shape[1] != self.embed_size:
-      raise ValueError("Matrix shape[1] of %d doesn't match expected shape of %d" % (matrix.shape[1], self.embed_size))
-    with tf.device('/cpu:0'):
-      with tf.variable_scope(self.name.title()):
-        self._embeddings = tf.Variable(matrix, name='Embeddings', dtype=tf.float32, trainable=True)
-    return
+  #@embeddings.setter
+  #def embeddings(self, matrix):
+  #  if matrix.shape[1] != self.embed_size:
+  #    raise ValueError("Matrix shape[1] of %d doesn't match expected shape of %d" % (matrix.shape[1], self.embed_size))
+  #  with tf.device('/cpu:0'):
+  #    with tf.variable_scope(self.name.title()):
+  #      self._embeddings = tf.Variable(matrix, name='Embeddings', dtype=tf.float32, trainable=True)
+  #  return
 
   #=============================================================
   def __getitem__(self, key):

parser/vocabs/index_vocab.py

@@ -59,6 +59,11 @@ def set_feed_dict(self, data, feed_dict):
     feed_dict[self.placeholder] = data
     return
 
+  #=============================================================
+  def setup(self):
+    self.placeholder = None
+    return
+
   #=============================================================
   def index(self, token):
     return 0 if token == '_' else int(token)

parser/vocabs/multivocab.py

@@ -57,6 +57,15 @@ def __call__(self, placeholder=None, moving_params=None):
     embeddings = [vocab(moving_params=moving_params) for vocab in self]
     return tf.add_n(embeddings)
 
+  #=============================================================
+  def setup(self):
+    """"""
+
+    self.placeholder = None
+    for vocab in self:
+      vocab.setup()
+    return
+
   #=============================================================
   def generate_placeholder(self):
     """"""

parser/vocabs/pretrained_vocab.py

@@ -66,6 +66,16 @@ def __call__(self, placeholder=None, moving_params=None):
     return matrix
     #return embeddings # changed in saves2/test8
 
+  #=============================================================
+  def setup(self):
+    """"""
+
+    self.placeholder = None
+    with tf.device('/cpu:0'):
+      with tf.variable_scope(self.name.title()):
+        self._embeddings = tf.Variable(self._embeddings_array, name='Embeddings', dtype=tf.float32, trainable=False)
+    return
+
   #=============================================================
   def load(self):
     """"""
@@ -93,7 +103,8 @@ def load(self):
     try:
       embeddings = np.stack(embeddings)
       embeddings = np.pad(embeddings, ( (len(self.special_tokens),0), (0,0) ), 'constant')
-      self.embeddings = np.stack(embeddings)
+      self._embeddings_array = np.stack(embeddings)
+      self._embed_size = self._embeddings_array.shape[1]
     except:
       shapes = set([embedding.shape for embedding in embeddings])
       raise ValueError("Couldn't stack embeddings with shapes in %s" % shapes)
@@ -123,13 +134,13 @@ def token_embed_size(self):
   @property
   def embeddings(self):
     return super(PretrainedVocab, self).embeddings
-  @embeddings.setter
-  def embeddings(self, matrix):
-    self._embed_size = matrix.shape[1]
-    with tf.device('/cpu:0'):
-      with tf.variable_scope(self.name.title()):
-        self._embeddings = tf.Variable(matrix, name='Embeddings', trainable=False)
-    return
+  #@embeddings.setter
+  #def embeddings(self, matrix):
+  #  self._embed_size = matrix.shape[1]
+  #  with tf.device('/cpu:0'):
+  #    with tf.variable_scope(self.name.title()):
+  #      self._embeddings = tf.Variable(matrix, name='Embeddings', trainable=False)
+  #  return
 
 #***************************************************************
 if __name__ == '__main__':

parser/vocabs/subtoken_vocab.py

@@ -61,11 +61,22 @@ def __init__(self, token_vocab, *args, **kwargs):
 
     embed_dims = [len(self), self.embed_size]
     if initialize_zero:
-      self.embeddings = np.zeros(embed_dims)
+      self._embeddings_array = np.zeros(embed_dims)
     else:
-      self.embeddings = np.random.randn(*embed_dims)
+      self._embeddings_array = np.random.randn(*embed_dims)
     return
 
+  #=============================================================
+  def setup(self):
+    """"""
+
+    self.placeholder = None
+    with tf.device('/cpu:0'):
+      with tf.variable_scope(self.name.title()):
+        self._embeddings = tf.Variable(self._embeddings_array, name='Embeddings', dtype=tf.float32, trainable=True)
+    self._multibucket.reset_placeholders()
+    return
+
   #=============================================================
   def __call__(self, placeholder=None, moving_params=None):
     """"""

parser/vocabs/token_vocab.py

@@ -56,11 +56,23 @@ def __init__(self, *args, **kwargs):
 
     embed_dims = [len(self), self.embed_size]
     if initialize_zero:
-      self.embeddings = np.zeros(embed_dims)
+      self._embeddings_array = np.zeros(embed_dims)
     else:
-      self.embeddings = np.random.randn(*embed_dims)
+      self._embeddings_array = np.random.randn(*embed_dims)
     return
 
+  #=============================================================
+  def setup(self):
+    """"""
+
+    self.placeholder = None
+    del self._embeddings
+    with tf.device('/cpu:0'):
+      with tf.variable_scope(self.name.title()):
+        self._embeddings = tf.Variable(self._embeddings_array, name='Embeddings', dtype=tf.float32, trainable=True)
+    return
+
+
   #=============================================================
   def count(self, conll_files=None):
     """"""