spikezoo 0.2.3.4__py3-none-any.whl → 0.2.3.6__py3-none-any.whl

spikezoo/archs/yourmodel/arch/net.py

@@ -0,0 +1,35 @@
+import torch.nn as nn
+
+def conv_layer(inDim, outDim, ks, s, p, norm_layer="none"):
+    ## convolutional layer
+    conv = nn.Conv2d(inDim, outDim, kernel_size=ks, stride=s, padding=p)
+    relu = nn.ReLU(True)
+    assert norm_layer in ("batch", "instance", "none")
+    if norm_layer == "none":
+        seq = nn.Sequential(*[conv, relu])
+    else:
+        if norm_layer == "instance":
+            norm = nn.InstanceNorm2d(outDim, affine=False, track_running_stats=False)  # instance norm
+        else:
+            momentum = 0.1
+            norm = nn.BatchNorm2d(outDim, momentum=momentum, affine=True, track_running_stats=True)
+        seq = nn.Sequential(*[conv, norm, relu])
+    return seq
+
+
+class YourNet(nn.Module):
+    """Borrow the structure from the SpikeCLIP. (https://arxiv.org/abs/2501.04477)"""
+
+    def __init__(self, inDim=41):
+        super(YourNet, self).__init__()
+        norm = "none"
+        outDim = 1
+        convBlock1 = conv_layer(inDim, 64, 3, 1, 1)
+        convBlock2 = conv_layer(64, 128, 3, 1, 1, norm)
+        convBlock3 = conv_layer(128, 64, 3, 1, 1, norm)
+        convBlock4 = conv_layer(64, 16, 3, 1, 1, norm)
+        conv = nn.Conv2d(16, outDim, 3, 1, 1)
+        self.seq = nn.Sequential(*[convBlock1, convBlock2, convBlock3, convBlock4, conv])
+
+    def forward(self, x):
+        return self.seq(x)

spikezoo/datasets/__init__.py

@@ -1,4 +1,5 @@
 from spikezoo.datasets.base_dataset import BaseDataset, BaseDatasetConfig
+
 from dataclasses import replace
 import importlib, inspect
 import os
@@ -12,23 +13,24 @@ dataset_list = [file.replace("_dataset.py", "") for file in files_list if file.e
 
 
 # todo register function
-def build_dataset_cfg(cfg: BaseDatasetConfig, split: Literal["train", "test"] = "test"):
+def build_dataset_cfg(cfg: BaseDatasetConfig):
     """Build the dataset from the given dataset config."""
-    # build new cfg according to split
-    cfg = replace(cfg, split=split, spike_length=cfg.spike_length_train if split == "train" else cfg.spike_length_test)
     # dataset module
-    module_name = cfg.dataset_name + "_dataset"
-    assert cfg.dataset_name in dataset_list, f"Given dataset {cfg.dataset_name} not in our dataset list {dataset_list}."
-    module_name = "spikezoo.datasets." + module_name
-    module = importlib.import_module(module_name)
-    # dataset,dataset_config
-    dataset_name = cfg.dataset_name
-    dataset_name = dataset_name + "Dataset" if dataset_name == "base" else dataset_name
-    dataset_cls: BaseDataset = getattr_case_insensitive(module, dataset_name)
+    if cfg.dataset_cls_local == None:
+        module_name = cfg.dataset_name + "_dataset"
+        assert cfg.dataset_name in dataset_list, f"Given dataset {cfg.dataset_name} not in our dataset list {dataset_list}."
+        module_name = "spikezoo.datasets." + module_name
+        module = importlib.import_module(module_name)
+        # dataset,dataset_config
+        dataset_name = cfg.dataset_name
+        dataset_name = dataset_name + "Dataset" if dataset_name == "base" else dataset_name
+        dataset_cls: BaseDataset = getattr_case_insensitive(module, dataset_name)
+    else:
+        dataset_cls = cfg.dataset_cls_local
     dataset = dataset_cls(cfg)
     return dataset
 
-def build_dataset_name(dataset_name: str, split: Literal["train", "test"] = "test"):
+def build_dataset_name(dataset_name: str):
     """Build the default dataset from the given name."""
     module_name = dataset_name + "_dataset"
     assert dataset_name in dataset_list, f"Given dataset {dataset_name} not in our dataset list {dataset_list}."
@@ -37,22 +39,19 @@ def build_dataset_name(dataset_name: str, split: Literal["train", "test"] = "tes
     # dataset,dataset_config
     dataset_name = dataset_name + "Dataset" if dataset_name == "base" else dataset_name
     dataset_cls: BaseDataset = getattr_case_insensitive(module, dataset_name)
-    dataset_cfg: BaseDatasetConfig = getattr_case_insensitive(module, dataset_name + "config")(split=split)
+    dataset_cfg: BaseDatasetConfig = getattr_case_insensitive(module, dataset_name + "config")()
     dataset = dataset_cls(dataset_cfg)
     return dataset
 
 
 # todo to modify according to the basicsr
-def build_dataloader(dataset: BaseDataset, cfg=None):
+def build_dataloader(dataset, cfg):
     # train dataloader
-    if dataset.cfg.split == "train":
-        if cfg is None:
-            return torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False, num_workers=1)
-        else:
-            return torch.utils.data.DataLoader(dataset, batch_size=cfg.bs_train, shuffle=True, num_workers=cfg.num_workers, pin_memory=cfg.pin_memory)
+    if dataset.split == "train" and cfg._mode == "train_mode":
+        return torch.utils.data.DataLoader(dataset, batch_size=cfg.bs_train, shuffle=True, num_workers=cfg.nw_train, pin_memory=cfg.pin_memory)
     # test dataloader
-    elif dataset.cfg.split == "test":
-        return torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False, num_workers=1)
+    else:
+        return torch.utils.data.DataLoader(dataset, batch_size=cfg.bs_test, shuffle=False, num_workers=cfg.nw_test,pin_memory=False)
 
 
 # dataset_size_dict = {}

spikezoo/datasets/base_dataset.py

@@ -11,6 +11,7 @@ import warnings
 import torch
 from tqdm import tqdm
 from spikezoo.utils.data_utils import Augmentor
+from typing import Optional
 
 
 @dataclass
@@ -30,31 +31,24 @@ class BaseDatasetConfig:
     spike_length_train: int = -1
     "Dataset spike length for the test data."
     spike_length_test: int = -1
-    "Dataset spike length for the instantiation dataclass."
-    spike_length: int = -1
     "Dir name for the spike."
     spike_dir_name: str = "spike"
     "Dir name for the image."
     img_dir_name: str = "gt"
+    "Rate. (-1 denotes variant)"
+    rate: float = 0.6
 
     # ------------- Config -------------
-    "Dataset split: train/test. Default set as the 'test' for evaluation."
-    split: Literal["train", "test"] = "test"
     "Use the data augumentation technique or not."
     use_aug: bool = False
     "Use cache mechanism."
     use_cache: bool = False
     "Crop size."
     crop_size: tuple = (-1, -1)
-    "Rate. (-1 denotes variant)"
-    rate: float = 0.6
-
-    # post process
-    def __post_init__(self):
-        self.spike_length = self.spike_length_train if self.split == "train" else self.spike_length_test
-        self.root_dir = Path(self.root_dir) if isinstance(self.root_dir, str) else self.root_dir
-        # todo try download
-        assert self.root_dir.exists(), f"No files found in {self.root_dir} for the specified dataset `{self.dataset_name}`."
+    "Load the dataset from local or spikezoo lib."
+    dataset_cls_local: Optional[Dataset] = None
+    "Spike load version. [python,cpp]"
+    spike_load_version: Literal["python", "cpp"] = "python"
 
 
 # todo cache mechanism
@@ -62,10 +56,6 @@ class BaseDataset(Dataset):
     def __init__(self, cfg: BaseDatasetConfig):
         super(BaseDataset, self).__init__()
         self.cfg = cfg
-        self.augmentor = Augmentor(self.cfg.crop_size) if self.cfg.use_aug == True and self.cfg.split == "train" else -1
-        self.prepare_data()
-        self.cache_data() if cfg.use_cache == True else -1
-        warnings.warn("Lengths of the image list and the spike list should be equal.") if len(self.img_list) != len(self.spike_list) else -1
 
     def __len__(self):
         return len(self.spike_list)
@@ -80,7 +70,7 @@ class BaseDataset(Dataset):
             img = self.get_img(idx)
 
         # process data
-        if self.cfg.use_aug == True and self.cfg.split == "train":
+        if self.cfg.use_aug == True and self.split == "train":
             spike, img = self.augmentor(spike, img)
 
         # rate
@@ -90,15 +80,29 @@ class BaseDataset(Dataset):
         batch = {"spike": spike, "gt_img": img, "rate": rate}
         return batch
 
+    def build_source(self, split: Literal["train", "test"] = "test"):
+        """Build the dataset source and prepare to be loaded files."""
+        # spike length
+        self.split = split
+        self.spike_length = self.cfg.spike_length_train if self.split == "train" else self.cfg.spike_length_test
+        # root dir
+        self.cfg.root_dir = Path(self.cfg.root_dir) if isinstance(self.cfg.root_dir, str) else self.cfg.root_dir
+        assert self.cfg.root_dir.exists(), f"No files found in {self.cfg.root_dir} for the specified dataset `{self.cfg.dataset_name}`."
+        # prepare
+        self.augmentor = Augmentor(self.cfg.crop_size) if self.cfg.use_aug == True and self.split == "train" else -1
+        self.prepare_data()
+        self.cache_data() if self.cfg.use_cache == True else -1
+        warnings.warn("Lengths of the image list and the spike list should be equal.") if len(self.img_list) != len(self.spike_list) else -1
+
     # todo: To be overridden
     def prepare_data(self):
         """Specify the spike and image files to be loaded."""
         # spike
-        self.spike_dir = self.cfg.root_dir / self.cfg.split / self.cfg.spike_dir_name
+        self.spike_dir = self.cfg.root_dir / self.split / self.cfg.spike_dir_name
         self.spike_list = self.get_spike_files(self.spike_dir)
         # gt
         if self.cfg.with_img == True:
-            self.img_dir = self.cfg.root_dir / self.cfg.split / self.cfg.img_dir_name
+            self.img_dir = self.cfg.root_dir / self.split / self.cfg.img_dir_name
             self.img_list = self.get_image_files(self.img_dir)
 
     # todo: To be overridden
@@ -116,12 +120,13 @@ class BaseDataset(Dataset):
             height=self.cfg.height,
             width=self.cfg.width,
             out_format="tensor",
+            version=self.cfg.spike_load_version
         )
         return spike
 
     def get_spike(self, idx):
         """Get and process the spike stream from the given idx."""
-        spike_length = self.cfg.spike_length
+        spike_length = self.spike_length
         spike = self.load_spike(idx)
         assert spike.shape[0] >= spike_length, f"Given spike length {spike.shape[0]} smaller than the required length {spike_length}"
         spike_mid = spike.shape[0] // 2

spikezoo/datasets/{realworld_dataset.py → realdata_dataset.py}

@@ -4,21 +4,19 @@ from dataclasses import dataclass
 
 
 @dataclass
-class RealWorldConfig(BaseDatasetConfig):
-    dataset_name: str = "realworld"
-    root_dir: Path = Path(__file__).parent.parent / Path("data/recVidarReal2019")
+class RealDataConfig(BaseDatasetConfig):
+    dataset_name: str = "realdata"
+    root_dir: Path = Path(__file__).parent.parent / Path("data/realdata")
     width: int = 400
     height: int = 250
     with_img: bool = False
     spike_length_train: int = -1
     spike_length_test: int = -1
     rate: float = 1
-    
 
-
-class RealWorld(BaseDataset):
+class RealData(BaseDataset):
     def __init__(self, cfg: BaseDatasetConfig):
-        super(RealWorld, self).__init__(cfg)
+        super(RealData, self).__init__(cfg)
 
     def prepare_data(self):
         self.spike_dir = self.cfg.root_dir

spikezoo/datasets/reds_base_dataset.py

@@ -8,7 +8,7 @@ import re
 @dataclass
 class REDS_BASEConfig(BaseDatasetConfig):
     dataset_name: str = "reds_base"
-    root_dir: Path = Path(__file__).parent.parent / Path("data/REDS_BASE")
+    root_dir: Path = Path(__file__).parent.parent / Path("data/reds_base")
     width: int = 400
     height: int = 250
     with_img: bool = True

spikezoo/datasets/szdata_dataset.py

@@ -9,7 +9,7 @@ import numpy as np
 @dataclass
 class SZDataConfig(BaseDatasetConfig):
     dataset_name: str = "szdata"
-    root_dir: Path = Path(__file__).parent.parent / Path("data/dataset")
+    root_dir: Path = Path(__file__).parent.parent / Path("data/szdata")
     width: int = 400
     height: int = 250
     with_img: bool = True
@@ -22,7 +22,3 @@ class SZDataConfig(BaseDatasetConfig):
 class SZData(BaseDataset):
     def __init__(self, cfg: BaseDatasetConfig):
         super(SZData, self).__init__(cfg)
-
-    def prepare_data(self):
-        super().prepare_data()
-        self.img_list = [self.img_dir / Path(str(s.name).replace('.dat','.png')) for s in self.spike_list]

spikezoo/datasets/uhsr_dataset.py

@@ -7,7 +7,7 @@ import torch
 @dataclass
 class UHSRConfig(BaseDatasetConfig):
     dataset_name: str = "uhsr"
-    root_dir: Path = Path(__file__).parent.parent / Path("data/U-CALTECH")
+    root_dir: Path = Path(__file__).parent.parent / Path("data/u_caltech")
     width: int = 224
     height: int = 224
     with_img: bool = False

spikezoo/datasets/yourdataset_dataset.py

@@ -0,0 +1,23 @@
+from torch.utils.data import Dataset
+from pathlib import Path
+from dataclasses import dataclass
+from typing import Literal, Union
+from typing import Optional
+from spikezoo.datasets.base_dataset import BaseDatasetConfig,BaseDataset
+
+@dataclass
+class YourDatasetConfig(BaseDatasetConfig):
+    dataset_name: str = "yourdataset"
+    root_dir: Union[str, Path] = Path(__file__).parent.parent / Path("data/your_data_path")
+    width: int = 400
+    height: int = 250
+    with_img: bool = True
+    spike_length_train: int = -1
+    spike_length_test: int = -1
+    spike_dir_name: str = "spike_data"
+    img_dir_name: str = "sharp_data"
+    rate: float = 1
+
+class YourDataset(BaseDataset):
+    def __init__(self, cfg: BaseDatasetConfig):
+        super(YourDataset, self).__init__(cfg)

spikezoo/models/__init__.py

@@ -1,6 +1,7 @@
 import importlib
 import inspect
 from spikezoo.models.base_model import BaseModel,BaseModelConfig
+
 from spikezoo.utils.other_utils import getattr_case_insensitive
 import os
 from pathlib import Path
@@ -13,14 +14,17 @@ model_list = [file.split("_")[0] for file in files_list if file.endswith("_model
 def build_model_cfg(cfg: BaseModelConfig):
     """Build the model from the given model config."""
     # model module name
-    module_name = cfg.model_name + "_model"
-    assert cfg.model_name in model_list, f"Given model {cfg.model_name} not in our model zoo {model_list}."
-    module_name = "spikezoo.models." + module_name
-    module = importlib.import_module(module_name)
-    # model,model_config
-    model_name = cfg.model_name
-    model_name = model_name + 'Model' if model_name == "base" else model_name
-    model_cls: BaseModel = getattr_case_insensitive(module,model_name)
+    if cfg.model_cls_local == None:
+        module_name = cfg.model_name + "_model"
+        assert cfg.model_name in model_list, f"Given model {cfg.model_name} not in our model zoo {model_list}."
+        module_name = "spikezoo.models." + module_name
+        module = importlib.import_module(module_name)
+        # model,model_config
+        model_name = cfg.model_name
+        model_name = model_name + 'Model' if model_name == "base" else model_name
+        model_cls: BaseModel = getattr_case_insensitive(module,model_name)
+    else:
+        model_cls: BaseModel = cfg.model_cls_local
     model = model_cls(cfg)
     return model
 

spikezoo/models/base_model.py

@@ -44,7 +44,10 @@ class BaseModelConfig:
     multi_gpu: bool = False
     "Base url."
     base_url: str = "https://github.com/chenkang455/Spike-Zoo/releases/download"
-
+    "Load the model from local class or spikezoo lib. (None)"
+    model_cls_local: Optional[nn.Module] = None
+    "Load the arch from local class or spikezoo lib. (None)"
+    arch_cls_local: Optional[nn.Module] = None
 
 class BaseModel(nn.Module):
     def __init__(self, cfg: BaseModelConfig):
@@ -71,8 +74,11 @@ class BaseModel(nn.Module):
     ):
         """Build the network and load the pretrained weight."""
         # network
-        module = importlib.import_module(f"spikezoo.archs.{self.cfg.model_name}.{self.cfg.model_file_name}")
-        model_cls = getattr(module, self.cfg.model_cls_name)
+        if self.cfg.arch_cls_local == None:
+            module = importlib.import_module(f"spikezoo.archs.{self.cfg.model_name}.{self.cfg.model_file_name}")
+            model_cls = getattr(module, self.cfg.model_cls_name)
+        else:
+            model_cls = self.cfg.arch_cls_local
         # load model config parameters
         if version == "local":
             model = model_cls(**self.cfg.model_params)
@@ -129,7 +135,7 @@ class BaseModel(nn.Module):
         """Crop the spike length."""
         spike_length = spike.shape[1]
         spike_mid = spike_length // 2
-        assert spike_length >= self.model_length, f"Spike input is not long enough, given {spike_length} frames < {self.cfg.model_length}."
+        assert spike_length >= self.model_length, f"Spike input is not long enough, given {spike_length} frames < {self.cfg.model_length} required by the {self.cfg.model_name}."
         # even length
         if self.model_length == self.model_half_length * 2:
             spike = spike[

spikezoo/models/bsf_model.py

@@ -4,7 +4,6 @@ from spikezoo.models.base_model import BaseModel, BaseModelConfig
 from torch.optim import Adam
 import torch.optim.lr_scheduler as lr_scheduler
 import torch.nn as nn
-from spikezoo.pipeline import TrainPipelineConfig
 from typing import List
 from spikezoo.archs.bsf.models.bsf.bsf import BSF 
 

spikezoo/models/spk2imgnet_model.py

@@ -1,7 +1,6 @@
 import torch
 from dataclasses import dataclass, field
 from spikezoo.models.base_model import BaseModel, BaseModelConfig
-from spikezoo.pipeline import TrainPipelineConfig
 import torch.nn as nn
 import torch.optim as optim
 import torch.optim.lr_scheduler as lr_scheduler

spikezoo/models/stir_model.py

@@ -4,7 +4,6 @@ from spikezoo.models.base_model import BaseModel, BaseModelConfig
 from torch.optim import Adam
 import torch.optim.lr_scheduler as lr_scheduler
 import torch.nn as nn
-from spikezoo.pipeline import TrainPipelineConfig
 from typing import List
 from spikezoo.archs.stir.metrics.losses import compute_per_loss_single
 from spikezoo.archs.stir.models.Vgg19 import Vgg19

spikezoo/models/wgse_model.py

@@ -1,7 +1,6 @@
 from dataclasses import dataclass, field
 from spikezoo.models.base_model import BaseModel, BaseModelConfig
 from typing import List
-from spikezoo.pipeline import TrainPipelineConfig
 import torch.nn as nn
 import torch.optim as optim
 import torch.optim.lr_scheduler as lr_scheduler

spikezoo/models/yourmodel_model.py

@@ -0,0 +1,22 @@
+from torch.utils.data import Dataset
+from pathlib import Path
+from dataclasses import dataclass
+from typing import Literal, Union
+from typing import Optional
+from spikezoo.models.base_model import BaseModel, BaseModelConfig
+from dataclasses import field
+import torch.nn as nn
+
+
+@dataclass
+class YourModelConfig(BaseModelConfig):
+    model_name: str = "yourmodel"  # 需与文件名保持一致
+    model_file_name: str = "arch.net"  # archs路径下的模块路径
+    model_cls_name: str = "YourNet"  # 模型类名
+    model_length: int = 41
+    require_params: bool = True
+    model_params: dict = field(default_factory=lambda: {"inDim": 41})
+
+class YourModel(BaseModel):
+    def __init__(self, cfg: BaseModelConfig):
+        super(YourModel, self).__init__(cfg)

spikezoo/pipeline/base_pipeline.py

@@ -34,11 +34,17 @@ class PipelineConfig:
     "Evaluate metrics or not."
     save_metric: bool = True
     "Metric names for evaluation."
-    metric_names: List[str] = field(default_factory=lambda: ["psnr", "ssim"])
+    metric_names: List[str] = field(default_factory=lambda: ["psnr", "ssim", "niqe", "brisque"])
     "Save recoverd images or not."
     save_img: bool = True
     "Normalizing recoverd images and gt or not."
-    save_img_norm: bool = False
+    img_norm: bool = False
+    "Batch size for the test dataloader."
+    bs_test: int = 1
+    "Num_workers for the test dataloader."
+    nw_test: int = 0
+    "Pin_memory true or false for the dataloader."
+    pin_memory: bool = False
     "Different modes for the pipeline."
     _mode: Literal["single_mode", "multi_mode", "train_mode"] = "single_mode"
 
@@ -63,7 +69,8 @@ class Pipeline:
         torch.set_grad_enabled(False)
         # dataset
         self.dataset: BaseDataset = build_dataset_name(dataset_cfg) if isinstance(dataset_cfg, str) else build_dataset_cfg(dataset_cfg)
-        self.dataloader = build_dataloader(self.dataset)
+        self.dataset.build_source(split="test")
+        self.dataloader = build_dataloader(self.dataset,self.cfg)
         # device
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -103,7 +110,7 @@ class Pipeline:
         """Function I---Save the recoverd image and calculate the metric from the given dataset."""
         # save folder
         self.logger.info("*********************** infer_from_dataset ***********************")
-        save_folder = self.save_folder / Path(f"infer_from_dataset/{self.dataset.cfg.dataset_name}_dataset/{self.dataset.cfg.split}/{idx:06d}")
+        save_folder = self.save_folder / Path(f"infer_from_dataset/{self.dataset.cfg.dataset_name}_dataset/{self.dataset.split}/{idx:06d}")
         os.makedirs(str(save_folder), exist_ok=True)
 
         # data process
@@ -117,7 +124,7 @@ class Pipeline:
             img = None
         return self.infer(spike, img, save_folder, rate)
 
-    def infer_from_file(self, file_path, height=-1, width=-1, img_path=None, rate=1, remove_head=False):
+    def infer_from_file(self, file_path, height=-1, width=-1, rate=1, img_path=None, remove_head=False):
         """Function II---Save the recoverd image and calculate the metric from the given input file."""
         # save folder
         self.logger.info("*********************** infer_from_file ***********************")
@@ -144,7 +151,7 @@ class Pipeline:
         spike = torch.from_numpy(spike)[None].to(self.device)
         return self.infer(spike, img, save_folder, rate)
 
-    def infer_from_spk(self, spike, img=None, rate=1):
+    def infer_from_spk(self, spike, rate=1, img=None):
         """Function III---Save the recoverd image and calculate the metric from the given spike stream."""
         # save folder
         self.logger.info("*********************** infer_from_spk ***********************")
@@ -181,7 +188,7 @@ class Pipeline:
         for idx in range(len(self.dataset)):
             self.infer_from_dataset(idx=idx)
         self.cfg.save_metric = base_setting
-        
+
     # TODO: To be overridden
     def cal_params(self):
         """Function VI---Calculate the parameters/flops/latency of the given method."""
@@ -228,8 +235,8 @@ class Pipeline:
         # With no GT
         if recon_img == None:
             return None
-        # TFP, TFI, spikeclip algorithms are normalized automatically, others are normalized based on the self.cfg.use_norm
-        if model_name in ["tfp", "tfi", "spikeclip"] or self.cfg.save_img_norm == True:
+        # spikeclip is normalized automatically
+        if model_name in ["spikeclip"] or self.cfg.img_norm == True:
             recon_img = (recon_img - recon_img.min()) / (recon_img.max() - recon_img.min())
         else:
             recon_img = recon_img / rate
@@ -253,7 +260,7 @@ class Pipeline:
             batch = model.feed_to_device(batch)
             outputs = model.get_outputs_dict(batch)
             recon_img, img = model.get_paired_imgs(batch, outputs)
-            recon_img, img = self._post_process_img(recon_img, model_name), self._post_process_img(img, "auto")
+            recon_img, img = self._post_process_img(recon_img, model_name), self._post_process_img(img, "gt")
             for metric_name in metrics_dict.keys():
                 if metric_name in metric_pair_names:
                     metrics_dict[metric_name].update(cal_metric_pair(recon_img, img, metric_name))

spikezoo/pipeline/ensemble_pipeline.py

@@ -48,7 +48,8 @@ class EnsemblePipeline(Pipeline):
         torch.set_grad_enabled(False)
         # data
         self.dataset: BaseDataset = build_dataset_name(dataset_cfg) if isinstance(dataset_cfg, str) else build_dataset_cfg(dataset_cfg)
-        self.dataloader = build_dataloader(self.dataset)
+        self.dataset.build_source(split = "test")
+        self.dataloader = build_dataloader(self.dataset,self.cfg)
         # device
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
 

spikezoo/pipeline/train_cfgs.py

@@ -22,7 +22,7 @@ class REDS_BASE_TrainConfig(TrainPipelineConfig):
 
     # dataloader setting
     bs_train: int = 8
-    num_workers: int = 4
+    nw_train: int = 4
     pin_memory: bool = False
 
     # train setting - optimizer & scheduler & loss_dict
@@ -30,6 +30,8 @@ class REDS_BASE_TrainConfig(TrainPipelineConfig):
     scheduler_cfg: Optional[SchedulerConfig] = MultiStepSchedulerConfig(milestones=[400], gamma=0.2) # from wgse
     loss_weight_dict: Dict = field(default_factory=lambda: {"l1": 1})
 
+
+
 # ! Train Config for each method on the official setting, not recommended to utilize their default parameters owing to the dataset setting.
 @dataclass
 class BSFTrainConfig(TrainPipelineConfig):