Better picklability for "spawn" mp context

tensorpack/dataflow/common.py

@@ -295,16 +295,17 @@ class MapDataComponent(MapData):
                 return None to discard this datapoint.
             index (int): index of the component.
         """
-        index = int(index)
+        self._index = int(index)
+        self._func = func
+        super(MapDataComponent, self).__init__(ds, self._mapper)
 
-        def f(dp):
-            r = func(dp[index])
+    def _mapper(self, dp):
+        r = self._func(dp[self._index])
         if r is None:
             return None
         dp = list(dp)   # shallow copy to avoid modifying the list
-            dp[index] = r
+        dp[self._index] = r
         return dp
-        super(MapDataComponent, self).__init__(ds, f)
 
 
 class RepeatedData(ProxyDataFlow):

tensorpack/dataflow/image.py

@@ -103,23 +103,22 @@ class AugmentImageComponent(MapDataComponent):
             self.augs = augmentors
         else:
             self.augs = AugmentorList(augmentors)
+        self._copy = copy
 
-        exception_handler = ExceptionHandler(catch_exceptions)
-
-        def func(x):
-            check_dtype(x)
-            with exception_handler.catch():
-                if copy:
-                    x = copy_mod.deepcopy(x)
-                return self.augs.augment(x)
-
-        super(AugmentImageComponent, self).__init__(
-            ds, func, index)
+        self._exception_handler = ExceptionHandler(catch_exceptions)
+        super(AugmentImageComponent, self).__init__(ds, self._aug_mapper, index)
 
     def reset_state(self):
         self.ds.reset_state()
         self.augs.reset_state()
 
+    def _aug_mapper(self, x):
+        check_dtype(x)
+        with self._exception_handler.catch():
+            if self._copy:
+                x = copy_mod.deepcopy(x)
+            return self.augs.augment(x)
+
 
 class AugmentImageCoordinates(MapData):
     """
@@ -142,27 +141,30 @@ class AugmentImageCoordinates(MapData):
         else:
             self.augs = AugmentorList(augmentors)
 
-        exception_handler = ExceptionHandler(catch_exceptions)
+        self._img_index = img_index
+        self._coords_index = coords_index
+        self._copy = copy
+        self._exception_handler = ExceptionHandler(catch_exceptions)
 
-        def func(dp):
-            with exception_handler.catch():
-                img, coords = dp[img_index], dp[coords_index]
+        super(AugmentImageCoordinates, self).__init__(ds, self._aug_mapper)
+
+    def reset_state(self):
+        self.ds.reset_state()
+        self.augs.reset_state()
+
+    def _aug_mapper(self, dp):
+        with self._exception_handler.catch():
+            img, coords = dp[self._img_index], dp[self._coords_index]
             check_dtype(img)
             validate_coords(coords)
-                if copy:
+            if self._copy:
                 img, coords = copy_mod.deepcopy((img, coords))
             img, prms = self.augs._augment_return_params(img)
-                dp[img_index] = img
+            dp[self._img_index] = img
             coords = self.augs._augment_coords(coords, prms)
-                dp[coords_index] = coords
+            dp[self._coords_index] = coords
             return dp
 
-        super(AugmentImageCoordinates, self).__init__(ds, func)
-
-    def reset_state(self):
-        self.ds.reset_state()
-        self.augs.reset_state()
-
 
 class AugmentImageComponents(MapData):
     """

tensorpack/dataflow/parallel.py

@@ -170,7 +170,9 @@ class MultiProcessPrefetchData(ProxyDataFlow):
             nr_proc (int): number of processes to use.
         """
         if os.name == 'nt':
-            logger.warn("MultiProcessPrefetchData may not support windows!")
+            logger.warn("MultiProcessPrefetchData does support windows. \
+However, windows requires more strict picklability on processes, which may \
+lead of failure on some of the code.")
         super(MultiProcessPrefetchData, self).__init__(ds)
         try:
             self._size = ds.size()