memusage_fit.py

#!/usr/bin/env python
from __future__ import division, print_function
import os
import sys
import time
import resource
import psutil
from sklearn import datasets
from sklearn.externals import joblib
from sklearn.feature_extraction.text import CountVectorizer, TfidfVectorizer, HashingVectorizer
from marisa_vectorizers import MarisaCountVectorizer, MarisaTfidfVectorizer

vectorizers = dict(
    count = CountVectorizer(),
    count2 = CountVectorizer(ngram_range=(1,2)),
    tfidf = TfidfVectorizer(),
    tfidf2 = TfidfVectorizer(ngram_range=(1,2)),
    hashing18 = HashingVectorizer(n_features=2**18),
    hashing20 = HashingVectorizer(n_features=2**20),
    marisa_count = MarisaCountVectorizer(),
    marisa_count2 = MarisaCountVectorizer(ngram_range=(1,2)),
    marisa_tfidf = MarisaTfidfVectorizer(),
    marisa_tfidf2 = MarisaTfidfVectorizer(ngram_range=(1,2)),
)

if __name__ == '__main__':
    if len(sys.argv) < 2:
        print("Available vectorizers:\n")
        print("\n".join(sorted(vectorizers.keys())))
        sys.exit()

    vecname = sys.argv[1]
    vec = vectorizers[vecname]

    newsgroups_train = datasets.fetch_20newsgroups(subset='train')

    p = psutil.Process(os.getpid())
    before = p.get_memory_info().rss / 2**20
    max_before = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 2**20
    start = time.time()

    vec.fit(newsgroups_train.data)

    end = time.time()
    after = p.get_memory_info().rss / 2**20
    max_after = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 2**20
    assert max_after >= max_before

    print("fit time: %0.1fs" % (end-start))
    print("fit memusage: %0.1fMB" % (max_after-before))

    before2 = p.get_memory_info().rss / 2**20

    joblib.dump(vec, vecname+".joblib")

    max_after2 = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 2**20
    assert max_after2 >= max_after

    print("dump time: %0.1fs" % (time.time()-end))
    print("dump memusage: %0.1fMB" % (max_after2-before2))