fix exiting behavior

examples/cifar10-convnet.py

@@ -97,7 +97,7 @@ def get_data(train_or_test):
     ds = AugmentImageComponent(ds, augmentors)
     ds = BatchData(ds, 128, remainder=not isTrain)
     if isTrain:
-        ds = PrefetchDataZMQ(ds, 5)
+        ds = PrefetchData(ds, 10, 5)
     return ds
 
 def get_config():
@@ -127,7 +127,7 @@ def get_config():
         session_config=sess_config,
         model=Model(),
         step_per_epoch=step_per_epoch,
-        max_epoch=300,
+        max_epoch=250,
     )
 
 if __name__ == '__main__':

tensorpack/dataflow/prefetch.py

@@ -67,14 +67,6 @@ class PrefetchData(ProxyDataFlow):
             dp = self.queue.get()
             yield dp
 
-    def __del__(self):
-        logger.info("Prefetch process exiting...")
-        self.queue.close()
-        for x in self.procs:
-            x.terminate()
-        logger.info("Prefetch process exited.")
-
-
 class PrefetchProcessZMQ(multiprocessing.Process):
     def __init__(self, ds, conn_name):
         """
@@ -118,6 +110,10 @@ class PrefetchDataZMQ(ProxyDataFlow):
         for x in self.procs:
             x.start()
 
+        # __del__ not guranteed to get called at exit
+        import atexit
+        atexit.register(lambda x: x.__del__(), self)
+
     def get_data(self):
         for _ in range(self._size):
             dp = loads(self.socket.recv(copy=False))
@@ -125,6 +121,7 @@ class PrefetchDataZMQ(ProxyDataFlow):
 
     def __del__(self):
         logger.info("Prefetch process exiting...")
+        if not self.context.closed:
             self.context.destroy(0)
         for x in self.procs:
             x.terminate()