From cc3c1e86868d9b5c65805ab47482721b1c191672 Mon Sep 17 00:00:00 2001 From: tdozat Date: Thu, 2 Nov 2017 19:24:52 -0700 Subject: [PATCH] Fixed memory leak --- main.py | 2 +- parser/bucket.py | 4 ++ parser/multibucket.py | 7 +++ parser/network.py | 58 +++++++++++++---------- parser/neural/models/embeds/base_embed.py | 5 ++ parser/neural/models/nn.py | 2 + parser/scripts/count_nonprojective.py | 3 +- parser/vocabs/base_vocab.py | 23 +++++---- parser/vocabs/index_vocab.py | 5 ++ parser/vocabs/multivocab.py | 9 ++++ parser/vocabs/pretrained_vocab.py | 27 +++++++---- parser/vocabs/subtoken_vocab.py | 15 +++++- parser/vocabs/token_vocab.py | 16 ++++++- 13 files changed, 129 insertions(+), 47 deletions(-) diff --git a/main.py b/main.py index f273037..7230cef 100644 --- a/main.py +++ b/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 diff --git a/parser/bucket.py b/parser/bucket.py index 5020fbf..f757c43 100644 --- a/parser/bucket.py +++ b/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 diff --git a/parser/multibucket.py b/parser/multibucket.py index adbca06..57ff869 100644 --- a/parser/multibucket.py +++ b/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): """""" diff --git a/parser/network.py b/parser/network.py index 61d21b2..9160617 100644 --- a/parser/network.py +++ b/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 diff --git a/parser/neural/models/embeds/base_embed.py b/parser/neural/models/embeds/base_embed.py index 19a3487..52a0fe5 100644 --- a/parser/neural/models/embeds/base_embed.py +++ b/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): """""" diff --git a/parser/neural/models/nn.py b/parser/neural/models/nn.py index 6ff2351..71a6d21 100644 --- a/parser/neural/models/nn.py +++ b/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 = {} diff --git a/parser/scripts/count_nonprojective.py b/parser/scripts/count_nonprojective.py index b876704..abc6996 100644 --- a/parser/scripts/count_nonprojective.py +++ b/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) diff --git a/parser/vocabs/base_vocab.py b/parser/vocabs/base_vocab.py index 207090c..dd72594 100644 --- a/parser/vocabs/base_vocab.py +++ b/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): diff --git a/parser/vocabs/index_vocab.py b/parser/vocabs/index_vocab.py index ab7426f..78d3bce 100644 --- a/parser/vocabs/index_vocab.py +++ b/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) diff --git a/parser/vocabs/multivocab.py b/parser/vocabs/multivocab.py index 3bc0b29..bf5400f 100644 --- a/parser/vocabs/multivocab.py +++ b/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): """""" diff --git a/parser/vocabs/pretrained_vocab.py b/parser/vocabs/pretrained_vocab.py index 3d42afa..9050edf 100644 --- a/parser/vocabs/pretrained_vocab.py +++ b/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__': diff --git a/parser/vocabs/subtoken_vocab.py b/parser/vocabs/subtoken_vocab.py index 0c9ff22..91ba423 100644 --- a/parser/vocabs/subtoken_vocab.py +++ b/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): """""" diff --git a/parser/vocabs/token_vocab.py b/parser/vocabs/token_vocab.py index 40f9773..bdd88b0 100644 --- a/parser/vocabs/token_vocab.py +++ b/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): """"""