Reduce fasttext memory usage by computing ngrams on the fly

Author: jbaiter

gensim/models/deprecated/fasttext_wrapper.py

@@ -70,9 +70,7 @@ def __init__(self):
         self.syn0_vocab_norm = None
         self.syn0_ngrams = None
         self.syn0_ngrams_norm = None
-        self.ngrams = {}
         self.hash2index = {}
-        self.ngrams_word = {}
         self.min_n = 0
         self.max_n = 0
 
@@ -99,17 +97,18 @@ def word_vec(self, word, use_norm=False):
             return super(FastTextKeyedVectors, self).word_vec(word, use_norm)
         else:
             word_vec = np.zeros(self.syn0_ngrams.shape[1], dtype=np.float32)
-            ngrams = compute_ngrams(word, self.min_n, self.max_n)
-            ngrams = [ng for ng in ngrams if ng in self.ngrams]
+            hashes = [ft_hash(ng) % self.bucket
+                      for ng in compute_ngrams(word, self.min_n, self.max_n)]
+            hashes = [h for h in hashes if h in self.hash2index]
             if use_norm:
                 ngram_weights = self.syn0_ngrams_norm
             else:
                 ngram_weights = self.syn0_ngrams
-            for ngram in ngrams:
-                word_vec += ngram_weights[self.ngrams[ngram]]
+            for ngram_hash in hashes:
+                word_vec += ngram_weights[self.hash2index[ngram_hash]]
             if word_vec.any():
-                return word_vec / len(ngrams)
-            else:  # No ngrams of the word are present in self.ngrams
+                return word_vec / len(hashes)
+            else:  # No hashes of any ngrams of the word are present in self.hash2index
                 raise KeyError('all ngrams for word %s absent from model' % word)
 
     def init_sims(self, replace=False):
@@ -143,7 +142,8 @@ def __contains__(self, word):
             return True
         else:
             char_ngrams = compute_ngrams(word, self.min_n, self.max_n)
-            return any(ng in self.ngrams for ng in char_ngrams)
+            return any(ft_hash(ng) % self.bucket in self.hash2index
+                       for ng in char_ngrams)
 
     @classmethod
     def load_word2vec_format(cls, *args, **kwargs):
@@ -277,6 +277,12 @@ def load(cls, *args, **kwargs):
         if hasattr(model.wv, 'syn0_all'):
             setattr(model.wv, 'syn0_ngrams', model.wv.syn0_all)
             delattr(model.wv, 'syn0_all')
+        setattr(model.wv, 'bucket', model.wv.syn0_ngrams.shape[0])
+        if not hasattr(model.wv, 'hash2index') and hasattr(model.wv, 'ngrams'):
+            model.wv.hash2index = {}
+            for i, ngram in enumerate(model.wv.ngrams):
+                ngram_hash = ft_hash(ngram) % model.wv.bucket
+                model.wv.hash2index[ngram_hash] = i
         return model
 
     @classmethod
@@ -316,6 +322,7 @@ def load_model_params(self, file_handle):
         self.hs = loss == 1
         self.sg = model == 2
         self.bucket = bucket
+        self.wv.bucket = bucket
         self.wv.min_n = minn
         self.wv.max_n = maxn
         self.sample = t
@@ -394,7 +401,6 @@ def init_ngrams(self):
         vectors are discarded here to save space.
 
         """
-        self.wv.ngrams = {}
         all_ngrams = []
         self.wv.syn0 = np.zeros((len(self.wv.vocab), self.vector_size), dtype=REAL)
 
@@ -406,9 +412,9 @@ def init_ngrams(self):
         self.num_ngram_vectors = len(all_ngrams)
         ngram_indices = []
         for i, ngram in enumerate(all_ngrams):
-            ngram_hash = ft_hash(ngram)
+            ngram_hash = ft_hash(ngram) % self.bucket
             ngram_indices.append(len(self.wv.vocab) + ngram_hash % self.bucket)
-            self.wv.ngrams[ngram] = i
+            self.wv.hash2index[ngram_hash] = i
         self.wv.syn0_ngrams = self.wv.syn0_ngrams.take(ngram_indices, axis=0)
 
         ngram_weights = self.wv.syn0_ngrams
@@ -421,7 +427,9 @@ def init_ngrams(self):
         for w, vocab in self.wv.vocab.items():
             word_ngrams = compute_ngrams(w, self.wv.min_n, self.wv.max_n)
             for word_ngram in word_ngrams:
-                self.wv.syn0[vocab.index] += np.array(ngram_weights[self.wv.ngrams[word_ngram]])
+                ng_hash = ft_hash(word_ngram) % self.bucket
+                self.wv.syn0[vocab.index] += np.array(
+                    ngram_weights[self.wv.hash2index[ng_hash]])
 
             self.wv.syn0[vocab.index] /= (len(word_ngrams) + 1)
         logger.info(

gensim/models/fasttext.py

@@ -39,7 +39,6 @@
 from gensim.models.base_any2vec import BaseWordEmbeddingsModel
 from gensim.models.utils_any2vec import _compute_ngrams, _ft_hash
 
-from six import iteritems
 from gensim.utils import deprecated, call_on_class_only
 from gensim import utils
 
@@ -93,10 +92,12 @@ def train_batch_cbow(model, sentences, alpha, work=None, neu1=None):
 
                 for index in word2_indices:
                     vocab_subwords_indices += [index]
-                    word2_subwords += model.wv.ngrams_word[model.wv.index2word[index]]
+                    word2_subwords += _compute_ngrams(model.wv.index2word[index],
+                                                      model.min_n, model.max_n)
 
                 for subword in word2_subwords:
-                    ngrams_subwords_indices.append(model.wv.ngrams[subword])
+                    ngrams_subwords_indices.append(
+                        model.wv.hash2index[_ft_hash(subword) % model.bucket])
 
                 l1_vocab = np_sum(model.wv.syn0_vocab[vocab_subwords_indices], axis=0)  # 1 x vector_size
                 l1_ngrams = np_sum(model.wv.syn0_ngrams[ngrams_subwords_indices], axis=0)  # 1 x vector_size
@@ -144,10 +145,11 @@ def train_batch_sg(model, sentences, alpha, work=None, neu1=None):
                 start = max(0, pos - model.window + reduced_window)
 
                 subwords_indices = [word.index]
-                word2_subwords = model.wv.ngrams_word[model.wv.index2word[word.index]]
+                word2_subwords = _compute_ngrams(model.wv.index2word[word.index],
+                                                 model.min_n, model.max_n)
 
                 for subword in word2_subwords:
-                    subwords_indices.append(model.wv.ngrams[subword])
+                    subwords_indices.append(model.wv.hash2index[_ft_hash(subword) % model.bucket])
 
                 for pos2, word2 in enumerate(word_vocabs[start:(pos + model.window + 1 - reduced_window)], start):
                     if pos2 != pos:  # don't train on the `word` itself
@@ -278,6 +280,7 @@ def __init__(self, sentences=None, sg=0, hs=0, size=100, alpha=0.025, window=5,
             sorted_vocab=bool(sorted_vocab), null_word=null_word)
         self.trainables = FastTextTrainables(
             vector_size=size, seed=seed, bucket=bucket, hashfxn=hashfxn)
+        self.wv.bucket = self.bucket
 
         super(FastText, self).__init__(
             sentences=sentences, workers=workers, vector_size=size, epochs=iter, callbacks=callbacks,
@@ -396,6 +399,36 @@ def _clear_post_train(self):
         self.wv.vectors_vocab_norm = None
         self.wv.vectors_ngrams_norm = None
 
+    def estimate_memory(self, vocab_size=None, report=None):
+        vocab_size = vocab_size or len(self.wv.vocab)
+        vec_size = self.vector_size * np.dtype(np.float32).itemsize
+        l1_size = self.layer1_size * np.dtype(np.float32).itemsize
+        report = report or {}
+        report['vocab'] = len(self.wv.vocab) * (700 if self.hs else 500)
+        report['syn0_vocab'] = len(self.wv.vocab) * vec_size
+        num_buckets = self.bucket
+        if self.hs:
+            report['syn1'] = len(self.wv.vocab) * l1_size
+        if self.negative:
+            report['syn1neg'] = len(self.wv.vocab) * l1_size
+        if self.word_ngrams > 0 and self.wv.vocab:
+            buckets = set()
+            for word in self.wv.vocab:
+                ngrams = _compute_ngrams(word, self.min_n, self.max_n)
+                buckets.update(_ft_hash(ng) % self.bucket for ng in ngrams)
+            num_buckets = len(buckets)
+            report['syn0_ngrams'] = len(buckets) * vec_size
+        elif self.word_ngrams > 0:
+            logger.warn('subword information is enabled, but no vocabulary '
+                        'could be found, estimated required memory might be '
+                        'inaccurate!')
+        report['total'] = sum(report.values())
+        logger.info(
+            "estimated required memory for %i words, %i buckets and %i dimensions: %i bytes",
+            len(self.wv.vocab), num_buckets, self.vector_size, report['total']
+        )
+        return report
+
     def _do_train_job(self, sentences, alpha, inits):
         """Train a single batch of sentences. Return 2-tuple `(effective word count after
         ignoring unknown words and sentence length trimming, total word count)`.
@@ -580,6 +613,7 @@ def _load_model_params(self, file_handle):
         self.hs = loss == 1
         self.sg = model == 2
         self.trainables.bucket = bucket
+        self.wv.bucket = bucket
         self.wv.min_n = minn
         self.wv.max_n = maxn
         self.vocabulary.sample = t
@@ -709,18 +743,8 @@ def prepare_vocab(self, hs, negative, wv, update=False, keep_raw_vocab=False, tr
         report_values = super(FastTextVocab, self).prepare_vocab(
             hs, negative, wv, update=update, keep_raw_vocab=keep_raw_vocab, trim_rule=trim_rule,
             min_count=min_count, sample=sample, dry_run=dry_run)
-        self.build_ngrams(wv, update=update)
         return report_values
 
-    def build_ngrams(self, wv, update=False):
-        if not update:
-            wv.ngrams_word = {}
-            for w, v in iteritems(wv.vocab):
-                wv.ngrams_word[w] = _compute_ngrams(w, wv.min_n, wv.max_n)
-        else:
-            for w, v in iteritems(wv.vocab):
-                wv.ngrams_word[w] = _compute_ngrams(w, wv.min_n, wv.max_n)
-
 
 class FastTextTrainables(Word2VecTrainables):
     def __init__(self, vector_size=100, seed=1, hashfxn=hash, bucket=2000000):
@@ -744,54 +768,43 @@ def init_ngrams_weights(self, wv, update=False, vocabulary=None):
 
         """
         if not update:
-            wv.ngrams = {}
             wv.vectors_vocab = empty((len(wv.vocab), wv.vector_size), dtype=REAL)
             self.vectors_vocab_lockf = ones((len(wv.vocab), wv.vector_size), dtype=REAL)
 
             wv.vectors_ngrams = empty((self.bucket, wv.vector_size), dtype=REAL)
             self.vectors_ngrams_lockf = ones((self.bucket, wv.vector_size), dtype=REAL)
 
-            all_ngrams = []
-            for w, ngrams in iteritems(wv.ngrams_word):
-                all_ngrams += ngrams
-
-            all_ngrams = list(set(all_ngrams))
-            wv.num_ngram_vectors = len(all_ngrams)
-            logger.info("Total number of ngrams is %d", len(all_ngrams))
-
             wv.hash2index = {}
             ngram_indices = []
             new_hash_count = 0
-            for i, ngram in enumerate(all_ngrams):
-                ngram_hash = _ft_hash(ngram) % self.bucket
-                if ngram_hash in wv.hash2index:
-                    wv.ngrams[ngram] = wv.hash2index[ngram_hash]
-                else:
-                    ngram_indices.append(ngram_hash % self.bucket)
-                    wv.hash2index[ngram_hash] = new_hash_count
-                    wv.ngrams[ngram] = wv.hash2index[ngram_hash]
-                    new_hash_count = new_hash_count + 1
+            wv.num_ngram_vectors = 0
+            for word in wv.vocab.keys():
+                for ngram in _compute_ngrams(word, wv.min_n, wv.max_n):
+                    ngram_hash = _ft_hash(ngram) % self.bucket
+                    if ngram_hash not in wv.hash2index:
+                        wv.num_ngram_vectors += 1
+                        ngram_indices.append(ngram_hash)
+                        wv.hash2index[ngram_hash] = new_hash_count
+                        new_hash_count = new_hash_count + 1
+
+            logger.info("Total number of ngrams is %d", wv.num_ngram_vectors)
 
             wv.vectors_ngrams = wv.vectors_ngrams.take(ngram_indices, axis=0)
             self.vectors_ngrams_lockf = self.vectors_ngrams_lockf.take(ngram_indices, axis=0)
             self.reset_ngrams_weights(wv)
         else:
-            new_ngrams = []
-            for w, ngrams in iteritems(wv.ngrams_word):
-                new_ngrams += [ng for ng in ngrams if ng not in wv.ngrams]
-
-            new_ngrams = list(set(new_ngrams))
-            wv.num_ngram_vectors += len(new_ngrams)
-            logger.info("Number of new ngrams is %d", len(new_ngrams))
             new_hash_count = 0
-            for i, ngram in enumerate(new_ngrams):
-                ngram_hash = _ft_hash(ngram) % self.bucket
-                if ngram_hash not in wv.hash2index:
-                    wv.hash2index[ngram_hash] = new_hash_count + self.old_hash2index_len
-                    wv.ngrams[ngram] = wv.hash2index[ngram_hash]
-                    new_hash_count = new_hash_count + 1
-                else:
-                    wv.ngrams[ngram] = wv.hash2index[ngram_hash]
+            num_new_ngrams = 0
+            for word in wv.vocab.keys():
+                for ngram in _compute_ngrams(word, wv.min_n, wv.max_n):
+                    ngram_hash = _ft_hash(ngram) % self.bucket
+                    if ngram_hash not in wv.hash2index:
+                        wv.hash2index[ngram_hash] = new_hash_count + self.old_hash2index_len
+                        new_hash_count = new_hash_count + 1
+                        num_new_ngrams += 1
+
+            wv.num_ngram_vectors += num_new_ngrams
+            logger.info("Number of new ngrams is %d", num_new_ngrams)
 
             rand_obj = np.random
             rand_obj.seed(self.seed)
@@ -833,10 +846,11 @@ def get_vocab_word_vecs(self, wv):
         """Calculate vectors for words in vocabulary and stores them in `vectors`."""
         for w, v in wv.vocab.items():
             word_vec = np.copy(wv.vectors_vocab[v.index])
-            ngrams = wv.ngrams_word[w]
+            ngrams = _compute_ngrams(w, wv.min_n, wv.max_n)
             ngram_weights = wv.vectors_ngrams
             for ngram in ngrams:
-                word_vec += ngram_weights[wv.ngrams[ngram]]
+                word_vec += ngram_weights[
+                    wv.hash2index[_ft_hash(ngram) % self.bucket]]
             word_vec /= (len(ngrams) + 1)
             wv.vectors[v.index] = word_vec
 
@@ -847,20 +861,21 @@ def init_ngrams_post_load(self, file_name, wv):
         vectors are discarded here to save space.
 
         """
-        all_ngrams = []
         wv.vectors = np.zeros((len(wv.vocab), wv.vector_size), dtype=REAL)
 
         for w, vocab in wv.vocab.items():
-            all_ngrams += _compute_ngrams(w, wv.min_n, wv.max_n)
             wv.vectors[vocab.index] += np.array(wv.vectors_ngrams[vocab.index])
 
-        all_ngrams = set(all_ngrams)
-        wv.num_ngram_vectors = len(all_ngrams)
         ngram_indices = []
-        for i, ngram in enumerate(all_ngrams):
-            ngram_hash = _ft_hash(ngram)
-            ngram_indices.append(len(wv.vocab) + ngram_hash % self.bucket)
-            wv.ngrams[ngram] = i
+        wv.num_ngram_vectors = 0
+        for word in wv.vocab.keys():
+            for ngram in _compute_ngrams(word, wv.min_n, wv.max_n):
+                ngram_hash = _ft_hash(ngram) % self.bucket
+                if ngram_hash in wv.hash2index:
+                    continue
+                ngram_indices.append(len(wv.vocab) + ngram_hash)
+                wv.hash2index[ngram_hash] = wv.num_ngram_vectors
+                wv.num_ngram_vectors += 1
         wv.vectors_ngrams = wv.vectors_ngrams.take(ngram_indices, axis=0)
 
         ngram_weights = wv.vectors_ngrams
@@ -873,7 +888,8 @@ def init_ngrams_post_load(self, file_name, wv):
         for w, vocab in wv.vocab.items():
             word_ngrams = _compute_ngrams(w, wv.min_n, wv.max_n)
             for word_ngram in word_ngrams:
-                wv.vectors[vocab.index] += np.array(ngram_weights[wv.ngrams[word_ngram]])
+                vec_idx = wv.hash2index[_ft_hash(word_ngram) % self.bucket]
+                wv.vectors[vocab.index] += np.array(ngram_weights[vec_idx])
 
             wv.vectors[vocab.index] /= (len(word_ngrams) + 1)
         logger.info(