Repository: intel/e2eAIOK (language: python)

Issue:
Refine DLRM README



Discussion / hints:

Fix patch:
diff --git a/Dockerfile-ubuntu18.04/DockerfilePytorch112 b/Dockerfile-ubuntu18.04/DockerfilePytorch112
index 881037712..fe1e5606d 100644
--- a/Dockerfile-ubuntu18.04/DockerfilePytorch112
+++ b/Dockerfile-ubuntu18.04/DockerfilePytorch112
@@ -7,7 +7,7 @@ ENV PATH /opt/intel/oneapi/intelpython/latest/condabin:$PATH
 RUN apt-get update -y && apt-get install -y openssh-server openjdk-8-jdk sshpass numactl
 
 SHELL ["conda", "run", "-n", "pytorch-1.12.0", "/bin/bash", "-c"]
-RUN python -m pip install torchsummary easydict sentencepiece opencv-python scikit-image timm boto3 ptflops sigopt e2eAIOK
+RUN python -m pip install torchsummary easydict sentencepiece opencv-python scikit-image tensorboard timm boto3 ptflops sigopt e2eAIOK
 RUN python -m pip install torchaudio~=0.12.0 -f https://download.pytorch.org/whl/cpu/torch_stable.html
 RUN python -m pip install oneccl_bind_pt==1.12.100 -f https://developer.intel.com/ipex-whl-stable
 


Test patch:
diff --git a/e2eAIOK/DeNas/requirements.txt b/e2eAIOK/DeNas/requirements.txt
index 429e7fa00..5c8bca4f2 100644
--- a/e2eAIOK/DeNas/requirements.txt
+++ b/e2eAIOK/DeNas/requirements.txt
@@ -7,4 +7,5 @@ timm
 boto3
 ptflops
 sigopt
+tensorboard
 torchaudio~=0.12.0 -f https://download.pytorch.org/whl/cpu/torch_stable.html
\ No newline at end of file
diff --git a/e2eAIOK/common/default.conf b/e2eAIOK/common/default.conf
new file mode 100644
index 000000000..8c42fb100
--- /dev/null
+++ b/e2eAIOK/common/default.conf
@@ -0,0 +1,75 @@
+## show all the configurations and default values in common lib
+## include required and optional settings
+
+########################### required ############################
+### training 
+train_epochs: 1
+eval_epochs: 1
+eval_step: 10 # optional
+output_dir: "" # optional
+device: "cpu" # optional
+dist_backend: "gloo"
+
+### model
+model_type: ""
+pretrain: "" # optional - load pretrain model from local path
+
+### dataset
+data_set: ""
+data_path: ""
+train_batch_size: 128
+eval_batch_size: 128
+num_workers: 1
+input_size: 32 # optional
+drop_last: False # optional
+pin_mem: False # optional
+
+### metric
+eval_metric: "accuracy"
+metric_threshold: 100.0 # optional
+
+### loss
+criterion: "CrossEntropyLoss"
+
+## optimizer
+optimizer: "SGD"
+learning_rate: 0.05
+momentum: 0.9
+weight_decay: 0.0001
+
+########################### optional ############################
+### optimize
+enable_ipex: False
+
+### tensorboard
+tensorboard_dir: ""
+
+### scheduler
+lr_scheduler: "" ## such as "CosineAnnealingLR"
+lr_scheduler_config:
+    decay_rate: 0.1
+    decay_stages: [] #for MultiStepLR
+    decay_patience: 10 # for ReduceLROnPlateau
+    T_max: 0 # for CosineAnnealingLR, default is train_epochs if T_max <=0
+
+### warm up scheduler
+warmup_scheduler: "" # such as "WarmUpLR"
+warmup_scheduler_epoch: 0
+
+### early stop
+early_stop: "" # such as "EarlyStopping"
+early_stop_config:
+    tolerance_epoch: 3
+    delta: 0.0001
+    is_max: True
+
+### profiler
+profiler: False
+profiler_config:
+    skip_first: 1
+    wait: 1
+    warmup: 1
+    active: 2
+    repeat: 1
+    trace_file: ""
+
diff --git a/e2eAIOK/common/trainer/data/asr/data_builder_librispeech.py b/e2eAIOK/common/trainer/data/asr/data_builder_librispeech.py
index ae82180a7..8dbfeaf20 100644
--- a/e2eAIOK/common/trainer/data/asr/data_builder_librispeech.py
+++ b/e2eAIOK/common/trainer/data/asr/data_builder_librispeech.py
@@ -10,4 +10,5 @@ def __init__(self, cfg, tokenizer):
 
     def prepare_dataset(self):
         train_data, valid_data, test_datasets = dataio_prepare(self.cfg, self.tokenizer)
-        return train_data, valid_data
\ No newline at end of file
+        self.dataset_train = train_data
+        self.dataset_val = valid_data
\ No newline at end of file
diff --git a/e2eAIOK/common/trainer/data/cv/__init__.py b/e2eAIOK/common/trainer/data/cv/__init__.py
index fe30424c0..5f6501292 100644
--- a/e2eAIOK/common/trainer/data/cv/__init__.py
+++ b/e2eAIOK/common/trainer/data/cv/__init__.py
@@ -1,1 +1,1 @@
-__all__ = ['data_builder_cifar']
\ No newline at end of file
+__all__ = ['data_builder_cifar', 'data_builder_uspavsminist']
\ No newline at end of file
diff --git a/e2eAIOK/common/trainer/data/cv/data_builder_cifar.py b/e2eAIOK/common/trainer/data/cv/data_builder_cifar.py
index fa92640c5..269dd5ae9 100644
--- a/e2eAIOK/common/trainer/data/cv/data_builder_cifar.py
+++ b/e2eAIOK/common/trainer/data/cv/data_builder_cifar.py
@@ -1,11 +1,8 @@
 import os
 import torch
 from torchvision import datasets, transforms
-from torchvision.datasets.folder import ImageFolder
-from timm.data.constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
-from timm.data import create_transform
-import e2eAIOK.common.trainer.utils.extend_distributed as ext_dist
 from e2eAIOK.common.trainer.data.data_builder_cv import DataBuilderCV
+from e2eAIOK.common.utils import check_mkdir
 
 class DataBuilderCIFAR(DataBuilderCV):
     def __init__(self, cfg):
@@ -13,24 +10,26 @@ def __init__(self, cfg):
     
     def prepare_dataset(self):
         """
-            prepare CV related dataset
+            prepare CIFAR dataset
         """
-        if self.cfg.data_set in ["CIFAR10","CIFAR100"]:
+        if self.cfg.data_set.lower() in ["cifar10","cifar100"]:
             dataset_train = self.build_dataset(is_train=True)
             dataset_val = self.build_dataset(is_train=False)
         else:
             raise RuntimeError(f"dataset {self.cfg.data_set} not supported")
-        return dataset_train, dataset_val
+        self.dataset_train = dataset_train
+        self.dataset_val = dataset_val
     
     def build_dataset(self, is_train = True):
         transform = self.build_transform(is_train)
-        if self.cfg.data_set == 'CIFAR10':
-            dataset = datasets.CIFAR10(self.cfg.data_path, train=is_train, transform=transform, download=True)
-        elif self.cfg.data_set == 'CIFAR100':
-            dataset = datasets.CIFAR100(self.cfg.data_path, train=is_train, transform=transform, download=True)
+        data_folder = check_mkdir(self.cfg.data_path)
+        if self.cfg.data_set.lower() == 'cifar10':
+            dataset = datasets.CIFAR10(data_folder, train=is_train, transform=transform, download=True)
+        elif self.cfg.data_set.lower() == 'cifar100':
+            dataset = datasets.CIFAR100(data_folder, train=is_train, transform=transform, download=True)
         return dataset
     def build_transform(self, is_train):
-        resize_im = self.cfg.input_size > 32
+        resize_im = "input_size" in self.cfg and self.cfg.input_size > 32
         if is_train:
             # this should always dispatch to transforms_imagenet_train
             transform = transforms.Compose([
diff --git a/e2eAIOK/common/trainer/data/cv/data_builder_uspsvsminist.py b/e2eAIOK/common/trainer/data/cv/data_builder_uspsvsminist.py
new file mode 100644
index 000000000..6d885a6b7
--- /dev/null
+++ b/e2eAIOK/common/trainer/data/cv/data_builder_uspsvsminist.py
@@ -0,0 +1,36 @@
+#!/usr/bin/python
+# -*- coding: utf-8 -*-
+import os
+import torch
+from torch.utils.data import random_split
+from e2eAIOK.common.trainer.data.data_utils.image_list import ImageList
+from torchvision import transforms
+from e2eAIOK.common.trainer.data.data_builder_cv import DataBuilderCV
+
+class DataBuilderUSPSMinist(DataBuilderCV):
+    def __init__(self, cfg):
+        super().__init__(cfg)
+
+    def prepare_dataset(self):
+        """
+            prepare USPS_vs_Minist dataset
+        """
+        transform = transforms.Compose([
+            transforms.Resize(28),
+            transforms.ToTensor(),
+            # transforms.Normalize(
+            # mean=[0.485, 0.456, 0.406],
+            # std=[0.229, 0.224, 0.225])
+        ])
+        test_dataset = ImageList(os.path.join(self.cfg.data_path,"MNIST"), 
+                                    open(os.path.join(self.cfg.data_path,"MNIST/mnist_test.txt")).readlines(),
+                                    transform, 'L')
+        train_dataset = ImageList(os.path.join(self.cfg.data_path,"MNIST"), 
+                                    open(os.path.join(self.cfg.data_path,"MNIST/mnist_train.txt")).readlines(),
+                                    transform, 'L')
+        test_len = len(test_dataset)
+        test_dataset, validation_dataset = random_split(test_dataset,[test_len // 2, test_len - test_len // 2])
+        self.dataset_train = train_dataset
+        self.dataset_val = validation_dataset
+        self.dataset_test = test_dataset
+    
\ No newline at end of file
diff --git a/e2eAIOK/common/trainer/data/data_builder_asr.py b/e2eAIOK/common/trainer/data/data_builder_asr.py
index 3e2f0500c..6dac4245d 100644
--- a/e2eAIOK/common/trainer/data/data_builder_asr.py
+++ b/e2eAIOK/common/trainer/data/data_builder_asr.py
@@ -10,23 +10,24 @@ def __init__(self, cfg, tokenizer):
         self.tokenizer = tokenizer
 
     def get_dataloader(self):
-        dataset_train, dataset_val = self.prepare_dataset()
+        if self.dataset_train is None or self.dataset_val is None:
+            self.prepare_dataset()
 
         if ext_dist.my_size > 1:
             num_tasks = ext_dist.dist.get_world_size()
             global_rank = ext_dist.dist.get_rank()
             sampler_train = torch.utils.data.DistributedSampler(
-                dataset_train, num_replicas=num_tasks, rank=global_rank
+                self.dataset_train, num_replicas=num_tasks, rank=global_rank
             )
             
             sampler_val = torch.utils.data.DistributedSampler(
-                dataset_val, num_replicas=num_tasks, rank=global_rank)
+                self.dataset_val, num_replicas=num_tasks, rank=global_rank)
         else:
             sampler_val = None
             sampler_train = None
 
         dataloader_train = torch.utils.data.DataLoader(
-            dataset_train, 
+            self.dataset_train, 
             sampler=sampler_train,
             batch_size=self.cfg.train_batch_size,
             num_workers=self.cfg.num_workers,
@@ -36,7 +37,7 @@ def get_dataloader(self):
         )
 
         dataloader_val = torch.utils.data.DataLoader(
-            dataset_val, 
+            self.dataset_val, 
             batch_size=self.cfg.eval_batch_size,
             sampler=sampler_val, 
             num_workers=self.cfg.num_workers,
diff --git a/e2eAIOK/common/trainer/data/data_builder_cv.py b/e2eAIOK/common/trainer/data/data_builder_cv.py
index f48828ece..135651f4d 100644
--- a/e2eAIOK/common/trainer/data/data_builder_cv.py
+++ b/e2eAIOK/common/trainer/data/data_builder_cv.py
@@ -1,12 +1,7 @@
 import os
 import torch
-from torchvision import datasets, transforms
-from torchvision.datasets.folder import ImageFolder
-from timm.data.constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
-from timm.data import create_transform
 import e2eAIOK.common.trainer.utils.extend_distributed as ext_dist
-
-
+# from e2eAIOK.common.trainer.data.data_utils.data_utils import channels_last_collate
 from e2eAIOK.common.trainer.data_builder import DataBuilder
 
 class DataBuilderCV(DataBuilder):
@@ -15,38 +10,66 @@ def __init__(self, cfg):
     
     def get_dataloader(self):
         """
-            create training/evaluation dataloader
+            create training/evaluation/test dataloader
         """
-        dataset_train, dataset_val = self.prepare_dataset()
+        ### generate datasets
+        if self.dataset_train is None or self.dataset_val is None:
+            self.prepare_dataset()
+
+        ### get configurations
+        drop_last =  self.cfg.drop_last if "drop_last" in self.cfg else False
+        pin_memory = self.cfg.pin_mem if "pin_mem" in self.cfg else False
 
+        ### check whether distributed or enable_ipex
         if ext_dist.my_size > 1:
             num_tasks = ext_dist.dist.get_world_size()
             global_rank = ext_dist.dist.get_rank()
             sampler_train = torch.utils.data.DistributedSampler(
-                dataset_train, num_replicas=num_tasks, rank=global_rank, shuffle=True, drop_last= True
-            )
-            
+                self.dataset_train, num_replicas=num_tasks, rank=global_rank, shuffle=True, drop_last= drop_last)
             sampler_val = torch.utils.data.DistributedSampler(
-                dataset_val, num_replicas=num_tasks, rank=global_rank,  shuffle=False)
+                self.dataset_val, num_replicas=num_tasks, rank=global_rank,  shuffle=False)
+            if self.dataset_test is not None:
+                sampler_test = torch.utils.data.DistributedSampler(
+                self.dataset_test, num_replicas=num_tasks, rank=global_rank,  shuffle=False)
         else:
-            sampler_val = torch.utils.data.SequentialSampler(dataset_val)
-            sampler_train = torch.utils.data.RandomSampler(dataset_train)
+            sampler_train = torch.utils.data.RandomSampler(self.dataset_train)
+            sampler_val = torch.utils.data.SequentialSampler(self.dataset_val)
+            sampler_test = torch.utils.data.SequentialSampler(self.dataset_test)
+        
+        # if "enable_ipex" in self.cfg and self.cfg.enable_ipex:
+        #     collate_fn = channels_last_collate
+        # else:
+        collate_fn = None
 
+        ### generate dataloader
         dataloader_train = torch.utils.data.DataLoader(
-            dataset_train, 
+            self.dataset_train, 
             sampler=sampler_train,
             batch_size=self.cfg.train_batch_size,
             num_workers=self.cfg.num_workers,
-            pin_memory=self.cfg.pin_mem
+            drop_last=drop_last,
+            collate_fn=collate_fn,
+            pin_memory=pin_memory
         )
 
         dataloader_val = torch.utils.data.DataLoader(
-            dataset_val, 
+            self.dataset_val, 
             batch_size=self.cfg.eval_batch_size,
             sampler=sampler_val, 
             num_workers=self.cfg.num_workers,
-            shuffle=False,
-            drop_last=False
+            drop_last=False,
+            collate_fn=collate_fn
         )
-        
-        return dataloader_train, dataloader_val
\ No newline at end of file
+
+        if self.dataset_test is None:
+            return dataloader_train, dataloader_val
+        else:
+            dataloader_test = torch.utils.data.DataLoader(
+                self.dataset_test, 
+                batch_size=self.cfg.eval_batch_size,
+                sampler=sampler_test, 
+                num_workers=self.cfg.num_workers,
+                drop_last=False,
+                collate_fn=collate_fn
+            )
+        return dataloader_train, dataloader_val, dataloader_test
\ No newline at end of file
diff --git a/e2eAIOK/common/trainer/data/data_utils/data_utils.py b/e2eAIOK/common/trainer/data/data_utils/data_utils.py
new file mode 100644
index 000000000..7d49cedc9
--- /dev/null
+++ b/e2eAIOK/common/trainer/data/data_utils/data_utils.py
@@ -0,0 +1,19 @@
+import torch
+
+#optimize for ipex
+def channels_last_collate(batch):
+    """Custom collate fn for channels_last.
+    Arguments:
+        batch: (tuple) A tuple of images and labels
+    Return:
+        A tuple containing:
+            1) (tensor) batch of images stacked on their 0 dim and to channels_last
+            2) (list of tensors) labels
+    """
+    data = [item[0] for item in batch]
+    target = [item[1] for item in batch]
+
+    data = torch.stack(data, 0).to(memory_format=torch.channels_last)
+    target = torch.LongTensor(target)
+    
+    return data, target
\ No newline at end of file
diff --git a/e2eAIOK/common/trainer/data/data_utils/image_list.py b/e2eAIOK/common/trainer/data/data_utils/image_list.py
new file mode 100644
index 000000000..fad75bd8c
--- /dev/null
+++ b/e2eAIOK/common/trainer/data/data_utils/image_list.py
@@ -0,0 +1,49 @@
+#!/usr/bin/python
+# -*- coding: utf-8 -*-
+
+from torch.utils.data import Dataset
+from .loader import rgb_loader,l_loader
+import logging
+
+class ImageList(Dataset):
+    ''' ImageList dataset
+
+    '''
+    def __init__(self, data_path,label_records,transform,img_mode='RGB'):
+        ''' init method
+
+        :param data_path: image dir
+        :param label_records: label record list, each record format: "image_name image_label"
+        :param transform: data transform
+        :param img_mode: img mode, must one "RGB" or "L" (ignore letter case)
+        '''
+        imgs = [("%s/%s"%(data_path,val.split()[0]), int(val.split()[1])) for val in label_records]
+        if len(imgs) == 0:
+            logging.error("Found 0 images in subfolders of: %s" %data_path)
+            raise (RuntimeError("Found 0 images in subfolders of: %s" %data_path))
+
+        self.imgs = imgs
+        self.transform = transform
+        self.img_mode = img_mode
+
+        if img_mode.upper() == 'RGB':
+            self.loader = rgb_loader
+        elif img_mode.upper()  == 'L':
+            self.loader = l_loader
+        else:
+            logging.error("mode muse be one of: 'RGB','L' (ignore letter case), but found :%s"%img_mode)
+            raise ValueError("mode muse be one of: 'RGB','L' (ignore letter case), but found :%s"%img_mode)
+
+    def __getitem__(self, index):
+        path, label = self.imgs[index]
+        img = self.loader(path)
+        if self.transform is not None:
+            img = self.transform(img)
+        return img, label
+
+    def __len__(self):
+        return len(self.imgs)
+
+    def __str__(self):
+        return 'ImageList: image num [%s], img_mode[%s], transform [%s]'%(
+            len(self.imgs),self.img_mode,self.transform)
\ No newline at end of file
diff --git a/e2eAIOK/common/trainer/data/data_utils/loader.py b/e2eAIOK/common/trainer/data/data_utils/loader.py
new file mode 100644
index 000000000..6901a3d87
--- /dev/null
+++ b/e2eAIOK/common/trainer/data/data_utils/loader.py
@@ -0,0 +1,24 @@
+#!/usr/bin/python
+# -*- coding: utf-8 -*-
+from PIL import Image
+import torch
+
+def rgb_loader(path):
+    ''' load rgb image
+
+    :param path: image path
+    :return: image data
+    '''
+    with open(path, 'rb') as f:
+        with Image.open(f) as img:
+            return img.convert('RGB')
+
+def l_loader(path):
+    ''' load grayscale image
+
+    :param path:  image path
+    :return: image data
+    '''
+    with open(path, 'rb') as f:
+        with Image.open(f) as img:
+            return img.convert('L')
\ No newline at end of file
diff --git a/e2eAIOK/common/trainer/data/nlp/data_builder_squad.py b/e2eAIOK/common/trainer/data/nlp/data_builder_squad.py
index 9aa6abc60..ba2928ed6 100644
--- a/e2eAIOK/common/trainer/data/nlp/data_builder_squad.py
+++ b/e2eAIOK/common/trainer/data/nlp/data_builder_squad.py
@@ -25,19 +25,21 @@ def __init__(self, cfg):
     def prepare_dataset(self):
         dataset_train, train_examples, train_dataset, labels = build_dataset(is_train=True, args=self.cfg)
         dataset_val, val_examples, val_dataset, val_features, tokenizer = build_dataset(is_train=False, args=self.cfg)
-        return dataset_train, dataset_val, (train_examples, val_examples, val_dataset, val_features, tokenizer)
+        self.dataset_train = dataset_train
+        self.dataset_val = dataset_val
+        return (train_examples, val_examples, val_dataset, val_features, tokenizer)
 
     def get_dataloader(self):
         """
             create training/evaluation dataloader
         """
-        dataset_train, dataset_val, other_data = self.prepare_dataset()
+        other_data = self.prepare_dataset()
 
         if ext_dist.my_size > 1:
             num_tasks = ext_dist.dist.get_world_size()
             global_rank = ext_dist.dist.get_rank()
             sampler_train = torch.utils.data.DistributedSampler(
-                dataset_train, num_replicas=num_tasks, rank=global_rank
+                self.dataset_train, num_replicas=num_tasks, rank=global_rank
             )
             
             #sampler_val = torch.utils.data.DistributedSampler(
@@ -52,7 +54,7 @@ def get_dataloader(self):
             shuffle = False
 
         dataloader_train = torch.utils.data.DataLoader(
-            dataset_train, 
+            self.dataset_train, 
             sampler=sampler_train,
             batch_size=self.cfg.train_batch_size,
             num_workers=self.cfg.num_workers,
@@ -61,7 +63,7 @@ def get_dataloader(self):
         )
 
         dataloader_val = torch.utils.data.DataLoader(
-            dataset_val, 
+            self.dataset_val, 
             batch_size=self.cfg.eval_batch_size,
             sampler=sampler_val, 
             num_workers=self.cfg.num_workers,
diff --git a/e2eAIOK/common/trainer/data_builder.py b/e2eAIOK/common/trainer/data_builder.py
index 2262614c4..b7a3cdeaa 100644
--- a/e2eAIOK/common/trainer/data_builder.py
+++ b/e2eAIOK/common/trainer/data_builder.py
@@ -11,6 +11,9 @@ class DataBuilder():
     """
     def __init__(self,cfg):
         self.cfg = cfg
+        self.dataset_train = None
+        self.dataset_val = None
+        self.dataset_test = None
     
     def prepare_dataset(self):
         """
@@ -22,23 +25,24 @@ def get_dataloader(self):
         """
             create training/evaluation dataloader
         """
-        dataset_train, dataset_val = self.prepare_dataset()
+        if self.dataset_train is None or self.dataset_val is None:
+            self.prepare_dataset()
 
         if ext_dist.my_size > 1:
             num_tasks = ext_dist.dist.get_world_size()
             global_rank = ext_dist.dist.get_rank()
             sampler_train = torch.utils.data.DistributedSampler(
-                dataset_train, num_replicas=num_tasks, rank=global_rank
+                self.dataset_train, num_replicas=num_tasks, rank=global_rank
             )
             
             sampler_val = torch.utils.data.DistributedSampler(
-                dataset_val, num_replicas=num_tasks, rank=global_rank)
+                self.dataset_val, num_replicas=num_tasks, rank=global_rank)
         else:
             sampler_val = None
             sampler_train = None
 
         dataloader_train = torch.utils.data.DataLoader(
-            dataset_train, 
+            self.dataset_train, 
             sampler=sampler_train,
             batch_size=self.cfg.train_batch_size,
             num_workers=self.cfg.num_workers,
@@ -47,7 +51,7 @@ def get_dataloader(self):
         )
 
         dataloader_val = torch.utils.data.DataLoader(
-            dataset_val, 
+            self.dataset_val, 
             batch_size=self.cfg.eval_batch_size,
             sampler=sampler_val, 
             num_workers=self.cfg.num_workers,
@@ -55,4 +59,15 @@ def get_dataloader(self):
             drop_last=False
         )
         
-        return dataloader_train, dataloader_val
+        if self.dataset_test is None:
+            return dataloader_train, dataloader_val
+        else:
+            dataloader_test = torch.utils.data.DataLoader(
+            self.dataset_test, 
+            batch_size=self.cfg.eval_batch_size,
+            sampler=sampler_val, 
+            num_workers=self.cfg.num_workers,
+            shuffle=False,
+            drop_last=False
+        )
+        return dataloader_train, dataloader_val, dataloader_test
diff --git a/e2eAIOK/common/trainer/model/cv/lenet.py b/e2eAIOK/common/trainer/model/cv/lenet.py
new file mode 100644
index 000000000..5fc5b7a88
--- /dev/null
+++ b/e2eAIOK/common/trainer/model/cv/lenet.py
@@ -0,0 +1,35 @@
+#!/usr/bin/python
+# -*- coding: utf-8 -*-
+import torch.nn as nn
+import logging
+
+class LeNet(nn.Module):
+    ''' LeNet backbone
+
+    '''
+    def __init__(self,num_classes):
+        ''' init method
+
+        :param num_classes: num classes
+        '''
+        super(LeNet, self).__init__()
+        self.num_classes = num_classes
+        self.conv_layers = nn.Sequential(
+            nn.Conv2d(1, 20, kernel_size=5),
+            nn.MaxPool2d(2),
+            nn.ReLU(),
+            nn.Conv2d(20, 50, kernel_size=5),
+            nn.Dropout2d(p=0.5),
+            nn.MaxPool2d(2),
+            nn.ReLU(),
+        )
+
+        self.fc_layers = nn.Sequential(nn.Linear(50 * 4 * 4, 500), nn.ReLU(), nn.Dropout(p=0.5))
+        self.classifier = nn.Linear(500, num_classes)
+
+    def forward(self, x):
+        x = self.conv_layers(x)
+        x = x.view(x.size(0), -1)
+        x = self.fc_layers(x)
+        y = self.classifier(x)
+        return y
\ No newline at end of file
diff --git a/e2eAIOK/common/trainer/model/cv/resnet_cifar.py b/e2eAIOK/common/trainer/model/cv/resnet_cifar.py
new file mode 100644
index 000000000..ac4796084
--- /dev/null
+++ b/e2eAIOK/common/trainer/model/cv/resnet_cifar.py
@@ -0,0 +1,140 @@
+'''
+Copyright (c) 2022, Intel Corporation
+Copyright (c) 2020, https://github.com/weiaicunzai/pytorch-cifar100
+
+ Licensed under the Apache License, Version 2.0 (the "License");
+ you may not use this file except in compliance with the License.
+ You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
+
+ResNet in PyTorch.
+
+For Pre-activation ResNet, see 'preact_resnet.py'.
+
+Reference:
+[1] Kaiming He, Xiangyu Zhang, Shaoqing Ren, Jian Sun
+    Deep Residual Learning for Image Recognition. arXiv:1512.03385
+'''
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, in_planes, planes, stride=1):
+        super(BasicBlock, self).__init__()
+        self.conv1 = nn.Conv2d(
+            in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3,
+                               stride=1, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+
+        self.shortcut = nn.Sequential()
+        if stride != 1 or in_planes != self.expansion*planes:
+            self.shortcut = nn.Sequential(
+                nn.Conv2d(in_planes, self.expansion*planes,
+                          kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(self.expansion*planes)
+            )
+
+    def forward(self, x):
+        out = F.relu(self.bn1(self.conv1(x)))
+        out = self.bn2(self.conv2(out))
+        out += self.shortcut(x)
+        out = F.relu(out)
+        return out
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, in_planes, planes, stride=1):
+        super(Bottleneck, self).__init__()
+        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3,
+                               stride=stride, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.conv3 = nn.Conv2d(planes, self.expansion *
+                               planes, kernel_size=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(self.expansion*planes)
+
+        self.shortcut = nn.Sequential()
+        if stride != 1 or in_planes != self.expansion*planes:
+            self.shortcut = nn.Sequential(
+                nn.Conv2d(in_planes, self.expansion*planes,
+                          kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(self.expansion*planes)
+            )
+
+    def forward(self, x):
+        out = F.relu(self.bn1(self.conv1(x)))
+        out = F.relu(self.bn2(self.conv2(out)))
+        out = self.bn3(self.conv3(out))
+        out += self.shortcut(x)
+        out = F.relu(out)
+        return out
+
+
+class ResNet(nn.Module):
+    def __init__(self, block, num_blocks, num_classes=10):
+        super(ResNet, self).__init__()
+        self.in_planes = 64
+
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=3,
+                               stride=1, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(64)
+        self.layer1 = self._make_layer(block, 64, num_blocks[0], stride=1)
+        self.layer2 = self._make_layer(block, 128, num_blocks[1], stride=2)
+        self.layer3 = self._make_layer(block, 256, num_blocks[2], stride=2)
+        self.layer4 = self._make_layer(block, 512, num_blocks[3], stride=2)
+        self.linear = nn.Linear(512*block.expansion, num_classes)
+
+    def _make_layer(self, block, planes, num_blocks, stride):
+        strides = [stride] + [1]*(num_blocks-1)
+        layers = []
+        for stride in strides:
+            layers.append(block(self.in_planes, planes, stride))
+            self.in_planes = planes * block.expansion
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        out = F.relu(self.bn1(self.conv1(x)))
+        out = self.layer1(out)
+        out = self.layer2(out)
+        out = self.layer3(out)
+        out = self.layer4(out)
+        out = F.avg_pool2d(out, 4)
+        out = out.view(out.size(0), -1)
+        out = self.linear(out)
+        return out
+
+
+def ResNet18(num_classes):
+    return ResNet(BasicBlock, [2, 2, 2, 2],num_classes=num_classes)
+
+
+def ResNet34(num_classes):
+    return ResNet(BasicBlock, [3, 4, 6, 3],num_classes=num_classes)
+
+
+def ResNet50(num_classes):
+    return ResNet(Bottleneck, [3, 4, 6, 3],num_classes=num_classes)
+
+
+def ResNet101(num_classes):
+    return ResNet(Bottleneck, [3, 4, 23, 3],num_classes=num_classes)
+
+
+def ResNet152(num_classes):
+    return ResNet(Bottleneck, [3, 8, 36, 3],num_classes=num_classes)
\ No newline at end of file
diff --git a/e2eAIOK/common/trainer/model/model_builder_asr.py b/e2eAIOK/common/trainer/model/model_builder_asr.py
index a9fe31b59..1ed09443f 100644
--- a/e2eAIOK/common/trainer/model/model_builder_asr.py
+++ b/e2eAIOK/common/trainer/model/model_builder_asr.py
@@ -4,4 +4,4 @@
 
 class ModelBuilderASR(ModelBuilder):
     def init(self, cfg):
-        super().init(cfg)
\ No newline at end of file
+        super().init(cfg)
diff --git a/e2eAIOK/common/trainer/model/model_builder_cv.py b/e2eAIOK/common/trainer/model/model_builder_cv.py
index df934f1a6..8617fedac 100644
--- a/e2eAIOK/common/trainer/model/model_builder_cv.py
+++ b/e2eAIOK/common/trainer/model/model_builder_cv.py
@@ -3,37 +3,5 @@
 from e2eAIOK.common.trainer.model_builder import ModelBuilder
 
 class ModelBuilderCV(ModelBuilder):
-    def __init__(self, cfg):
-        super().__init__(cfg)
-    
-    def _init_model(self):
-        if self.cfg.best_model_structure != None:
-            with open(self.cfg.best_model_structure, 'r') as f:
-                arch = f.readlines()[-1]
-        else:
-            raise RuntimeError(f"model structure string not found")
-        
-        if self.cfg.domain == 'cnn':
-            model = DeMainNet(num_classes=self.cfg.num_classes, plainnet_struct=arch, no_create=False)
-        elif self.cfg.domain == 'vit':
-            model = Vision_TransformerSuper(img_size=self.cfg.input_size,
-                                patch_size=self.cfg.patch_size,
-                                embed_dim=self.cfg['SUPERNET']['EMBED_DIM'], depth=self.cfg['SUPERNET']['DEPTH'],
-                                num_heads=self.cfg['SUPERNET']['NUM_HEADS'],mlp_ratio=self.cfg['SUPERNET']['MLP_RATIO'],
-                                qkv_bias=True, drop_rate=self.cfg.drop,
-                                drop_path_rate=self.cfg.drop_path,
-                                gp=self.cfg.gp,
-                                num_classes=self.cfg.num_classes,
-                                max_relative_position=self.cfg.max_relative_position,
-                                relative_position=self.cfg.relative_position,
-                                change_qkv=self.cfg.change_qkv, abs_pos=not self.cfg.no_abs_pos)
-            depth, mlp_ratio, num_heads, embed_dim = decode_arch_tuple(arch)
-            model_config = {}
-            model_config['layer_num'] = depth
-            model_config['mlp_ratio'] = mlp_ratio
-            model_config['num_heads'] = num_heads
-            model_config['embed_dim'] = [embed_dim]*depth
-            n_parameters = model.get_sampled_params_numel(model_config)
-            print(f"model parameters size: {n_parameters}")
-        return model
-        
+    def __init__(self, cfg, model=None):
+        super().__init__(cfg, model)
diff --git a/e2eAIOK/common/trainer/model/model_builder_nlp.py b/e2eAIOK/common/trainer/model/model_builder_nlp.py
index 124a45c39..8e3df4c3b 100644
--- a/e2eAIOK/common/trainer/model/model_builder_nlp.py
+++ b/e2eAIOK/common/trainer/model/model_builder_nlp.py
@@ -4,4 +4,4 @@
 
 class ModelBuilderNLP(ModelBuilder):
     def init(self, cfg):
-        super().init(cfg)
\ No newline at end of file
+        super().init(cfg)
diff --git a/e2eAIOK/common/trainer/model/model_utils/model_utils.py b/e2eAIOK/common/trainer/model/model_utils/model_utils.py
new file mode 100644
index 000000000..34d50b919
--- /dev/null
+++ b/e2eAIOK/common/trainer/model/model_utils/model_utils.py
@@ -0,0 +1,32 @@
+import logging
+import re
+import torch
+import torch.nn as nn
+import os
+
+def initWeights(layer):
+    ''' Initialize layer parameters
+
+    :param layer: the layer to be initialized
+    :return:
+    '''
+    classname = layer.__class__.__name__
+    if classname.find('Conv2d') != -1 or classname.find('ConvTranspose2d') != -1:
+        nn.init.kaiming_uniform_(layer.weight)
+        if layer.bias is not None:
+            nn.init.zeros_(layer.bias)
+        logging.info("init layer [%s] with kaiming_uniform"% classname)
+    elif classname.find('BatchNorm') != -1:
+        nn.init.normal_(layer.weight, 1.0, 0.02)
+        if layer.bias is not None:
+            nn.init.zeros_(layer.bias)
+        logging.info("init layer [%s] with normal" % classname)
+    elif classname.find('Linear') != -1:
+        if hasattr(layer, "weight"):
+            nn.init.xavier_normal_(layer.weight)
+        if hasattr(layer, "bias") and layer.bias is not None:
+            nn.init.zeros_(layer.bias)
+        logging.info("init layer [%s] with xavier_normal" % classname)
+    else:
+        logging.info("no init layer[%s]"%classname)
+        # print("no init layer[%s]"%classname)    
diff --git a/e2eAIOK/common/trainer/model_builder.py b/e2eAIOK/common/trainer/model_builder.py
index b9d5c9d56..2055017cc 100644
--- a/e2eAIOK/common/trainer/model_builder.py
+++ b/e2eAIOK/common/trainer/model_builder.py
@@ -3,6 +3,7 @@
 import os
 import torch
 import e2eAIOK.common.trainer.utils.extend_distributed as ext_dist
+from e2eAIOK.common.trainer.utils.utils import get_device
 
 class ModelBuilder():
     """
@@ -35,26 +36,31 @@ def _post_process(self):
         """
         self.logger.info(f"model created: {self.model}")
 
-    def create_model(self):
+    def create_model(self,pretrain=None):
         """
             create model, load pre-trained model
+            :param pretrain: load pretrained model, if None, use cfg.pretrain as default
         """
         self._pre_process()
-        if self.model is not None and self.cfg.pretrain:
-            self.load_model()
-        else:
+        if self.model is not None:
+            ##  check if have pretrain model
+            cfg_pretrain = self.cfg.pretrain if ("pretrain" in self.cfg and self.cfg.pretrain != "") else None
+            pretrain = pretrain if pretrain is not None else cfg_pretrain
+            if pretrain:
+                self.load_model(pretrain)
+        if self.model is None:
             self.model = self._init_model()
 
         self._post_process()
         return self.model
 
-    def load_model(self):
+    def load_model(self,pretrain):
         """
             load pre-trained model
         """
-        if not os.path.exists(self.cfg.pretrain):
-            raise RuntimeError(f"Can not find {self.cfg.pretrain}!")
-        self.logger.info(f"loading pretrained model at {self.cfg.pretrain}")
-        state_dict = torch.load(self.cfg.pretrain, map_location=torch.device(self.cfg.device))
+        if not os.path.exists(pretrain):
+            raise RuntimeError(f"Can not find {pretrain}!")
+        self.logger.info(f"loading pretrained model at {pretrain}")
+        state_dict = torch.load(pretrain, map_location=torch.device(get_device(self.cfg)))
         state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict
         self.model.load_state_dict(state_dict, strict=True)
diff --git a/e2eAIOK/common/trainer/torch_trainer.py b/e2eAIOK/common/trainer/torch_trainer.py
index 2f113fa96..4cb5659ac 100644
--- a/e2eAIOK/common/trainer/torch_trainer.py
+++ b/e2eAIOK/common/trainer/torch_trainer.py
@@ -53,7 +53,8 @@ def _is_early_stop(self, metric):
         """
             check whether training achieved pre-defined metric threshold
         """
-        return metric >= self.cfg["metric_threshold"]
+        
+        return metric >= self.cfg["metric_threshold"] if "metric_threshold" in self.cfg else False
 
     def train_one_epoch(self, epoch):
         """
@@ -111,7 +112,7 @@ def evaluate(self, epoch):
             )
         self.logger.info(output_str)
 
-        if self.cfg.output_dir:
+        if "output_dir" in self.cfg:
             checkpoint_paths = [self.cfg.output_dir + '/checkpoint.pth']
             for checkpoint_path in checkpoint_paths:
                 utils.save_model({
diff --git a/e2eAIOK/common/trainer/utils/__init__.py b/e2eAIOK/common/trainer/utils/__init__.py
index de37cdcdf..49071abd2 100644
--- a/e2eAIOK/common/trainer/utils/__init__.py
+++ b/e2eAIOK/common/trainer/utils/__init__.py
@@ -1,1 +1,1 @@
-__all__ = ['early_stop','utils','extend_distributed']
\ No newline at end of file
+__all__ = ['utils','extend_distributed','data_utils', 'model_utils']
\ No newline at end of file
diff --git a/e2eAIOK/common/trainer/utils/early_stop.py b/e2eAIOK/common/trainer/utils/early_stop.py
deleted file mode 100644
index e69de29bb..000000000
diff --git a/e2eAIOK/common/trainer/utils/utils.py b/e2eAIOK/common/trainer/utils/utils.py
index 1d2c08927..6f51e7668 100644
--- a/e2eAIOK/common/trainer/utils/utils.py
+++ b/e2eAIOK/common/trainer/utils/utils.py
@@ -22,25 +22,9 @@
 from timm.utils import accuracy
 from collections import defaultdict, deque
 from easydict import EasyDict as edict
+from torch.utils.tensorboard import SummaryWriter
 import logging
-
-def is_main_process():
-    if ext_dist.my_size > 1:
-        return ext_dist.dist.get_rank() == 0
-    return 0
-
-def save_model(*args, **kwargs):
-    if is_main_process():
-        torch.save(*args, **kwargs)
-
-def init_log():
-    if ext_dist.my_size > 1:
-        if ext_dist.my_rank == 0:
-            logging.getLogger().setLevel(logging.INFO)
-        else:
-            logging.getLogger().setLevel(logging.WARNING)
-    else:
-        logging.getLogger().setLevel(logging.INFO)
+from torch.optim.lr_scheduler import _LRScheduler
 
 class SmoothedValue(object):
     """Track a series of values and provide access to smoothed values over a
@@ -102,7 +86,6 @@ def __str__(self):
             max=self.max,
             value=self.value)
 
-
 class MetricLogger(object):
     def __init__(self, delimiter="\t"):
         self.meters = defaultdict(SmoothedValue)
@@ -184,13 +167,119 @@ def log_every(self, iterable, print_freq, header=None):
         print('{} Total time: {} ({:.4f} s / it)'.format(
             header, total_time_str, total_time / len(iterable)))
 
-    
-def create_optimizer(model, cfg):
-    print(F"model:{model}")
-    if cfg.optimizer == "SGD":
-        optimizer = torch.optim.SGD(model.parameters(), lr=cfg.learning_rate,
-                momentum=cfg.momentum, weight_decay=cfg.weight_decay)
-    return optimizer
+class EarlyStopping():
+    ''' Early Stopping
+
+    '''
+    def __init__(self, tolerance_epoch = 5, delta=0, is_max = False, metric_threshold = None):
+        ''' Init method
+
+        :param tolerance_epoch: tolarance epoch
+        :param delta: delta for difference
+        :param is_max: max or min
+        :param metric_threshold: when metric up/down limiation then stop training. None means ignore.
+        '''
+
+        self._tolerance_epoch = tolerance_epoch
+        self._delta = delta
+        self._is_max = is_max
+        self._metric_threshold = metric_threshold
+
+        self._counter = 0
+        self.early_stop = False
+
+    def __call__(self, validation_metric, optimal_metric):
+        ############## absolute level #################
+        if self._metric_threshold is not None:
+            if (self._is_max and validation_metric >= self._metric_threshold) or\
+            ((not self._is_max) and validation_metric <= self._metric_threshold):
+                self.early_stop = True
+                print("Earlystop when meet metric_threshold [%s]"%self._metric_threshold)
+                logging.info("Earlystop when meet metric_threshold [%s]"%self._metric_threshold)
+                return
+        ############## relative level #################
+        if (self._is_max and (validation_metric < optimal_metric - self._delta)) \
+                or ((not self._is_max) and (validation_metric > optimal_metric + self._delta)): # less optimal
+            self._counter += 1
+        else: # more optimal
+            logging.info("Reset earlystop counter")
+            self._counter = 0
+        if self._counter >= self._tolerance_epoch:
+            self.early_stop = True
+
+    def __str__(self):
+        _str = 'EarlyStopping:%s\n'%self._tolerance_epoch
+        _str += '\tdelta:%s\n'%self._delta
+        _str += '\tis_max:%s\n' % self._is_max
+        _str += '\tmetric_threshold:%s\n' % self._metric_threshold
+        return _str
+
+class WarmUpLR(_LRScheduler):
+    """warmup_training learning rate scheduler
+    Args:
+        optimizer: optimzier(e.g. SGD)
+        total_iters: totoal_iters of warmup phase
+    """
+    def __init__(self, optimizer, total_iters, last_epoch=-1):
+
+        self.total_iters = total_iters
+        super().__init__(optimizer, last_epoch)
+
+    def get_lr(self):
+        """we will use the first m batches, and set the learning
+        rate to base_lr * m / total_iters
+        """
+        return [base_lr * self.last_epoch / (self.total_iters + 1e-8) for base_lr in self.base_lrs]
+
+class Timer:
+    ''' Timer to stat elapsed time
+
+    '''
+    def __enter__(self):
+        self.start = datetime.datetime.now()
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.end = datetime.datetime.now()
+        total_seconds = (self.end - self.start).total_seconds()
+        _str = "Total seconds:%s" % (total_seconds)
+        print(_str)
+        logging.info(_str)
+
+def get_device(cfg):
+    if 'device' in cfg and cfg.device in ['gpu', 'cuda'] and torch.cuda.is_available():
+        device = 'cuda'
+    else:
+        device = 'cpu'
+    return device
+
+def is_main_process():
+    if ext_dist.my_size > 1:
+        return ext_dist.dist.get_rank() == 0
+    return 0
+
+def save_model(*args, **kwargs):
+    if is_main_process():
+        torch.save(*args, **kwargs)
+
+def init_log():
+    if ext_dist.my_size > 1:
+        if ext_dist.my_rank == 0:
+            logging.getLogger().setLevel(logging.INFO)
+        else:
+            logging.getLogger().setLevel(logging.WARNING)
+    else:
+        logging.getLogger().setLevel(logging.INFO)
+
+
+def acc(output,label):
+    pred = output.data.cpu().max(1)[1]
+    label = label.data.cpu()
+    if label.shape == output.shape:
+        label = label.max(1)[1]
+
+    if pred.shape != label.shape:
+        logging.error('pred shape[%s] and label shape[%s] not match' % (pred.shape, label.shape))
+        raise RuntimeError('pred shape[%s] and label shape[%s] not match' % (pred.shape, label.shape))
+    return torch.mean((pred == label).float())
 
 def create_metric(cfg):
     if cfg.eval_metric == "accuracy":
@@ -200,9 +289,129 @@ def create_metric(cfg):
 def create_criterion(cfg):
     if cfg.criterion == "CrossEntropyLoss":
         criterion = torch.nn.CrossEntropyLoss()
+    else:
+        logging.error("[%s] is not supported"%cfg.criterion)
+        raise NotImplementedError("[%s] is not supported"%cfg.criterion)
     return criterion
 
+def create_optimizer(model=None, cfg=None, parameters=None):
+    ''' create optimizer
+
+    :param model: model, will be disabled if 'parameters' is not None
+    :param cfg: configurations
+    :param parameters: parameters for optimizer
+
+    :return: a optimizer
+    '''
+    logging.info(F"model:{model}")
+    logging.info(F"parameters:{parameters}")
+    parameters = parameters if parameters is not None else model.parameters()
+    if cfg.optimizer == "SGD":
+        optimizer = torch.optim.SGD(parameters, lr=cfg.learning_rate,
+                momentum=cfg.momentum, weight_decay=cfg.weight_decay)
+    else:
+        logging.error("[%s] is not supported"%cfg.optimizer)
+        raise NotImplementedError("[%s] is not supported"%cfg.optimizer)
+    return optimizer
+
 def create_scheduler(optimizer, cfg):
-    if cfg.lr_scheduler == "CosineAnnealingLR":
-        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=cfg.train_epochs)
+    ''' create scheduler
+
+    :return: a scheduler
+    '''
+    if "lr_scheduler" not in cfg or cfg.lr_scheduler == "":
+        scheduler = None
+    elif cfg.lr_scheduler == "ExponentialLR":
+        scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=cfg.lr_scheduler_config.decay_rate)
+    elif cfg.lr_scheduler == "MultiStepLR":
+        scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=cfg.lr_scheduler_config.decay_stages, gamma=cfg.lr_scheduler_config.decay_rate)
+    elif cfg.lr_scheduler == "CosineAnnealingLR":
+        t_max = cfg.lr_scheduler.T_max if "T_max" in cfg.lr_scheduler and cfg.lr_scheduler.T_max > 0 else cfg.train_epochs
+        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=t_max)
+    elif cfg.lr_scheduler == "ReduceLROnPlateau":
+        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode="max",factor=cfg.lr_scheduler_config.decay_rate, patience=cfg.lr_scheduler_config.decay_patience, \
+                                                        verbose=True, threshold_mode='abs')
+    else:
+        logging.error("[%s] is not supported"%cfg.lr_scheduler)
+        raise NotImplementedError("[%s] is not supported"%cfg.lr_scheduler)
     return scheduler
+
+def create_warmup_scheduler(optimizer, total_iters, cfg):
+    ''' create warm up scheduler
+
+    :param optimizer: optimizer
+    :param total_iters: iterations in one epoch
+    :return: a warm up scheduler
+    '''
+    if "warmup_scheduler" not in cfg or cfg.warmup_scheduler == "" or \
+        "warmup_scheduler_epoch" not in cfg or cfg.warmup_scheduler_epoch <= 0:
+        warmup_scheduler = None
+    elif cfg.warmup_scheduler == "WarmUpLR":
+        warmup_scheduler = WarmUpLR(optimizer, total_iters * cfg.warmup_scheduler_epoch)
+    else:
+        logging.error("[%s] is not supported"%cfg.warmup_scheduler)
+        raise NotImplementedError("[%s] is not supported"%cfg.warmup_scheduler)
+    
+    return warmup_scheduler
+
+def create_early_stopping(cfg):
+    ''' create early_stopping
+
+    :return: an early_stopping
+    '''
+    if "early_stop" not in cfg or cfg.early_stop == "":
+        early_stopping = None
+    elif cfg.early_stop == "EarlyStopping":
+        early_stopping = EarlyStopping(tolerance_epoch=cfg.early_stop_config.tolerance_epoch, delta=cfg.early_stop_config.delta, 
+                                    is_max=cfg.early_stop_config.is_max, metric_threshold=cfg.metric_threshold)
+        logging.info('early_stopping :%s' % early_stopping)
+    else:
+        logging.error("[%s] is not supported"%cfg.early_stop)
+        raise NotImplementedError("[%s] is not supported"%cfg.early_stop)
+    return early_stopping
+
+def create_tensorboard_writer(cfg):
+    ''' create tensorboard_writer
+
+    :return: a tensorboard_writer
+    '''
+    if "tensorboard_dir" not in cfg or cfg.tensorboard_dir == "":
+        tensorboard_writer  = None 
+    else:
+        tensorboard_writer = SummaryWriter(cfg.tensorboard_dir)
+        logging.info('tensorboard_writer :%s' % tensorboard_writer)
+    return tensorboard_writer
+
+def trace_handler(trace_file,p):
+    ''' profile trace handler
+
+    :param trace_file: output trace file
+    :param p: profiler
+    :return:
+    '''
+    logging.info("trace_output:%s" % p.key_averages().table(sort_by="cpu_time_total", row_limit=20))
+    p.export_chrome_trace(trace_file + str(p.step_num) + ".json")
+
+def create_profiler(cfg):
+    ''' create a profiler
+
+    :return: profiler
+    '''
+    if "profiler" not in cfg or not cfg.profiler:
+        return None
+    def parse_activities(activity_str):
+        result = set()
+        for item in activity_str.lower().split(","):
+            if item == 'cpu':
+                result.add(ProfilerActivity.CPU)
+            elif item == 'gpu' or item == 'cuda':
+                result.add(ProfilerActivity.CUDA)
+        return result
+
+    activities = parse_activities(get_device(cfg))
+    schedule = torch.profiler.schedule(skip_first=cfg.profiler_config.skip_first,wait=cfg.profiler_config.wait,
+        warmup=cfg.profiler_config.warmup,active=cfg.profiler_config.active,repeat=cfg.profiler_config.repeat)
+    profiler = profile(activities=activities,schedule=schedule,
+                        on_trace_ready=partial(trace_handler,cfg.profiler_config.trace_file))
+    logging.info("profiler:%s"%profiler)
+    return profiler
\ No newline at end of file
diff --git a/e2eAIOK/common/utils.py b/e2eAIOK/common/utils.py
index bdcd7c869..987368288 100644
--- a/e2eAIOK/common/utils.py
+++ b/e2eAIOK/common/utils.py
@@ -14,6 +14,11 @@ def update_list(orig, diff):
             orig[i] = dict_diff[orig[i]['name']]
     return orig
 
+def check_mkdir(path):
+    if not os.path.isdir(path):
+        os.makedirs(path)
+    return path
+
 def mkdir(dest_path):
     new_name = datetime.now().strftime("%Y%m%d_%H%M%S")
     new_path = os.path.join(dest_path, new_name)
@@ -88,3 +93,22 @@ def parse_size(size):
 def get_estimate_size_of_dtype(dtype_name):
     units = {'byte': 1, 'short': 2, 'int': 4, 'long': 8, 'float': 4, 'double': 8, 'string': 10}
     return units[dtype_name] if dtype_name in units else 4
+
+## Update multi-level dict, merge dict2 into dict1
+def update_dict(dict1, dict2):
+    keys = set(dict1.keys()) | set(dict2.keys())
+    dict3 = {}
+    for key in keys:
+        if key in dict1.keys() and key in dict2.keys():
+            if isinstance(dict1[key], dict) and isinstance(dict2[key], dict):
+                dict3[key] = update_dict(dict1[key],dict2[key])
+            elif (isinstance(dict1[key], dict) and not isinstance(dict2[key], dict)) or \
+                    (not isinstance(dict1[key], dict) and isinstance(dict2[key], dict)):
+                raise ValueError(f"{key} in two dicts have different types, one is dict, another is not!")
+            else:
+                dict3[key] = dict2[key]
+        elif key in dict1.keys():
+            dict3[key] = dict1[key]
+        else:
+            dict3[key] = dict2[key]
+    return dict3