Python Trick: Lazy Module Loading

Dboy Liao

4 min read

·

Feb 10, 2020

--

Listen

Share

Press enter or click to view image in full size

Introduction

I am heavily involved in the development of utensor_cgen, a CLI tool for uTensor, a neural network inference library for micro controllers.

Before we dive into the trick in detail in this post, it’s better for the readers to have some context of the issue first.

Briefly speaking, utensor_cgen can take a pre-trained model file such as a protobuff model file trained using Tensorflow and transpile it as equivalent implementation on uTensor runtime, i.e a bunch of C++ source and header files. As a result, I need to import libraries such as Tensorflow when the CLI is up and running.

However, it takes like forever to load Tensorflow. To be specific, it will take me ~3 seconds to load it on my machine:

Press enter or click to view image in full size

I use time to get the loading time of the cli.

As you can see in the real section, it took me 3.5 seconds to load the cli and print the help message.

As I inspect the program, I realize that my program spent a lot of time loading Tensorflow. That’s why it’s slow. Thank you very much Tensorflow, you really ruin my performance.

3 seconds seems harmless first glance, things will feel worse when you have to use this CLI over and over again on a daily basis. No one wants to wait 3 seconds for help messages. Fortunately, it did not take me long to figure out how to fix it.

The idea is simple: Let’s not to import Tensorflow!

Well, that’s not fair. What I mean is not to import Tensorflow right away but only import it when it’s required.

But how? You may ask.

Since python is an interpreted language, you can do something like this:

def my_func():
    import tensorflow as tf
    # do stuff with tf

or

class MyClass:
    def magic(self):
        import tensorflow as tf
        # do stuff with tf

In both cases, the import statement will be executed only when the functions/methods get invoked.

However, this is not a good idea since Tensorflow is such an essential part of the CLI. It is referenced in multiple places so we’ll have to copy and paste the import statement all around and it’s not good to maintain.

So far, the requirements are clear:

1. A single reference to Tensorflow, just like import tensorflow as tf at the top of my every python script so all members in the script can share the same reference to Tensorflow.
2. Only execute the import statement when attributes of tf get accessed.

Lazy Module Loading

Yes, let’s make our module lazy!

The plan is to write a class which takes a library/package name as string with customized attribute accessing, i.e implement __getattrib__ method. In the body of __getattrib__, it will check if the target library/package is loaded. If it is not loaded, we use importlib.import_module to load the library/package and get the attribute we need after then.

A simplified implementation may look like this:

import importlibclass LazyLoader:
    def __init__(self, lib_name):
        self.lib_name = lib_name
        self._mod = None    def __getattrib__(self, name):
        if self._mod is None:
            self._mod = importlib.import_module(self.lib_name)
        return getattr(self._mod, name)

It works as following:

1. Initially, the module, self._mod, is set as None and the target module name is given as lib_name as a string
2. When there is any attribute lookup to this object, i.e __getattrib__ invoked, we check if self._mod is still None and load the library with importlib.import_module if the module is missing.
3. delegate the attribute lookup to self._mod with getattr.

As a result, this lazy loader object will works just like a normal module except that it will not load the module util its attribute is required.

With this LazyLoader, we simply replace all the import tensorflow as tf with tf = LazyLoader('tensorflow') and here is the resulting loading time:

Press enter or click to view image in full size

Boooom, the loading time drop to less than 0.5 seconds! A huge boost.

Of course, this trick is not exclusive to Tensorflow. You can apply it to any import statements in your python script where you think the statements are slow. You can find my implementation here.

That being said, my implementation may not cover all use cases. For example, I do not handle the from ... import ... statement properly in my lazy loader.

For those who need this, I suggest you take a look of importlib documentation to see how to do it.

I hope you will like this trick and enjoy python programming!

Next time, we will talk about how to write plugins for utensor_cgen so we can inject a new backend to the CLI.