refine weight parameter

Author: HydrogenSulfate

ppsci/data/__init__.py

@@ -83,10 +83,9 @@ def build_dataloader(_dataset, cfg):
     # build collate_fn if specified
     batch_transforms_cfg = cfg.pop("batch_transforms", None)
 
+    collate_fn = None
     if isinstance(batch_transforms_cfg, dict) and batch_transforms_cfg:
         collate_fn = batch_transform.build_batch_transforms(batch_transforms_cfg)
-    else:
-        collate_fn = batch_transform.default_collate_fn_allow_none
 
     # build init function
     init_fn = partial(
@@ -97,7 +96,7 @@ def build_dataloader(_dataset, cfg):
     )
 
     # build dataloader
-    dataloader = io.DataLoader(
+    dataloader_ = io.DataLoader(
         dataset=_dataset,
         places=device.get_device(),
         batch_sampler=sampler,
@@ -107,4 +106,4 @@ def build_dataloader(_dataset, cfg):
         worker_init_fn=init_fn,
     )
 
-    return dataloader
+    return dataloader_

ppsci/data/dataset/array_dataset.py

@@ -56,11 +56,7 @@ def __init__(
     def __getitem__(self, idx):
         input_item = {key: value[idx] for key, value in self.input.items()}
         label_item = {key: value[idx] for key, value in self.label.items()}
-        weight_item = (
-            {key: value[idx] for key, value in self.weight.items()}
-            if self.weight is not None
-            else None
-        )
+        weight_item = {key: value[idx] for key, value in self.weight.items()}
 
         # TODO(sensen): Transforms may be applied on label and weight.
         if self.transforms is not None:

ppsci/data/dataset/era5_dataset.py

@@ -68,14 +68,12 @@ def __init__(
         self.label_keys = label_keys
         self.precip_file_path = precip_file_path
 
-        self.weight_dict = weight_dict
+        self.weight_dict = {} if weight_dict is None else weight_dict
         if weight_dict is not None:
             self.weight_dict = {key: 1.0 for key in self.label_keys}
             self.weight_dict.update(weight_dict)
 
-        self.vars_channel = (
-            vars_channel if vars_channel is not None else [i for i in range(20)]
-        )
+        self.vars_channel = list(range(20)) if vars_channel is None else vars_channel
         self.num_label_timestamps = num_label_timestamps
         self.transforms = transforms
         self.training = training
@@ -127,6 +125,7 @@ def __getitem__(self, global_idx):
             input_idx, label_idx = local_idx, local_idx + step
 
         input_item = {self.input_keys[0]: input_file[input_idx, self.vars_channel]}
+
         label_item = {}
         for i in range(self.num_label_timestamps):
             if self.precip_file_path is not None:
@@ -138,14 +137,11 @@ def __getitem__(self, global_idx):
                     label_idx + i, self.vars_channel
                 ]
 
-        if self.weight_dict is not None:
-            weight_shape = [1] * len(next(iter(label_item.values)).shape)
-            weight_item = {
-                key: np.full(weight_shape, value, paddle.get_default_dtype())
-                for key, value in self.weight_dict.items()
-            }
-        else:
-            weight_item = None
+        weight_shape = [1] * len(next(iter(label_item.values)).shape)
+        weight_item = {
+            key: np.full(weight_shape, value, paddle.get_default_dtype())
+            for key, value in self.weight_dict.items()
+        }
 
         if self.transforms is not None:
             input_item, label_item, weight_item = self.transforms(
@@ -187,7 +183,7 @@ def __init__(
         self.input_keys = input_keys
         self.label_keys = label_keys
 
-        self.weight_dict = weight_dict
+        self.weight_dict = {} if weight_dict is None else weight_dict
         if weight_dict is not None:
             self.weight_dict = {key: 1.0 for key in self.label_keys}
             self.weight_dict.update(weight_dict)
@@ -201,8 +197,8 @@ def read_data(self, path: str):
         paths = glob.glob(path + "/*.h5")
         paths.sort()
         files = []
-        for path in paths:
-            _file = h5py.File(path, "r")
+        for _path in paths:
+            _file = h5py.File(_path, "r")
             files.append(_file)
         return files
 
@@ -217,20 +213,18 @@ def __getitem__(self, global_idx):
             input_item[key] = np.asarray(
                 _file["input_dict"][key], paddle.get_default_dtype()
             )
+
         label_item = {}
         for key in _file["label_dict"]:
             label_item[key] = np.asarray(
                 _file["label_dict"][key], paddle.get_default_dtype()
             )
 
-        if self.weight_dict is not None:
-            weight_shape = [1] * len(next(iter(label_item.values)).shape)
-            weight_item = {
-                key: np.full(weight_shape, value, paddle.get_default_dtype())
-                for key, value in self.weight_dict.items()
-            }
-        else:
-            weight_item = None
+        weight_shape = [1] * len(next(iter(label_item.values)).shape)
+        weight_item = {
+            key: np.full(weight_shape, value, paddle.get_default_dtype())
+            for key, value in self.weight_dict.items()
+        }
 
         if self.transforms is not None:
             input_item, label_item, weight_item = self.transforms(

ppsci/data/dataset/vtu_dataset.py

@@ -83,7 +83,7 @@ def __init__(
     def __getitem__(self, idx):
         input_item = {key: value[idx] for key, value in self.input.items()}
         label_item = {key: value[idx] for key, value in self.label.items()}
-        return (input_item, label_item, None)
+        return (input_item, label_item, {})
 
     def __len__(self):
         return self.num_samples

ppsci/data/process/batch_transform/__init__.py

@@ -25,13 +25,13 @@
 
 from ppsci.data.process import transform
 
-__all__ = ["build_batch_transforms", "default_collate_fn_allow_none"]
+__all__ = ["build_batch_transforms"]
 
 
-def default_collate_fn_allow_none(batch: List[Any]) -> Any:
-    """Modified collate function to allow some fields to be None, such as weight field.
+def default_collate_fn(batch: List[Any]) -> Any:
+    """Default_collate_fn for paddle dataloader.
 
-    ref: https://github.com/PaddlePaddle/Paddle/blob/develop/python/paddle/fluid/dataloader/collate.py#L24
+    ref: https://github.com/PaddlePaddle/Paddle/blob/develop/python/paddle/io/dataloader/collate.py#L25
 
     Args:
         batch (List[Any]): Batch of samples to be collated.
@@ -40,11 +40,6 @@ def default_collate_fn_allow_none(batch: List[Any]) -> Any:
         Any: Collated batch data.
     """
     sample = batch[0]
-
-    # allow field to be None
-    if sample is None:
-        return None
-
     if isinstance(sample, np.ndarray):
         batch = np.stack(batch, axis=0)
         return batch
@@ -56,15 +51,12 @@ def default_collate_fn_allow_none(batch: List[Any]) -> Any:
     elif isinstance(sample, (str, bytes)):
         return batch
     elif isinstance(sample, Mapping):
-        return {
-            key: default_collate_fn_allow_none([d[key] for d in batch])
-            for key in sample
-        }
+        return {key: default_collate_fn([d[key] for d in batch]) for key in sample}
     elif isinstance(sample, Sequence):
         sample_fields_num = len(sample)
         if not all(len(sample) == sample_fields_num for sample in iter(batch)):
             raise RuntimeError("fileds number not same among samples in a batch")
-        return [default_collate_fn_allow_none(fields) for fields in zip(*batch)]
+        return [default_collate_fn(fields) for fields in zip(*batch)]
 
     raise TypeError(
         "batch data can only contains: tensor, numpy.ndarray, "
@@ -80,6 +72,6 @@ def collate_fn_batch_transforms(batch: List[Any]):
         # apply batch transform on uncollated data
         batch = batch_transforms(batch)
         # then do collate
-        return default_collate_fn_allow_none(batch)
+        return default_collate_fn(batch)
 
     return collate_fn_batch_transforms