spikezoo 0.2.2__py3-none-any.whl → 0.2.3__py3-none-any.whl

spikezoo/models/wgse_model.py

@@ -1,6 +1,12 @@
 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
+from typing import List
+from spikezoo.archs.wgse.dwtnets import Dwt1dResnetX_TCN
 
 
 @dataclass
@@ -9,21 +15,9 @@ class WGSEConfig(BaseModelConfig):
     model_name: str = "wgse"
     model_file_name: str = "dwtnets"
     model_cls_name: str = "Dwt1dResnetX_TCN"
-    model_win_length: int = 41
+    model_length: int = 41
+    model_length_dict: dict = field(default_factory=lambda: {"v010": 41, "v023": 41})
     require_params: bool = True
-    ckpt_path: str = "weights/wgse.pt"
-    model_params: dict = field(
-        default_factory=lambda: {
-            "wvlname": "db8",
-            "J": 5,
-            "yl_size": "15",
-            "yh_size": [28, 21, 18, 16, 15],
-            "num_residual_blocks": 3,
-            "norm": None,
-            "ks": 3,
-            "store_features": True,
-        }
-    )
 
 
 class WGSE(BaseModel):

spikezoo/pipeline/base_pipeline.py

@@ -18,24 +18,27 @@ from tqdm import tqdm
 from spikezoo.models import build_model_cfg, build_model_name, BaseModel, BaseModelConfig
 from spikezoo.datasets import build_dataset_cfg, build_dataset_name, BaseDataset, BaseDatasetConfig, build_dataloader
 from typing import Optional, Union, List
+import shutil
+from spikingjelly.clock_driven import functional
+import spikezoo as sz
 
 
 @dataclass
 class PipelineConfig:
-    "Evaluate metrics or not."
-    save_metric: bool = True
-    "Save recoverd images or not."
-    save_img: bool = True
-    "Normalizing recoverd images or not."
-    save_img_norm: bool = False
-    "Normalizing gt or not."
-    gt_img_norm: bool = False
+    "Loading weights from local or version on the url."
+    version: Literal["local", "v010", "v023"] = "local"
     "Save folder for the code running result."
     save_folder: str = ""
     "Saved experiment name."
     exp_name: str = ""
+    "Evaluate metrics or not."
+    save_metric: bool = True
     "Metric names for evaluation."
     metric_names: List[str] = field(default_factory=lambda: ["psnr", "ssim"])
+    "Save recoverd images or not."
+    save_img: bool = True
+    "Normalizing recoverd images and gt or not."
+    save_img_norm: bool = False
     "Different modes for the pipeline."
     _mode: Literal["single_mode", "multi_mode", "train_mode"] = "single_mode"
 
@@ -44,8 +47,8 @@ class Pipeline:
     def __init__(
         self,
         cfg: PipelineConfig,
-        model_cfg: Union[str, BaseModelConfig],
-        dataset_cfg: Union[str, BaseDatasetConfig],
+        model_cfg: Union[sz.METHOD, BaseModelConfig],
+        dataset_cfg: Union[sz.DATASET, BaseDatasetConfig],
     ):
         self.cfg = cfg
         self._setup_model_data(model_cfg, dataset_cfg)
@@ -53,9 +56,9 @@ class Pipeline:
 
     def _setup_model_data(self, model_cfg, dataset_cfg):
         """Model and Data setup."""
-        # model
+        # model [1] build the model. [2] build the network.
         self.model: BaseModel = build_model_name(model_cfg) if isinstance(model_cfg, str) else build_model_cfg(model_cfg)
-        self.model = self.model.eval()
+        self.model.build_network(mode="eval", version=self.cfg.version)
         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)
@@ -74,15 +77,19 @@ class Pipeline:
             if len(self.cfg.exp_name) == 0
             else self.save_folder / Path(f"{mode_name}/{self.cfg.exp_name}")
         )
-        save_folder = self.save_folder
-        os.makedirs(str(save_folder), exist_ok=True)
+        # remove and establish folder
+        save_folder = str(self.save_folder)
+        if os.path.exists(save_folder):
+            shutil.rmtree(save_folder)
+        os.makedirs(save_folder)
+        save_folder = Path(save_folder)
         # logger result
         self.logger = setup_logging(save_folder / Path("result.log"))
         self.logger.info(f"Info logs are saved on the {save_folder}/result.log")
         # pipeline config
         save_config(self.cfg, save_folder / Path("cfg_pipeline.log"))
         # model config
-        if self.cfg._mode == "single_mode":
+        if self.cfg._mode in ["single_mode", "train_mode"]:
             save_config(self.model.cfg, save_folder / Path("cfg_model.log"))
         elif self.cfg._mode == "multi_mode":
             for model in self.model_list:
@@ -90,28 +97,29 @@ class Pipeline:
         # dataset config
         save_config(self.dataset.cfg, save_folder / Path("cfg_dataset.log"))
 
-    def spk2img_from_dataset(self, idx=0):
-        """Func---Save the recoverd image and calculate the metric from the given dataset."""
+    def infer_from_dataset(self, idx=0):
+        """Function I---Save the recoverd image and calculate the metric from the given dataset."""
         # save folder
-        self.logger.info("*********************** spk2img_from_dataset ***********************")
-        save_folder = self.save_folder / Path(f"spk2img_from_dataset/{self.dataset.cfg.dataset_name}_dataset/{self.dataset.cfg.split}/{idx:06d}")
+        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}")
         os.makedirs(str(save_folder), exist_ok=True)
 
         # data process
+        # todo
         batch = self.dataset[idx]
-        spike, img = batch["spike"], batch["img"]
+        spike, img, rate = batch["spike"], batch["gt_img"], batch["rate"]
         spike = spike[None].to(self.device)
         if self.dataset.cfg.with_img == True:
             img = img[None].to(self.device)
         else:
             img = None
-        return self._spk2img(spike, img, save_folder)
+        return self.infer(spike, img, save_folder, rate)
 
-    def spk2img_from_file(self, file_path, height = -1, width  = -1, img_path=None, remove_head=False):
-        """Func---Save the recoverd image and calculate the metric from the given input file."""
+    def infer_from_file(self, file_path, height=-1, width=-1, img_path=None, rate=1, remove_head=False):
+        """Function II---Save the recoverd image and calculate the metric from the given input file."""
         # save folder
-        self.logger.info("*********************** spk2img_from_file ***********************")
-        save_folder = self.save_folder / Path(f"spk2img_from_file/{os.path.basename(file_path)}")
+        self.logger.info("*********************** infer_from_file ***********************")
+        save_folder = self.save_folder / Path(f"infer_from_file/{os.path.basename(file_path)}")
         os.makedirs(str(save_folder), exist_ok=True)
 
         # load spike from .dat
@@ -132,13 +140,13 @@ class Pipeline:
         else:
             img = img_path
         spike = torch.from_numpy(spike)[None].to(self.device)
-        return self._spk2img(spike, img, save_folder)
+        return self.infer(spike, img, save_folder, rate)
 
-    def spk2img_from_spk(self, spike, img=None):
-        """Func---Save the recoverd image and calculate the metric from the given spike stream."""
+    def infer_from_spk(self, spike, img=None, rate=1):
+        """Function III---Save the recoverd image and calculate the metric from the given spike stream."""
         # save folder
-        self.logger.info("*********************** spk2img_from_spk ***********************")
-        save_folder = self.save_folder / Path(f"spk2img_from_spk/{self.thistime}")
+        self.logger.info("*********************** infer_from_spk ***********************")
+        save_folder = self.save_folder / Path(f"infer_from_spk")
         os.makedirs(str(save_folder), exist_ok=True)
 
         # spike process
@@ -157,36 +165,40 @@ class Pipeline:
                 img = torch.from_numpy(img)[None, None].to(self.device)
             else:
                 raise RuntimeError("Not recognized image input type.")
-        return self._spk2img(spike, img, save_folder)
+        return self.infer(spike, img, save_folder, rate)
+
+    # TODO: To be overridden
+    def infer(self, spike, img, save_folder, rate):
+        """Function IV---Spike-to-image conversion interface, input data format: spike [bs,c,h,w] (0-1), img [bs,1,h,w] (0-1)"""
+        return self._infer_model(self.model, spike, img, save_folder, rate)
 
     def save_imgs_from_dataset(self):
-        """Func---Save all images from the given dataset."""
+        """Function V---Save all images from the given dataset."""
+        base_setting = self.cfg.save_metric
+        self.cfg.save_metric = False
         for idx in range(len(self.dataset)):
-            self.spk2img_from_dataset(idx=idx)
-
+            self.infer_from_dataset(idx=idx)
+        self.cfg.save_metric = base_setting
+        
     # TODO: To be overridden
     def cal_params(self):
-        """Func---Calculate the parameters/flops/latency of the given method."""
+        """Function VI---Calculate the parameters/flops/latency of the given method."""
         self._cal_prams_model(self.model)
 
     # TODO: To be overridden
     def cal_metrics(self):
-        """Func---Calculate the metric of the given method."""
-        self._cal_metrics_model(self.model)
-
-    # TODO: To be overridden
-    def _spk2img(self, spike, img, save_folder):
-        """Spike-to-image: spike:[bs,c,h,w] (0-1), img:[bs,1,h,w] (0-1)"""
-        return self._spk2img_model(self.model, spike, img, save_folder)
+        """Function VII---Calculate the metric of the given method."""
+        return self._cal_metrics_model(self.model)
 
-    def _spk2img_model(self, model, spike, img, save_folder):
+    def _infer_model(self, model, spike, img, save_folder, rate):
         """Spike-to-image from the given model."""
         # spike2image conversion
         model_name = model.cfg.model_name
-        recon_img = model(spike)
+        recon_img = model.spk2img(spike)
         recon_img_copy = recon_img.clone()
         # normalization
-        recon_img, img = self._post_process_img(model_name, recon_img, img)
+        recon_img, img = self._post_process_img(recon_img, model_name, rate), self._post_process_img(img, None, 1)
+        self._state_reset(model)
         # metric
         if self.cfg.save_metric == True:
             self.logger.info(f"----------------------Method: {model_name.upper()}----------------------")
@@ -209,19 +221,24 @@ class Pipeline:
             self.logger.info(f"Images are saved on the {save_folder}")
         return recon_img_copy
 
-    def _post_process_img(self, model_name, recon_img, gt_img):
+    def _post_process_img(self, recon_img, model_name, rate=1):
         """Post process the reconstructed image."""
-        # TFP and TFI algorithms are normalized automatically, others are normalized based on the self.cfg.use_norm
-        if model_name in ["tfp", "tfi", "spikeformer", "spikeclip"]:
-            recon_img = (recon_img - recon_img.min()) / (recon_img.max() - recon_img.min())
-        elif self.cfg.save_img_norm == True:
+        # 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:
             recon_img = (recon_img - recon_img.min()) / (recon_img.max() - recon_img.min())
+        else:
+            recon_img = recon_img / rate
         recon_img = recon_img.clip(0, 1)
-        gt_img = (gt_img - gt_img.min()) / (gt_img.max() - gt_img.min()) if self.cfg.gt_img_norm == True and gt_img is not None else gt_img
-        return recon_img, gt_img
+        return recon_img
 
     def _cal_metrics_model(self, model: BaseModel):
         """Calculate the metrics for the given model."""
+        # metric state reset (since get_outputs_dict from the training state is utilized)
+        model_state = model.net.training
+        model.net.training = True
         # metrics construct
         model_name = model.cfg.model_name
         metrics_dict = {}
@@ -234,17 +251,19 @@ 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(model_name, recon_img, img)
+            recon_img, img = self._post_process_img(recon_img, model_name), self._post_process_img(img, "auto")
             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))
                 elif metric_name in metric_single_names:
                     metrics_dict[metric_name].update(cal_metric_single(recon_img, metric_name))
-
+        self._state_reset(model)
+        model.net.training = model_state
         # metrics self.logger.info
         self.logger.info(f"----------------------Method: {model_name.upper()}----------------------")
         for metric_name in metrics_dict.keys():
             self.logger.info(f"{metric_name.upper()}: {metrics_dict[metric_name].avg}")
+        return metrics_dict
 
     def _cal_prams_model(self, model):
         """Calculate the parameters for the given model."""
@@ -256,7 +275,7 @@ class Pipeline:
         spike = torch.zeros((1, 200, 250, 400)).cuda()
         start_time = time.time()
         for _ in range(100):
-            model(spike)
+            model.spk2img(spike)
         latency = (time.time() - start_time) / 100
         # flop # todo thop bug for BSF
         flops, _ = profile((model), inputs=(spike,))
@@ -267,3 +286,8 @@ class Pipeline:
         )
         self.logger.info(f"----------------------Method: {model_name}----------------------")
         self.logger.info(re_msg)
+
+    def _state_reset(self, model):
+        """State reset for the snn-based method."""
+        if model.cfg.model_name == "ssir":
+            functional.reset_net(model.net)

spikezoo/pipeline/ensemble_pipeline.py

@@ -19,6 +19,7 @@ from spikezoo.models import build_model_cfg, build_model_name, BaseModel, BaseMo
 from spikezoo.datasets import build_dataset_cfg, build_dataset_name, BaseDataset, BaseDatasetConfig, build_dataloader
 from typing import Optional, Union, List
 from spikezoo.pipeline.base_pipeline import Pipeline, PipelineConfig
+import spikezoo as sz
 
 
 @dataclass
@@ -30,30 +31,30 @@ class EnsemblePipeline(Pipeline):
     def __init__(
         self,
         cfg: PipelineConfig,
-        model_cfg_list: Union[List[str], List[BaseModelConfig]],
-        dataset_cfg: Union[str, BaseDatasetConfig],
+        model_cfg_list: Union[List[sz.METHOD], List[BaseModelConfig]],
+        dataset_cfg: Union[sz.DATASET, BaseDatasetConfig],
     ):
         self.cfg = cfg
-        self._setup_model_data(model_cfg_list,dataset_cfg)
+        self._setup_model_data(model_cfg_list, dataset_cfg)
         self._setup_pipeline()
 
-    def _setup_model_data(self,model_cfg_list,dataset_cfg):
+    def _setup_model_data(self, model_cfg_list, dataset_cfg):
         """Model and Data setup."""
         # model
         self.model_list: List[BaseModel] = (
-            [build_model_name(name) for name in model_cfg_list] if isinstance(model_cfg_list[0],str) else [build_model_cfg(cfg) for cfg in model_cfg_list]
+            [build_model_name(name) for name in model_cfg_list] if isinstance(model_cfg_list[0], str) else [build_model_cfg(cfg) for cfg in model_cfg_list]
         )
-        self.model_list = [model.eval() for model in self.model_list]
+        self.model_list = [model.build_network(mode="eval", version=self.cfg.version) for model in self.model_list]
         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)
         # device
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
-    
-    def _spk2img(self, spike, img, save_folder):
+
+    def infer(self, spike, img, save_folder, rate):
         for model in self.model_list:
-            self._spk2img_model(model, spike, img, save_folder)
+            self._infer_model(model, spike, img, save_folder, rate)
 
     def cal_params(self):
         for model in self.model_list:

spikezoo/pipeline/train_cfgs.py

@@ -0,0 +1,89 @@
+import torch.nn as nn
+import torch.optim as optimizer
+import torch.optim.lr_scheduler as lr_scheduler
+import functools
+from spikezoo.utils.optimizer_utils import OptimizerConfig, AdamOptimizerConfig
+from spikezoo.utils.scheduler_utils import SchedulerConfig, MultiStepSchedulerConfig, CosineAnnealingLRConfig
+from dataclasses import dataclass, field
+from spikezoo.pipeline.train_pipeline import TrainPipelineConfig
+from typing import Optional, Dict, List
+
+
+@dataclass
+class REDS_BASE_TrainConfig(TrainPipelineConfig):
+    """Training setting for methods on the REDS-BASE dataset."""
+
+    # parameters setting
+    epochs: int = 600
+    steps_per_save_imgs: int = 200
+    steps_per_save_ckpt: int = 500
+    steps_per_cal_metrics: int = 100
+    metric_names: List[str] = field(default_factory=lambda: ["psnr", "ssim","lpips","niqe","brisque","piqe"])
+
+    # dataloader setting
+    bs_train: int = 8
+    num_workers: int = 4
+    pin_memory: bool = False
+
+    # train setting - optimizer & scheduler & loss_dict
+    optimizer_cfg: OptimizerConfig = AdamOptimizerConfig(lr=1e-4)
+    scheduler_cfg: Optional[SchedulerConfig] = MultiStepSchedulerConfig(milestones=[400], gamma=0.2) # from wgse
+    loss_weight_dict: Dict = field(default_factory=lambda: {"l1": 1})
+
+# @dataclass
+# class REDS_BASE_TrainConfig(TrainPipelineConfig):
+#     """Training setting for methods on the REDS-BASE dataset."""
+
+#     # parameters setting
+#     epochs: int = 700
+#     steps_per_save_imgs: int = 200
+#     steps_per_save_ckpt: int = 500
+#     steps_per_cal_metrics: int = 100
+#     metric_names: List[str] = field(default_factory=lambda: ["psnr", "ssim"])
+
+#     # dataloader setting
+#     bs_train: int = 8
+#     num_workers: int = 4
+#     pin_memory: bool = False
+
+#     # train setting - optimizer & scheduler & loss_dict
+#     optimizer_cfg: OptimizerConfig = AdamOptimizerConfig(lr=1e-4)
+#     scheduler_cfg: Optional[SchedulerConfig] = MultiStepSchedulerConfig(milestones=[400, 600], 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):
+    """Training setting for BSF. https://github.com/ruizhao26/BSF"""
+
+    optimizer_cfg: OptimizerConfig = AdamOptimizerConfig(lr=1e-4, weight_decay=0.0)
+    scheduler_cfg: Optional[SchedulerConfig] = MultiStepSchedulerConfig(milestones=[10, 20, 30, 40, 50, 60, 70, 80, 90, 100])
+    loss_weight_dict: Dict = field(default_factory=lambda: {"l1": 1})
+
+
+@dataclass
+class WGSETrainConfig(TrainPipelineConfig):
+    """Training setting for WGSE. https://github.com/Leozhangjiyuan/WGSE-SpikeCamera"""
+
+    optimizer_cfg: OptimizerConfig = AdamOptimizerConfig(lr=1e-4, betas=(0.9, 0.99), weight_decay=0)
+    scheduler_cfg: Optional[SchedulerConfig] = MultiStepSchedulerConfig(milestones=[400, 600], gamma=0.2)
+    loss_weight_dict: Dict = field(default_factory=lambda: {"l1": 1})
+
+
+@dataclass
+class STIRTrainConfig(TrainPipelineConfig):
+    """Training setting for STIR. https://github.com/GitCVfb/STIR"""
+
+    optimizer_cfg: OptimizerConfig = AdamOptimizerConfig(lr=1e-4, betas=(0.9, 0.999))
+    scheduler_cfg: Optional[SchedulerConfig] = MultiStepSchedulerConfig(milestones=[20, 25, 30, 35, 40, 45, 50, 55, 60, 65, 70], gamma=0.7)
+    loss_weight_dict: Dict = field(default_factory=lambda: {"l1": 1})
+
+
+@dataclass
+class Spk2ImgNetTrainConfig(TrainPipelineConfig):
+    """Training setting for Spk2ImgNet. https://github.com/Vspacer/Spk2ImgNet"""
+
+    optimizer_cfg: OptimizerConfig = AdamOptimizerConfig(lr=1e-4)
+    scheduler_cfg: Optional[SchedulerConfig] = MultiStepSchedulerConfig(milestones=[20], gamma=0.1)
+    loss_weight_dict: Dict = field(default_factory=lambda: {"l1": 1})

spikezoo/pipeline/train_pipeline.py

@@ -1,27 +1,64 @@
 import torch
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 import os
 from spikezoo.utils.img_utils import tensor2npy
 import cv2
 from pathlib import Path
-from typing import Literal
+from typing import Literal, Dict
 from tqdm import tqdm
 from spikezoo.models import build_model_cfg, build_model_name, BaseModel, BaseModelConfig
 from spikezoo.datasets import build_dataset_cfg, build_dataset_name, BaseDataset, BaseDatasetConfig, build_dataloader
-from typing import  Union
+from typing import Union, List, Optional
 from spikezoo.pipeline.base_pipeline import Pipeline, PipelineConfig
+import torch.nn as nn
+import torch.optim as optimizer
+import torch.optim.lr_scheduler as lr_scheduler
+import functools
+from spikezoo.utils.optimizer_utils import OptimizerConfig, AdamOptimizerConfig
+from spikezoo.utils.scheduler_utils import SchedulerConfig, MultiStepSchedulerConfig, CosineAnnealingLRConfig
+from torch.utils.tensorboard import SummaryWriter
+import subprocess
+import webbrowser
+import time
+import re
+from spikezoo.utils.other_utils import set_random_seed
+from spikingjelly.clock_driven import functional
 
 
 @dataclass
 class TrainPipelineConfig(PipelineConfig):
-    bs_train: int = 4
-    epochs: int = 100
-    lr: float = 1e-3
+    # parameters setting
+    "Training epochs."
+    epochs: int = 1000
+    "Steps per to save images."
+    steps_per_save_imgs: int = 200
+    "Steps per to save model weights."
+    steps_per_save_ckpt: int = 500
+    "Steps per to calculate the metrics."
+    steps_per_cal_metrics: int = 100
+    "Step for gradient accumulation. (for snn methods)"
+    steps_grad_accumulation: int = 4
+    "Pipeline mode."
+    _mode: Literal["single_mode", "multi_mode", "train_mode"] = "train_mode"
+    "Use tensorboard or not"
+    use_tensorboard: bool = True
+    "Random seed."
+    seed: int = 521
+    # dataloader setting
+    "Batch size for the train dataloader."
+    bs_train: int = 8
+    "Num_workers for the train dataloader."
     num_workers: int = 4
+    "Pin_memory true or false for the train dataloader."
     pin_memory: bool = False
-    steps_per_save_imgs = 10
-    steps_per_cal_metrics = 10
-    _mode: Literal["single_mode", "multi_mode", "train_mode"] = "train_mode"
+
+    # train setting - optimizer & scheduler & loss_dict
+    "Optimizer config."
+    optimizer_cfg: OptimizerConfig = AdamOptimizerConfig(lr=1e-3)
+    "Scheduler config."
+    scheduler_cfg: Optional[SchedulerConfig] = None
+    "Loss dict {loss_name,weight}."
+    loss_weight_dict: Dict[Literal["l1", "l2"], float] = field(default_factory=lambda: {"l1": 1})
 
 
 class TrainPipeline(Pipeline):
@@ -34,13 +71,20 @@ class TrainPipeline(Pipeline):
         self.cfg = cfg
         self._setup_model_data(model_cfg, dataset_cfg)
         self._setup_pipeline()
-        self.model.setup_training(cfg)
+        self._setup_training()
+
+    def _setup_pipeline(self):
+        super()._setup_pipeline()
+        set_random_seed(self.cfg.seed)
+        if self.cfg.use_tensorboard:
+            self.writer = SummaryWriter(self.save_folder / Path(""))
+            subprocess.Popen(["tensorboard", f"--logdir={self.save_folder}"])
 
     def _setup_model_data(self, model_cfg, dataset_cfg):
         """Model and Data setup."""
         # model
         self.model: BaseModel = build_model_name(model_cfg) if isinstance(model_cfg, str) else build_model_cfg(model_cfg)
-        self.model = self.model.train()
+        self.model.build_network(mode = "train",version="local")
         torch.set_grad_enabled(True)
         # data
         if isinstance(dataset_cfg, str):
@@ -54,41 +98,96 @@ class TrainPipeline(Pipeline):
         # device
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
 
+    def _setup_training(self):
+        """Setup training optimizer."""
+        self.optimizer = self.cfg.optimizer_cfg.setup(self.model.net.parameters())
+        self.scheduler = self.cfg.scheduler_cfg.setup(self.optimizer) if self.cfg.scheduler_cfg != None else None
+        self.cnt_grad = 0
+
     def save_network(self, epoch):
         """Save the network."""
         save_folder = self.save_folder / Path("ckpt")
         os.makedirs(save_folder, exist_ok=True)
         self.model.save_network(save_folder / f"{epoch:06d}.pth")
 
+    def train(self):
+        """Training code."""
+        self.logger.info("Start Training!")
+        for epoch in range(self.cfg.epochs):
+            # training
+            for batch_idx, batch in enumerate(tqdm(self.train_dataloader)):
+                batch = self.model.feed_to_device(batch)
+                outputs = self.model.get_outputs_dict(batch)
+                loss_dict, loss_values_dict = self.model.get_loss_dict(outputs, batch, self.cfg.loss_weight_dict)
+                self.optimize_parameters(loss_dict, batch_idx == len(self.train_dataloader) - 1)
+            self.update_learning_rate()
+            self.write_log_train(epoch, loss_values_dict)
+
+            #  save visual results & save ckpt & evaluate metrics
+            with torch.no_grad():
+                if epoch % self.cfg.steps_per_save_imgs == 0 or epoch == self.cfg.epochs - 1:
+                    self.save_visual(epoch)
+                if epoch % self.cfg.steps_per_save_ckpt == 0 or epoch == self.cfg.epochs - 1:
+                    self.save_network(epoch)
+                if epoch % self.cfg.steps_per_cal_metrics == 0 or epoch == self.cfg.epochs - 1:
+                    metrics_dict = self.cal_metrics()
+                    self.write_log_test(epoch, metrics_dict)
+
     def save_visual(self, epoch):
         """Save the visual results."""
         self.logger.info("Saving visual results...")
         save_folder = self.save_folder / Path("imgs") / Path(f"{epoch:06d}")
         os.makedirs(save_folder, exist_ok=True)
         for batch_idx, batch in enumerate(tqdm(self.dataloader)):
-            if batch_idx % (len(self.dataloader) // 4) != 0:
-                continue 
+            if batch_idx % (len(self.dataloader) // 3) != 0:
+                continue
             batch = self.model.feed_to_device(batch)
             outputs = self.model.get_outputs_dict(batch)
             visual_dict = self.model.get_visual_dict(batch, outputs)
+            self._state_reset(self.model)
             # save
             for key, img in visual_dict.items():
+                img = self._post_process_img(img, model_name=self.model.cfg.model_name)
                 cv2.imwrite(str(save_folder / Path(f"{batch_idx:06d}_{key}.png")), tensor2npy(img))
+                if self.cfg.use_tensorboard == True and batch_idx == 0:
+                    self.writer.add_image(f"imgs/{key}", img[0].detach().cpu(), epoch)
 
-    def train(self):
-        """Training code."""
-        self.logger.info("Start Training!")
-        for epoch in range(self.cfg.epochs):
-            # training
-            for batch_idx, batch in enumerate(tqdm(self.train_dataloader)):
-                batch = self.model.feed_to_device(batch)
-                outputs = self.model.get_outputs_dict(batch)
-                loss_dict, loss_values_dict = self.model.get_loss_dict(outputs, batch)
-                self.model.optimize_parameters(loss_dict)
-            self.model.update_learning_rate()
-            self.logger.info(f"EPOCH {epoch}/{self.cfg.epochs}: Train Loss: {loss_values_dict}")
-            #  save visual results & evaluate metrics
-            if epoch % self.cfg.steps_per_save_imgs == 0 or epoch == self.cfg.epochs - 1:
-                self.save_visual(epoch)
-            if epoch % self.cfg.steps_per_cal_metrics == 0 or epoch == self.cfg.epochs - 1:
-                self.cal_metrics()
+    def optimize_parameters(self, loss_dict, final_flag):
+        """Optimize the parameters from the loss_dict."""
+        loss = functools.reduce(torch.add, loss_dict.values())
+        step_grad = self.cfg.steps_grad_accumulation
+        # for snn methods
+        if self.model.cfg.model_name == "ssir":
+            if self.cnt_grad % step_grad != step_grad - 1 and final_flag == False:
+                loss.backward(retain_graph=True)
+                self.cnt_grad += 1
+            else:
+                loss.backward(retain_graph=False)
+                self.optimizer.step()
+                self.optimizer.zero_grad()
+                self._state_reset(self.model)
+                self.cnt_grad = 0
+        else:
+            # for cnn methods
+            self.optimizer.zero_grad()
+            loss.backward()
+            self.optimizer.step()
+
+    def update_learning_rate(self):
+        """Update the learning rate."""
+        self.scheduler.step() if self.cfg.scheduler_cfg != None else None
+
+    def write_log_train(self, epoch, loss_values_dict):
+        """Write the train log information."""
+        self.logger.info(f"EPOCH {epoch}/{self.cfg.epochs}: Train Loss: {loss_values_dict}")
+        if self.cfg.use_tensorboard:
+            for name, val in loss_values_dict.items():
+                self.writer.add_scalar(f"Loss/{name}", val, epoch)
+            lr = self.optimizer.param_groups[0]["lr"]
+            self.writer.add_scalar(f"Loss/lr", lr, epoch)
+
+    def write_log_test(self, epoch, metrics_dict):
+        """Write the test log information."""
+        if self.cfg.use_tensorboard:
+            for name in metrics_dict.keys():
+                self.writer.add_scalar(f"Test/{name}", metrics_dict[name].avg, epoch)