Merge pull request #19 from delta-mpc/optimize

Optimize

Author: mh739025250

delta/dataset/dataset.py

@@ -31,8 +31,8 @@ def __init__(
 
 
 class FileDataset(TorchDataset):
-    def __init__(self, filename: str) -> None:
-        result = load_file(filename)
+    def __init__(self, filename: str, **kwargs) -> None:
+        result = load_file(filename, **kwargs)
         if isinstance(result, Image.Image):
             raise ValueError("file dataset does not support image file")
         self._result = result
@@ -50,9 +50,10 @@ def __len__(self) -> int:
 
 
 class DirectoryDataset(TorchDataset):
-    def __init__(self, directory: str) -> None:
+    def __init__(self, directory: str, **kwargs) -> None:
         self._xs = []
         self._ys = []
+        self._kwargs = kwargs
         root, dirnames, filenames = next(os.walk(directory))
         if len(filenames) > 0 and len(dirnames) == 0:
             self._xs.extend([os.path.join(root, filename) for filename in filenames])
@@ -77,7 +78,7 @@ def __init__(self, directory: str) -> None:
 
     def __getitem__(self, index):
         filename = self._xs[index]
-        x = load_file(filename)
+        x = load_file(filename, **self._kwargs)
         y = None
         if len(self._ys) > 0:
             y = self._ys[index]
@@ -121,12 +122,12 @@ def split_dataset(
 
 
 def load_dataset(
-    dataset_name: str,
+    dataset_name: str, **kwargs
 ) -> TorchDataset | Tuple[TorchDataset, TorchDataset]:
     if not os.path.exists(dataset_name):
         raise FileNotFoundError(errno.ENOENT, os.strerror(errno.ENOENT), dataset_name)
     if os.path.isfile(dataset_name):
-        dataset = FileDataset(dataset_name)
+        dataset = FileDataset(dataset_name, **kwargs)
         return dataset
     else:
         train_path = os.path.join(dataset_name, "train")
@@ -140,12 +141,12 @@ def load_dataset(
             train_root, _, train_files = next(os.walk(train_path))
             val_root, _, val_files = next(os.walk(val_path))
             if len(train_files) == 1 and len(val_files) == 1:
-                train_dataset = FileDataset(os.path.join(train_root, train_files[0]))
-                val_dataset = FileDataset(os.path.join(val_root, val_files[0]))
+                train_dataset = FileDataset(os.path.join(train_root, train_files[0]), **kwargs)
+                val_dataset = FileDataset(os.path.join(val_root, val_files[0]), **kwargs)
             else:
-                train_dataset = DirectoryDataset(train_path)
-                val_dataset = DirectoryDataset(val_path)
+                train_dataset = DirectoryDataset(train_path, **kwargs)
+                val_dataset = DirectoryDataset(val_path, **kwargs)
             return train_dataset, val_dataset
         else:
-            dataset = DirectoryDataset(dataset_name)
+            dataset = DirectoryDataset(dataset_name, **kwargs)
             return dataset

delta/dataset/file.py

@@ -39,6 +39,8 @@ def load_file(
         result = pd.read_csv(filename, sep=r"\s+", **kwargs)
     elif filename.endswith(".xls") or filename.endswith(".xlsx"):
         result = pd.read_excel(filename, **kwargs)
+    elif filename.endswith(".json"):
+        result = pd.read_json(filename, **kwargs)
     else:
         try:
             result = Image.open(filename, **kwargs)

delta/delta_node.py

@@ -23,10 +23,25 @@ def __init__(self, url: str) -> None:
 
     def create_task(self, task: Task) -> int:
         url = f"{self._url}/v1/task"
-        with TemporaryFile(mode="w+b") as file:
-            serialize.dump_task(file, task)
-            file.seek(0)
-            resp = httpx.post(url, files={"file": file}, timeout=None)
+        with TemporaryFile(mode="w+b") as task_file, TemporaryFile(mode="w+b") as config_file:
+            task_config = {
+                "name": task.name,
+                "dataset": task.dataset,
+                "type": task.type,
+                "enable_verify": task.enable_verify,
+                "options": task.options
+            }
+            pickle.dump(task_config, config_file)
+            config_file.seek(0)
+
+            serialize.dump_task(task_file, task)
+            task_file.seek(0)
+            files = {
+                "file": ("task_file.pkl", task_file, "application/octet-stream"),
+                "config": ("task_config_file.pkl", config_file, "application/pickle")
+            }
+
+            resp = httpx.post(url, files=files, timeout=None)
             resp.raise_for_status()
             data = resp.json()
             task_id = data["task_id"]

delta/pandas/dataframe.py

@@ -91,7 +91,7 @@ def map(self, data: pandas.DataFrame) -> pandas.Series:
 
     def _dispatch_binary_op(
         self,
-        other: "DataFrame" | "Series" | List[float] | float,
+        other: "DataFrame | Series | List[float] | float",
         op_name: str,
         op: Callable[..., Any],
         **kwargs: Any,

delta/task/learning.py

@@ -121,27 +121,55 @@ def __init__(
         self.epoch = epoch
         self.iteration = iteration
         self.strategy = strategy
+        self.batch_sampler = None
 
     def __iter__(self):
         return self._get_iter()
 
+    def _make_dataloader(self) -> DataLoader:
+        if self.batch_sampler is None:
+            return self.dataloader
+        else:
+            return DataLoader(
+                dataset=self.dataloader.dataset,
+                batch_sampler=self.batch_sampler,
+                num_workers=self.dataloader.num_workers,
+                collate_fn=self.dataloader.collate_fn,
+                pin_memory=self.dataloader.pin_memory,
+                timeout=self.dataloader.timeout,
+                worker_init_fn=self.dataloader.worker_init_fn,
+                multiprocessing_context=self.dataloader.multiprocessing_context,
+                generator=self.dataloader.generator,
+                prefetch_factor=self.dataloader.prefetch_factor,
+                persistent_workers=self.dataloader.persistent_workers,
+                pin_memory_device=self.dataloader.pin_memory_device,
+            )
+
     def _get_iter(self):
         finished = False
 
         while not finished:
-            for batch in self.dataloader:
+            count = 0
+            dataloader = self._make_dataloader()
+            for batch in dataloader:
                 if finished:
                     break
 
                 _logger.info(f"Training epoch {self.epoch} iteration {self.iteration}")
 
                 yield batch
-
+                
+                count += 1
                 if self.strategy.should_merge(self.epoch, self.iteration, False):
                     _logger.info(f"iteration {self.iteration}, start to merge")
+                    assert dataloader.batch_sampler is not None
+                    if self.batch_sampler is None:
+                        self.batch_sampler = list(dataloader.batch_sampler)
+                    self.batch_sampler = self.batch_sampler[count:]
                     finished = True
                 self.iteration += 1
-
+            
+            self.batch_sampler = None
             if self.strategy.should_merge(self.epoch, self.iteration, True):
                 _logger.info(f"epoch {self.epoch}, start to merge")
                 finished = True
@@ -372,9 +400,14 @@ def map(
                 epoch: int,
                 iteration: int,
             ) -> Tuple[Dict[str, np.ndarray], int, int]:
-                self.learning.strategy.weight_to_params(
-                    weight, self.learning.state_dict()
-                )
+                if len(weight) > 0:
+                    self.learning.strategy.weight_to_params(
+                        weight, self.learning.state_dict()
+                    )
+                else:
+                    weight = self.learning.strategy.params_to_weight(
+                        self.learning.state_dict()
+                    )
                 _logger.info(f"Round {self.round} training")
                 train_iter = TrainIterator(
                     dataloader, epoch, iteration, self.learning.strategy
@@ -517,7 +550,10 @@ def reduce(
                 self.learning.strategy.weight_to_params(
                     weight, self.learning.state_dict()
                 )
-                return self.learning.state_dict()
+                res: Dict[str, Any] = {"weight": self.learning.state_dict()}
+                if metrics is not None:
+                    res["metrics"] = metrics
+                return res
 
         input_nodes: List[DataNode] = [weight_node]
         if metrics_node is not None:
@@ -541,9 +577,8 @@ def _build_graph(self) -> Tuple[List[delta.dataset.Dataset], List[GraphNode]]:
         iteration_node = InputGraphNode(
             name="iteration", location=DataLocation.CLIENT, default=1
         )
-        weight_arr = self.strategy.params_to_weight(self.state_dict())
         weight_node = InputGraphNode(
-            name="weight_0", location=DataLocation.SERVER, default=weight_arr
+            name="weight_0", location=DataLocation.SERVER, default=np.empty(0)
         )
         metrics_node = None
         inputs = [dataset_node, epoch_node, iteration_node, weight_node]

setup.py

@@ -30,20 +30,20 @@ def run_tests(self):
 
 setup(
     name="delta-task",
-    version="0.8.3",
+    version="0.8.4rc1",
     license_files=("LICENSE"),
     packages=find_packages(),
     include_package_data=True,
     exclude_package_data={"": [".gitignore"]},
     install_requires=[
-        "cloudpickle==1.6.0",
-        "httpx==0.23.0",
-        "numpy==1.22.0",
-        "Pillow==9.1.1",
-        "pandas==1.2.3",
-        "pytest==6.2.5",
-        "torch==1.8.2+cpu",
-        "networkx==2.7.1"
+        "cloudpickle>=1.6.0",
+        "httpx>=0.23.0",
+        "numpy>=1.22.0",
+        "Pillow>=9.1.1",
+        "pandas>=1.2.3",
+        "pytest>=6.2.5",
+        "torch>=1.8.2",
+        "networkx>=2.7.1"
     ],
     tests_require=["pytest"],
     cmdclass={"test": PyTest},