Upload 14 files

configs/data_configs.py

@@ -0,0 +1,58 @@
+def get_dataset_class(dataset_name):
+    """
+    :param dataset_name:
+    :return: The algorithm class with the given name
+    """
+    if dataset_name not in globals():
+        raise NotImplementedError("Dataset not found : {}".format(dataset_name))
+    return globals()[dataset_name]
+
+
+class mit:
+    def __init__(self):
+        super().__init__()
+        # data params
+        self.num_classes = 5
+        self.class_names = ['N', 'S', 'V', 'F', 'Q']
+        self.sequence_len = 186
+
+        # model configs
+        self.input_channels = 1
+        self.kernel_size = 8
+        self.stride = 1
+        self.dropout = 0.2
+
+        # features
+        self.mid_channels = 32
+        self.final_out_channels = 128
+
+        # MoE Transformer
+        self.trans_dim = 25
+        self.num_heads = 5
+        self.num_experts = 8
+
+
+class ptb:
+    def __init__(self):
+        super().__init__()
+        # data params
+        self.num_classes = 2
+        self.class_names = ['normal', 'abnormal']
+        self.sequence_len = 188
+
+        # model_configs
+        self.input_channels = 1
+        self.kernel_size = 8
+        self.stride = 1
+        self.dropout = 0.2
+
+        # features
+        self.mid_channels = 32
+        self.final_out_channels = 128
+
+        # MoE Transformer
+        self.trans_dim = 25
+        self.num_heads = 5
+        self.num_experts = 8
+
+

configs/hparams.py

@@ -0,0 +1,17 @@
+def get_hparams_class(dataset_name):
+    """Return the algorithm class with the given name."""
+    if dataset_name not in globals():
+        raise NotImplementedError("Algorithm not found: {}".format(dataset_name))
+    return globals()[dataset_name]
+
+
+class Supervised:
+    def __init__(self):
+        super(Supervised, self).__init__()
+        self.train_params = {
+            'num_epochs': 100,
+            'batch_size': 128,
+            'weight_decay': 1e-4,
+            'learning_rate': 1e-3,
+            'feature_dim': 1*128
+        }

data/dataloader.py

@@ -0,0 +1,56 @@
+import torch
+from torch.utils.data import DataLoader
+from data.dataset import load_ECG_Dataset
+
+import os
+import numpy as np
+import math
+
+
+def get_class_weight(labels_dict):
+    total = sum(labels_dict.values())
+    max_num = max(labels_dict.values())
+    mu = 1.0 / (total / max_num)
+    class_weight = dict()
+    for key, value in labels_dict.items():
+        score = math.log(mu * total / float(value))
+        class_weight[key] = score if score > 1.0 else 1.0
+    return class_weight
+
+
+class ECGDataloader:
+    testdata_path: str
+    traindata_path: str
+    valdata_path: str
+
+    def __init__(self, data_path, data_type, hparams):
+        self.traindata_path = os.path.join(data_path, data_type, f'train.pt')
+        self.testdata_path = os.path.join(data_path, data_type, f'test.pt')
+        self.validdata_path = os.path.join(data_path, data_type, f'val.pt')
+        self.batch_size = hparams['batch_size']
+
+    def train_dataloader(self):
+        train_dataset = torch.load(self.traindata_path)
+        train_dataset = load_ECG_Dataset(train_dataset)
+        cw = train_dataset.y_data.numpy().tolist()
+        cw_dict = {}
+        for i in range(len(np.unique(train_dataset.y_data.numpy()))):
+            cw_dict[i] = cw.count(i)
+
+        train_loader = DataLoader(dataset=train_dataset, batch_size=self.batch_size, shuffle=True,
+                                  drop_last=True, num_workers=4)
+        return train_loader, get_class_weight(cw_dict)
+
+    def test_dataloader(self):
+        test_dataset = torch.load(self.testdata_path)
+        test_dataset = load_ECG_Dataset(test_dataset)
+        test_loader = DataLoader(dataset=test_dataset, batch_size=self.batch_size, shuffle=False,
+                                 drop_last=False, num_workers=4)
+        return test_loader
+
+    def valid_dataloader(self):
+        valid_dataset = torch.load(self.validdata_path)
+        valid_dataset = load_ECG_Dataset(valid_dataset)
+        valid_loader = DataLoader(dataset=valid_dataset, batch_size=self.batch_size, shuffle=True,
+                                  drop_last=False, num_workers=4)
+        return valid_loader

data/dataset.py

@@ -0,0 +1,39 @@
+import torch
+from torch.utils.data import Dataset
+
+import numpy as np
+
+
+class load_ECG_Dataset(Dataset):
+    # Initialize dataset
+    def __init__(self, dataset):
+
+        # Load sample
+        x_data = dataset["samples"]
+        # Convert to pytorch tensor
+        if isinstance(x_data, np.ndarray):
+            x_data = torch.from_numpy(x_data)
+
+        # Load labels
+        y_data = dataset.get("labels")
+        if y_data is not None and isinstance(y_data, np.ndarray):
+            y_data = torch.from_numpy(y_data)
+
+        self.x_data = x_data.float()
+        self.y_data = y_data.long() if y_data is not None else None
+
+        self.len = x_data.shape[0]
+
+    def get_labels(self):
+        return self.y_data
+
+    def __getitem__(self, idx):
+        sample = {
+            'samples': self.x_data[idx].squeeze(-1),
+            'labels': int(self.y_data[idx])
+        }
+        return sample
+
+    def __len__(self):
+        return self.len
+

main.py

@@ -0,0 +1,35 @@
+import os
+import argparse
+import warnings
+from train import Trainer
+import sklearn.exceptions
+warnings.filterwarnings("ignore", category=sklearn.exceptions.UndefinedMetricWarning)
+
+parser = argparse.ArgumentParser()
+
+# ========  Experiments Name ================
+parser.add_argument('--save_dir',               default='experiments_logs',         type=str,
+                    help='Directory containing all experiments')
+parser.add_argument('--experiment_description', default='Exp1',   type=str, help='experiment name')
+parser.add_argument('--run_description',        default='run1',     type=str, help='run name')
+
+# ========= Select the DATASET ==============
+parser.add_argument('--dataset',                default='mit',           type=str, help='mit, ptb')
+parser.add_argument('--seed_id',                default='0',             type=str,
+                    help='to fix a seed while training')
+
+# ========= Experiment settings ===============
+parser.add_argument('--data_path',              default=r'/Users/splendor1811/datn/ECGTransForm/datasets',
+                    type=str, help='Path containing dataset')
+
+parser.add_argument('--num_runs',               default=1,                 type=int,
+                    help='Number of consecutive run with different seeds')
+parser.add_argument('--device',                 default='mps',            type=str,
+                    help='cpu or cuda')
+
+
+args = parser.parse_args()
+
+if __name__ == "__main__":
+    trainer = Trainer(args)
+    trainer.train()

models/MoE_ECGFormer.py

@@ -0,0 +1,144 @@
+import torch
+from torch import nn
+
+
+class MoE_ECGFormer(nn.Module):
+    def __init__(self, configs, hparams):
+        super().__init__()
+
+        filter_sizes = [5, 9, 11]
+        self.conv1 = nn.Conv1d(configs.input_channels, configs.mid_channels, kernel_size=filter_sizes[0],
+                               stride=configs.stride, bias=False, padding=(filter_sizes[0] // 2))
+
+        self.conv2 = nn.Conv1d(configs.input_channels, configs.mid_channels, kernel_size=filter_sizes[1],
+                               stride=configs.stride, bias=False, padding=(filter_sizes[1] // 2))
+
+        self.conv3 = nn.Conv1d(configs.input_channels, configs.mid_channels, kernel_size=filter_sizes[2],
+                               stride=configs.stride, bias=False, padding=(filter_sizes[2] // 2))
+
+        self.bn = nn.BatchNorm1d(configs.mid_channels)
+        self.relu = nn.ReLU()
+        self.maxpool = nn.MaxPool1d(kernel_size=2, stride=2, padding=1)
+        self.dropout = nn.Dropout(configs.dropout)
+
+        self.conv_block2 = nn.Sequential(
+            nn.Conv1d(configs.mid_channels, configs.mid_channels * 2, kernel_size=8, stride=1, bias=False,
+                      padding=4),
+            nn.BatchNorm1d(configs.mid_channels * 2),
+            nn.ReLU(),
+            nn.MaxPool1d(kernel_size=2, stride=2, padding=1)
+        )
+
+        self.conv_block3 = nn.Sequential(
+            nn.Conv1d(configs.mid_channels * 2, configs.final_out_channels, kernel_size=8, stride=1, bias=False,
+                      padding=4),
+            nn.BatchNorm1d(configs.final_out_channels),
+            nn.ReLU(),
+            nn.MaxPool1d(kernel_size=2, stride=2, padding=1),
+        )
+
+        self.inplanes = 128
+        self.crm = self._make_layer(SEBasicBlock, 128, 3)
+
+        # Transformer_layer
+        self.encoder_layer = nn.TransformerEncoderLayer(d_model=configs.trans_dim,
+                                                        nhead=configs.num_heads,
+                                                        batch_first=True)
+        self.transformer_encoder = nn.TransformerEncoder(self.encoder_layer, num_layers=3)
+        self.aap = nn.AdaptiveAvgPool1d(1)
+        self.clf = nn.Linear(hparams['feature_dim'], configs.num_classes)
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv1d(self.inplanes, planes * block.expansion,
+                          kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm1d(planes * block.expansion),
+            )
+
+        layers = [block(self.inplanes, planes, stride, downsample)]
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x_in):
+
+        # Multi-scale Convolutions
+        x1 = self.conv1(x_in)
+        x2 = self.conv2(x_in)
+        x3 = self.conv3(x_in)
+
+        x_concat = torch.mean(torch.stack([x1, x2, x3], dim=2), dim=2)
+        x_concat = self.dropout(self.maxpool(self.relu(self.bn(x_concat))))
+
+        x = self.conv_block2(x_concat)
+        x = self.conv_block3(x)
+
+        # Channel Recalibration Module
+        x = self.crm(x)
+
+        # Bidirectional MoE Transformer
+        x1 = self.transformer_encoder(x)
+        x2 = self.transformer_encoder(torch.flip(x, [2]))
+        x = x1 + x2
+
+        x = self.aap(x)
+        x_flat = x.reshape(x.size(0), -1)
+        x_out = self.clf(x_flat)
+        return x_out
+
+
+class SELayer(nn.Module):
+    def __init__(self, channel, reduction=4):
+        super(SELayer, self).__init__()
+        self.avg_pool = nn.AdaptiveAvgPool1d(1)
+        self.fc = nn.Sequential(
+            nn.Linear(channel, channel // reduction, bias=False),
+            nn.ReLU(inplace=True),
+            nn.Linear(channel // reduction, channel, bias=False),
+            nn.Sigmoid()
+        )
+
+    def forward(self, x):
+        b, c, _ = x.size()
+        y = self.avg_pool(x).view(b, c)
+        y = self.fc(y).view(b, c, 1)
+        return x * y.expand_as(x)
+
+
+class SEBasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
+                 base_width=64, dilation=1, norm_layer=None,
+                 *, reduction=4):
+        super(SEBasicBlock, self).__init__()
+        self.conv1 = nn.Conv1d(inplanes, planes, stride)
+        self.bn1 = nn.BatchNorm1d(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = nn.Conv1d(planes, planes, 1)
+        self.bn2 = nn.BatchNorm1d(planes)
+        self.se = SELayer(planes, reduction)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.se(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out

test.py

@@ -0,0 +1,3 @@
+from deepface import DeepFace
+
+print(DeepFace.verify('/Users/splendor1811/Desktop/a1.jpg' ,'/Users/splendor1811/Desktop/a1.jpg'))

train.py

@@ -0,0 +1,192 @@
+from models.MoE_ECGFormer import MoE_ECGFormer
+from data.dataloader import ECGDataloader
+from configs.data_configs import get_dataset_class
+from configs.hparams import get_hparams_class
+from utils import AverageMeter, to_device, _save_metrics, copy_files
+from utils import fix_randomness, starting_logs, save_checkpoint, _calc_metrics
+import torch
+import torch.nn.functional as F
+import datetime
+import os
+import collections
+import numpy as np
+
+import warnings
+import sklearn.exceptions
+
+warnings.filterwarnings("ignore", category=sklearn.exceptions.UndefinedMetricWarning)
+warnings.simplefilter(action='ignore', category=FutureWarning)
+
+
+class Trainer(object):
+    def __init__(self, args):
+        # dataset parameters
+        self.dataset = args.dataset
+        self.seed_id = args.seed_id
+        self.device = torch.device(args.device)
+
+        # Exp Description
+        self.run_description = f"{args.run_description}_{datetime.datetime.now().strftime('%H_%M')}"
+        self.experiment_description = args.experiment_description
+
+        # paths
+        self.home_path = os.getcwd()
+        self.save_dir = os.path.join(os.getcwd(), "experiments_logs")
+        self.exp_log_dir = os.path.join(self.save_dir, self.experiment_description, self.run_description)
+        os.makedirs(self.exp_log_dir, exist_ok=True)
+
+        self.data_path = args.data_path
+
+        # Specify runs
+        self.num_runs = args.num_runs
+
+        # get dataset and base model configs
+        self.dataset_configs, self.hparams_class = self.get_configs()
+
+        # Specify hparams
+        self.hparams = self.hparams_class.train_params
+
+    def get_configs(self):
+        dataset_class = get_dataset_class(self.dataset)
+        hparams_class = get_hparams_class("Supervised")
+        return dataset_class(), hparams_class()
+
+    def load_data(self, data_type):
+        self.train_dl, self.cw_dict = ECGDataloader(self.data_path, data_type, self.hparams).train_dataloader()
+        self.test_dl = ECGDataloader(self.data_path, data_type, self.hparams).test_dataloader()
+        self.valid_dl = ECGDataloader(self.data_path, data_type, self.hparams).valid_dataloader()
+
+    def calc_results_per_run(self):
+        acc, f1 = _calc_metrics(self.pred_labels, self.true_labels, self.dataset_configs.class_names)
+        return acc, f1
+
+    def train(self):
+        copy_files(self.exp_log_dir)  # save a copy of training files
+
+        self.metrics = {'accuracy': [], 'f1_score': []}
+
+        # fixing random seed
+        fix_randomness(int(self.seed_id))
+
+        # Logging
+        self.logger, self.scenario_log_dir = starting_logs(self.dataset, self.exp_log_dir, self.seed_id)
+        self.logger.debug(self.hparams)
+
+        # Load data
+        self.load_data(self.dataset)
+
+        model = MoE_ECGFormer(configs=self.dataset_configs, hparams=self.hparams)
+        model.to(self.device)
+
+        # Average meters
+        loss_avg_meters = collections.defaultdict(lambda: AverageMeter())
+
+        self.optimizer = torch.optim.Adam(
+            model.parameters(),
+            lr=self.hparams["learning_rate"],
+            weight_decay=self.hparams["weight_decay"],
+            betas=(0.9, 0.99)
+        )
+
+        weights = [float(value) for value in self.cw_dict.values()]
+        # Now convert the list of floats to a numpy array, then to a PyTorch tensor
+        weights_array = np.array(weights).astype(np.float32)  # Ensuring the correct dtype
+        weights_tensor = torch.tensor(weights_array).to(self.device)
+        self.cross_entropy = torch.nn.CrossEntropyLoss(weight=weights_tensor)
+
+        best_acc = 0
+        best_f1 = 0
+
+        # training..
+        ts_acc = 0
+        ts_f1 = 0
+        for epoch in range(1, self.hparams["num_epochs"] + 1):
+            model.train()
+
+            for step, batches in enumerate(self.train_dl):
+                batches = to_device(batches, self.device)
+
+                data = batches['samples'].float()
+                labels = batches['labels'].long()
+
+                # ====== Source =====================
+                self.optimizer.zero_grad()
+
+                # Src original features
+                logits = model(data)
+
+                # Cross-Entropy loss
+                x_ent_loss = self.cross_entropy(logits, labels)
+
+                x_ent_loss.backward()
+                self.optimizer.step()
+
+                losses = {'Total_loss': x_ent_loss.item()}
+                for key, val in losses.items():
+                    loss_avg_meters[key].update(val, self.hparams["batch_size"])
+
+            self.evaluate(model, self.valid_dl)
+            tr_acc, tr_f1 = self.calc_results_per_run()
+            # logging
+            self.logger.debug(f'[Epoch : {epoch}/{self.hparams["num_epochs"]}]')
+            for key, val in loss_avg_meters.items():
+                self.logger.debug(f'{key}\t: {val.avg:2.4f}')
+            self.logger.debug(f'TRAIN: Acc:{tr_acc:2.4f} \t F1:{tr_f1:2.4f}')
+
+            # VALIDATION part
+            self.evaluate(model, self.valid_dl)
+            ts_acc, ts_f1 = self.calc_results_per_run()
+            if ts_f1 > best_f1:  # save best model based on best f1.
+                best_f1 = ts_f1
+                best_acc = ts_acc
+                save_checkpoint(self.exp_log_dir, model, self.dataset, self.dataset_configs, self.hparams, "best")
+                _save_metrics(self.pred_labels, self.true_labels, self.exp_log_dir, "validation_best")
+
+            # logging
+            self.logger.debug(f'VAL  : Acc:{ts_acc:2.4f} \t F1:{ts_f1:2.4f} (best: {best_f1:2.4f})')
+            self.logger.debug(f'-------------------------------------')
+
+            # LAST EPOCH
+        _save_metrics(self.pred_labels, self.true_labels, self.exp_log_dir, "validation_last")
+        self.logger.debug("LAST EPOCH PERFORMANCE on validation set...")
+        self.logger.debug(f'Acc:{ts_acc:2.4f} \t F1:{ts_f1:2.4f}')
+
+        self.logger.debug(":::::::::::::")
+        # BEST EPOCH
+        self.logger.debug("BEST EPOCH PERFORMANCE on validation set ...")
+        self.logger.debug(f'Acc:{best_acc:2.4f} \t F1:{best_f1:2.4f}')
+        save_checkpoint(self.exp_log_dir, model, self.dataset, self.dataset_configs, self.hparams, "last")
+
+        # TESTING
+        print(" === Evaluating on TEST set ===")
+        self.evaluate(model, self.test_dl)
+        test_acc, test_f1 = self.calc_results_per_run()
+        _save_metrics(self.pred_labels, self.true_labels, self.exp_log_dir, "test_last")
+        self.logger.debug(f'Acc:{test_acc:2.4f} \t F1:{test_f1:2.4f}')
+
+    def evaluate(self, model, dataset):
+        model.to(self.device).eval()
+
+        total_loss_ = []
+
+        self.pred_labels = np.array([])
+        self.true_labels = np.array([])
+
+        with torch.no_grad():
+            for batches in dataset:
+                batches = to_device(batches, self.device)
+                data = batches['samples'].float()
+                labels = batches['labels'].long()
+
+                # forward pass
+                predictions = model(data)
+
+                # compute loss
+                loss = F.cross_entropy(predictions, labels)
+                total_loss_.append(loss.item())
+                pred = predictions.detach().argmax(dim=1)  # get the index of the max log-probability
+
+                self.pred_labels = np.append(self.pred_labels, pred.cpu().numpy())
+                self.true_labels = np.append(self.true_labels, labels.data.cpu().numpy())
+
+        self.trg_loss = torch.tensor(total_loss_).mean()  # average loss

utils.py

@@ -0,0 +1,204 @@
+import torch
+import random
+import os
+import sys
+import logging
+import numpy as np
+import pandas as pd
+from shutil import copy
+from datetime import datetime
+import matplotlib.pyplot as plt
+import collections
+import umap
+import umap.plot
+from matplotlib.colors import ListedColormap
+
+
+from sklearn.metrics import classification_report, accuracy_score
+
+
+def count_parameters(model):
+    return sum(p.numel() for p in model.parameters() if p.requires_grad)
+
+
+class AverageMeter(object):
+    """Computes and stores the average and current value"""
+
+    def __init__(self):
+        self.reset()
+
+    def reset(self):
+        self.val = 0
+        self.avg = 0
+        self.sum = 0
+        self.count = 0
+
+    def update(self, val, n=1):
+        self.val = val
+        self.sum += val * n
+        self.count += n
+        self.avg = self.sum / self.count
+
+
+def fix_randomness(SEED):
+    random.seed(SEED)
+    np.random.seed(SEED)
+    torch.manual_seed(SEED)
+    torch.cuda.manual_seed(SEED)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+
+def _logger(logger_name, level=logging.DEBUG):
+    """
+    Method to return a custom logger with the given name and level
+    :param logger_name:
+    :param level:
+    :return:
+    """
+    logger = logging.getLogger(logger_name)
+    logger.setLevel(level)
+    format_string = "%(message)s"
+    log_format = logging.Formatter(format_string)
+
+    console_handler = logging.StreamHandler(sys.stdout)
+    console_handler.setFormatter(log_format)
+    logger.addHandler(console_handler)
+    # Creating and adding the file handler
+    file_handler = logging.FileHandler(logger_name, mode='a')
+    file_handler.setFormatter(log_format)
+    logger.addHandler(file_handler)
+    return logger
+
+
+def starting_logs(data_type, exp_log_dir, seed_id):
+    log_dir = os.path.join(exp_log_dir, "_seed_"+ str(seed_id))
+    os.makedirs(log_dir,exist_ok=True)
+    log_file_name = os.path.join(log_dir, f"logs_{datetime.now().strftime('%d_%m_%Y_%H_%M_%S')}.log")
+    logger = _logger(log_file_name)
+    logger.debug('=' * 45)
+    logger.debug(f'Logging data type {data_type}')
+    logger.debug("=" * 45)
+    logger.debug(f'Logging seed id {seed_id}')
+    logger.debug("=" * 45)
+    return logger, log_dir
+
+
+def save_checkpoint(exp_log_dir, model, dataset, dataset_configs, hparams, status):
+    save_dict = {
+        "dataset": dataset,
+        "configs": dataset_configs.__dict__,
+        "hparams": dict(hparams),
+        "model": model.state_dict(),
+    }
+    # save classification report
+    save_path = os.path.join(exp_log_dir, f"checkpoint_{status}.pt")
+    torch.save(save_dict, save_path)
+
+
+def _calc_metrics(pred_labels, true_labels, classes_names):
+    pred_labels = np.array(pred_labels).astype(int)
+    true_labels = np.array(true_labels).astype(int)
+
+    r = classification_report(true_labels, pred_labels, target_names=classes_names, digits=6, output_dict=True)
+    accuracy = accuracy_score(true_labels, pred_labels)
+
+    return accuracy * 100, r["macro avg"]["f1-score"] * 100
+
+
+def _save_metrics(pred_labels, true_labels, log_dir, status):
+    pred_labels = np.array(pred_labels).astype(int)
+    true_labels = np.array(true_labels).astype(int)
+
+    r = classification_report(true_labels, pred_labels, digits=6, output_dict=True)
+
+    df = pd.DataFrame(r)
+    accuracy = accuracy_score(true_labels, pred_labels)
+    df["accuracy"] = accuracy
+    df = df * 100
+
+    # save classification report
+    file_name = f"classification_report_{status}.xlsx"
+    report_save_path = os.path.join(log_dir, file_name)
+    df.to_excel(report_save_path)
+
+
+def to_device(input, device):
+    if torch.is_tensor(input):
+        return input.to(device=device)
+    elif isinstance(input, str):
+        return input
+    elif isinstance(input, collections.abc.Mapping):
+        return {k: to_device(sample, device=device) for k, sample in input.items()}
+    elif isinstance(input, collections.abc.Sequence):
+        return [to_device(sample, device=device) for sample in input]
+    else:
+        raise TypeError("Input must contain tensor, dict or list, found {type(input)}")
+
+
+def copy_files(destination):
+    destination_dir = os.path.join(destination, "MODEL_BACKUP_FILES")
+    os.makedirs(destination_dir, exist_ok=True)
+    copy("main.py", os.path.join(destination_dir, "main.py"))
+    copy("data/dataloader.py", os.path.join(destination_dir, "dataloader.py"))
+    copy("data/dataset.py", os.path.join(destination_dir, "dataset.py"))
+    copy(f"models/MoE_ECGFormer.py", os.path.join(destination_dir, f"models.py"))
+    copy(f"configs/data_configs.py", os.path.join(destination_dir, f"data_configs.py"))
+    copy(f"configs/hparams.py", os.path.join(destination_dir, f"hparams.py"))
+    copy(f"train.py", os.path.join(destination_dir, f"train.py"))
+    copy("utils.py", os.path.join(destination_dir, "utils.py"))
+
+
+def _plot_umap(model, data_loader, device, save_dir):
+
+    classes_names = ['N', 'S', 'V', 'F', 'Q']
+
+    font = {'family': 'Times New Roman',
+            'weight': 'bold',
+            'size': 17}
+    plt.rc('font', **font)
+
+    with torch.no_grad():
+        # Source flow
+        data = data_loader.dataset.x_data.float().to(device)
+        labels = data_loader.dataset.y_data.view((-1)).long()
+        out = model[0](data)
+        features = model[1](out)
+
+    if not os.path.exists(os.path.join(save_dir, "umap_plots")):
+        os.mkdir(os.path.join(save_dir, "umap_plots"))
+
+    # cmaps = plt.get_cmap('jet')
+    model_reducer = umap.UMAP()  # n_neighbors=3, min_dist=0.3, metric='correlation', random_state=42)
+    embedding = model_reducer.fit_transform(features.detach().cpu().numpy())
+
+    # Normalize the labels to [0, 1] for colormap
+    norm_labels = labels / 4.0
+
+    # Create a new colormap by extracting the first 5 colors from "Paired"
+    paired = plt.cm.get_cmap('Paired', 12)  # 12 distinct colors
+    new_colors = [paired(0), paired(1), paired(2), paired(4),
+                  paired(6)]  # Skip every second color, but take both from the first pair
+    new_cmap = ListedColormap(new_colors)
+
+    print("Plotting UMAP ...")
+    plt.figure(figsize=(16, 10))
+    # scatter = plt.scatter(embedding[:, 0], embedding[:, 1], c=labels,  s=10, cmap='Spectral')
+    scatter = plt.scatter(embedding[:, 0], embedding[:, 1], c=norm_labels, cmap=new_cmap, s=15)
+
+    handles, _ = scatter.legend_elements(prop='colors')
+    plt.legend(handles, classes_names, title="Classes")
+    file_name = "umap_.png"
+    fig_save_name = os.path.join(save_dir, "umap_plots", file_name)
+    plt.xticks([])
+    plt.yticks([])
+    plt.savefig(fig_save_name, bbox_inches='tight')
+    plt.close()
+
+
+
+
+
+
+
+