add mmap option

It will be handy when debug or memory is limited
lshiwjx · May 22, 2018 · e102534 · e102534
1 parent d245bd1
commit e102534
Showing 1 changed file with 7 additions and 3 deletions.
diff --git a/st_gcn/feeder/feeder.py b/st_gcn/feeder/feeder.py
@@ -40,7 +40,8 @@ def __init__(self,
                  random_move=False,
                  window_size=-1,
                  normalization=False,
-                 debug=False):
+                 debug=False,
+                 mmap=False):
         self.debug = debug
         self.data_path = data_path
         self.label_path = label_path
@@ -75,8 +76,11 @@ def load_data(self):
             raise ValueError()
 
         # load data
-        self.data = np.load(self.data_path, mmap_mode='r')
-
+        if mmap:
+            self.data = np.load(self.data_path, mmap_mode='r')
+        else:
+            self.data = np.load(self.data_path)
+
         if self.debug:
             self.label = self.label[0:100]
             self.data = self.data[0:100]