pythonmxnet

Python multiprocessing.Pool cannot pickle mxnet.mod.Module object


here is approximately what I did:

import mxnet as mx
import cv2
from multiprocessing import Pool
from itertools import repeat

num_worker=4
CNNNet=[]
img = cv2.imread('../datasets/1.jpg')
sym, arg_params, aux_params = mx.model.load_checkpoint('det1', 0)
for i in range(num_worker):
    worker_net = mx.mod.Module(symbol=sym,label_names=None)
    worker_net.bind(data_shapes=[('data', (1, 3, 1000, 1000))],for_training=False)
    worker_net.set_params(arg_params,aux_params)
    CNNNet.append(worker_net)
pool = Pool(num_worker)
threshold = 0.6
res = pool.map(do_work_warpper,zip(repeat(img),CNNNet[:num_worker],repeat(threshold)))

and the do_work_warpper() function is:

def do_work_warpper(args):
    return do_work(*args)
def do_work(img,net,threshold):
    #do image predict job here
    return res

I am puzzled by the question that when using multiprocessing.Pool with the mx.mod.Module object, I get the error in python3.6:

TypeError: can't pickle module objects

or in python2.7:

PicklingError: Can't pickle <type 'module'>: attribute lookup __builtin__.module failed

any suggestion will be appreciated.


Solution

  • The reason why you're getting this exception is because multiprocessing needs to be able to pickle the variables you pass to your workers in order to pass them between various processes it spawns.

    The error:

    TypeError: can't pickle module objects
    

    Suggests that one of the variables you're passing to your Pool contains a module (or a class that has a module as an attribute).

    To demonstrate the issue, have a look at these two classes:

    import os
    
    class Pickable: 
        a = 1
    
    class UnPickable:
        def __init__(self):
            self.mod = os
    

    If you try to pickle instances of these two classes, you'll get:

    In [11]: pickle.dumps(Pickable())
    Out[11]: b'\x80\x03c__main__\nPickable\nq\x00)\x81q\x01.'
    
    In [10]: pickle.dumps(UnPickable())
    ---------------------------------------------------------------------------
    TypeError                                 Traceback (most recent call last)
    <ipython-input-10-7d4d725a1c6d> in <module>()
    ----> 1 pickle.dumps(UnPickable())
    
    TypeError: can't pickle module objects
    

    That being said - either you create your own class that mimicks the functionality of mx.mod.Module, but is serializable - OR (better solution in my opinion) use simple (https://docs.python.org/3.1/library/pickle.html#what-can-be-pickled-and-unpickled) python-builtin types to pass variables into workers of your Pool and construct mx.mod.Module instances inside them on their own.