update DataStrategy

Author: iiGray

loomtrain/core/data/dataloader/iter.py

@@ -54,15 +54,22 @@ def __next__(self):
     def __iter__(self):
         for epoch in range(self.num_epochs):
             self._current_epoch = epoch
-            if epoch < self.start_epoch:continue
-            self.load_sampler_state(epoch, self.start_epoch)
+            if epoch < self.consumed_epoch:continue
+            self.set_sampler_state(epoch, self.consumed_epoch, self.consumed_samples)
             yield from iter(super().__iter__())
-    
 
-    def set_state(self, start_epoch: int, consumed_samples = 0):
-        self.start_epoch = start_epoch
+
+
+
+    def set_sampler_state(self, current_epoch: int, consumed_epoch:int, consumed_samples):
+        raise NotImplementedError
+
+
+    def set_state(self, consumed_epoch: int, consumed_samples = 0):
+        self.consumed_epoch = consumed_epoch
         self.consumed_samples = consumed_samples
 
 
     def get_state(self) -> dict:
-        ...
+        raise NotImplementedError
+

loomtrain/core/datamodule.py

@@ -210,10 +210,10 @@ def _setup_dataset(self):
 
 
     def _setup_train_data_iter(self):
-        self.train_data_iter = self.strategy.setup_data_iter(self.train_dataset)
+        self.train_data_iter = self.strategy._setup_train_data_iter(self.train_dataset)
 
     def _setup_val_data_iter(self):
-        self.val_data_iter = self.strategy.setup_data_iter(self.val_dataset)
+        self.val_data_iter = self.strategy._setup_val_data_iter(self.val_dataset)
 
 
     def get_saved_sub_dir(self): return "data_iter"

loomtrain/core/strategy.py

@@ -63,10 +63,26 @@ def rank(self):
             self._rank = parallel.get_dp_rank()
         return self._rank
 
-    def setup_data_iter(self, 
-                        dataset: "tud.Dataset") -> "LoomDataIter":
-        raise NotImplementedError
-    
+    def _setup_train_data_iter(self, train_dataset: "tud.Dataset"):
+        self.train_data_iter = self.setup_data_iter(train_dataset)
+        return self.train_data_iter
+
+    def _setup_val_data_iter(self, val_dataset: "tud.Dataset"):
+        self.val_data_iter = self.setup_data_iter(val_dataset)
+        return self.val_data_iter
+
+
+    def config_loomDataModule_method(self, datamodule: "LoomDataModule"):
+        try:
+            self.loomDataModule_load_ckpt(None, None)
+        except NotImplementedError: ...
+        except Exception as e: 
+            datamodule.load_ckpt = self.loomDataModule_load_ckpt
+        try:
+            self.loomDataModule_save_ckpt(None, None)
+        except NotADirectoryError: ...
+        except Exception as e: 
+            datamodule.load_ckpt = self.loomDataModule_load_ckpt
 
 
     def loomDataModule_save_ckpt(self, save_dir: str, tag: str):
@@ -75,12 +91,11 @@ def loomDataModule_save_ckpt(self, save_dir: str, tag: str):
     def loomDataModule_load_ckpt(self, saved_dir: str, tag: str):
         raise NotImplementedError
 
+    def setup_data_iter(self, 
+                        dataset: "tud.Dataset") -> "LoomDataIter":
+        raise NotImplementedError
+
 
-    def config_loomDataModule_method(self, datamodule: "LoomDataModule"):
-        datamodule.save_ckpt = self.loomDataModule_save_ckpt
-        datamodule.load_ckpt = self.loomDataModule_load_ckpt
-        
-        
 
 # TBD
 class TrainStrategy: