Deleter Thread

del_thread.py

import mxnet as mx
from mxnet.base import check_call, _LIB
from multiprocessing.pool import ThreadPool
import time

num_workers = 16

old_deleter = mx.nd.NDArray.__del__
del_pool = ThreadPool(num_workers)

def deleter_fn(handle):
    check_call(_LIB.MXNDArrayFree(handle))


small_size = 1000000


def new_deleter(self):
    global small_size
    shape = self.shape
    ctx  = self.context
    size = self.size
    tic = time.time()
    flag = 0
    if self.size < small_size:
        check_call(_LIB.MXNDArrayFree(self.handle))
    else:
        flag = 1
        del_pool.apply_async(deleter_fn, (self.handle, ))
    cost = time.time() - tic
    if cost > 1e-3:
        print("DEL:", cost, shape, ctx, size, flag)
        small_size = size

mx.nd.NDArray.__del__ = new_deleter

if __name__ == '__main__':
    import time
    a = mx.nd.array([1,2,3])
    del a
    time.sleep(3)