Repository: torcheeg/torcheeg (language: python)

Issue:
CCNN cross subjects



Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index f583206..0e66b3f 100644
--- a/README.md
+++ b/README.md
@@ -7,7 +7,7 @@
 [gitter-image]: https://badges.gitter.im/torcheeg/community.svg
 [gitter-url]: https://gitter.im/torcheeg/community?utm_source=badge&utm_medium=badge&utm_campaign=pr-badge&utm_content=badge
 
-![TorchEEG Logo](https://github.com/tczhangzhi/torcheeg/blob/main/docs/source/_static/torcheeg_logo_dark.png)
+![TorchEEG Logo](https://github.com/torcheeg/torcheeg/blob/main/docs/source/_static/torcheeg_logo_dark.png)
 
 --------------------------------------------------------------------------------
 
@@ -16,7 +16,7 @@
 [![Downloads][pepy-image]][pepy-url]
 [![Join the Chat][gitter-image]][gitter-url]
 
-**[Documentation](https://torcheeg.readthedocs.io/)** | **[TorchEEG Examples](https://github.com/tczhangzhi/torcheeg/tree/main/examples)**
+**[Documentation](https://torcheeg.readthedocs.io/)** | **[TorchEEG Examples](https://github.com/torcheeg/torcheeg/tree/main/examples)**
 
 TorchEEG is a library built on PyTorch for EEG signal analysis. TorchEEG aims to provide a plug-and-play EEG analysis tool, so that researchers can quickly reproduce EEG analysis work and start new EEG analysis research without paying attention to technical details unrelated to the research focus.
 
@@ -63,7 +63,7 @@ pip install torcheeg
 In case you want to experiment with the latest TorchEEG features which are not fully released yet, please run the following command to install from the main branch on github:
 
 ```bash
-pip install git+https://github.com/tczhangzhi/torcheeg.git
+pip install git+https://github.com/torcheeg/torcheeg.git
 ```
 
 ### Plugin
@@ -138,25 +138,32 @@ We list currently supported datasets, transforms, data splitting, and deep learn
 
 ## Quickstart
 
-In this quick tour, we highlight the ease of starting an EEG analysis research with only modifying a few lines of [PyTorch tutorial](https://pytorch.org/tutorials/beginner/basics/quickstart_tutorial.html).
+In this quick tour, we highlight the ease of starting an EEG analysis research with only a few lines.
 
 The `torcheeg.datasets` module contains dataset classes for many real-world EEG datasets. In this tutorial, we use the `DEAP` dataset. We first go to the official website to apply for data download permission according to the introduction of [DEAP dataset](https://www.eecs.qmul.ac.uk/mmv/datasets/deap/), and download the dataset. Next, we need to specify the download location of the dataset in the `root_path` parameter. For the DEAP dataset, we specify the path to the `data_preprocessed_python` folder, e.g. `./tmp_in/data_preprocessed_python`.
 
 ```python
 from torcheeg.datasets import DEAPDataset
-from torcheeg.datasets.constants.emotion_recognition.deap import DEAP_CHANNEL_LOCATION_DICT
-
-dataset = DEAPDataset(io_path=f'./tmp_out/deap',
-                      root_path='./tmp_in/data_preprocessed_python',
-                      offline_transform=transforms.Compose(
-                          [transforms.BandDifferentialEntropy(),
-                           transforms.ToGrid(DEAP_CHANNEL_LOCATION_DICT)]),
-                      online_transform=transforms.Compose([transforms.BaselineRemoval(),
-                                                           transforms.ToTensor()]),
-                      label_transform=transforms.Compose([
-                          transforms.Select('valence'),
-                          transforms.Binary(5.0),
-                      ]), num_worker=4)
+from torcheeg import transforms
+
+from torcheeg.datasets.constants.emotion_recognition.deap import \
+    DEAP_CHANNEL_LOCATION_DICT
+
+dataset = DEAPDataset(
+    io_path=f'./tmp_out/examples_pipeline/deap',
+    root_path='./tmp_in/data_preprocessed_python',
+    offline_transform=transforms.Compose([
+        transforms.BandDifferentialEntropy(apply_to_baseline=True),
+        transforms.ToGrid(DEAP_CHANNEL_LOCATION_DICT, apply_to_baseline=True)
+    ]),
+    online_transform=transforms.Compose(
+        [transforms.BaselineRemoval(),
+         transforms.ToTensor()]),
+    label_transform=transforms.Compose([
+        transforms.Select('valence'),
+        transforms.Binary(5.0),
+    ]),
+    num_worker=8)
 ```
 
 The `DEAPDataset` API further contains three parameters: `online_transform`, `offline_transform`, and `label_transform`, which are used to modify samples and labels, respectively.
@@ -168,123 +175,47 @@ Here, `offline_transform` will only be called once when the dataset is initializ
 Next, we need to divide the dataset into a training set and a test set. In the field of EEG analysis, commonly used data partitioning methods include k-fold cross-validation and leave-one-out cross-validation. In this tutorial, we use k-fold cross-validation on the entire dataset (`KFold`) as an example of dataset splitting.
 
 ```python
-from torcheeg.model_selection import KFold
+from torcheeg.datasets import KFoldGroupbyTrial
 
-k_fold = KFold(n_splits=10,
-                      split_path=f'./tmp_out/split',
-                      shuffle=True,
-                      random_state=42)
+k_fold = KFoldGroupbyTrial(n_splits=10,
+                           split_path='./tmp_out/examples_pipeline/split',
+                           shuffle=True,
+                           random_state=42)
 ```
 
-Let\'s define a simple but effective CNN model according to [CCNN](https://link.springer.com/chapter/10.1007/978-3-030-04239-4_39):
+We loop through each cross-validation set, and for each one, we initialize the CCNN model and define its hyperparameters. For instance, each EEG sample contains 4-channel features from 4 sub-bands, and the grid size is 9x9. We then train the model for 50 epochs using the `ClassifierTrainer`.
 
 ```python
-class CNN(torch.nn.Module):
-    def __init__(self, in_channels=4, num_classes=3):
-        super().__init__()
-        self.conv1 = nn.Sequential(
-            nn.ZeroPad2d((1, 2, 1, 2)),
-            nn.Conv2d(in_channels, 64, kernel_size=4, stride=1),
-            nn.ReLU()
-        )
-        self.conv2 = nn.Sequential(
-            nn.ZeroPad2d((1, 2, 1, 2)),
-            nn.Conv2d(64, 128, kernel_size=4, stride=1),
-            nn.ReLU()
-        )
-        self.conv3 = nn.Sequential(
-            nn.ZeroPad2d((1, 2, 1, 2)),
-            nn.Conv2d(128, 256, kernel_size=4, stride=1),
-            nn.ReLU()
-        )
-        self.conv4 = nn.Sequential(
-            nn.ZeroPad2d((1, 2, 1, 2)),
-            nn.Conv2d(256, 64, kernel_size=4, stride=1),
-            nn.ReLU()
-        )
-
-        self.lin1 = nn.Linear(9 * 9 * 64, 1024)
-        self.lin2 = nn.Linear(1024, num_classes)
-
-    def forward(self, x):
-        x = self.conv1(x)
-        x = self.conv2(x)
-        x = self.conv3(x)
-        x = self.conv4(x)
-
-        x = x.flatten(start_dim=1)
-        x = self.lin1(x)
-        x = self.lin2(x)
-        return x
-```
-
-Specify the device and loss function used during training and test.
-
-```python
-device = "cuda" if torch.cuda.is_available() else "cpu"
-loss_fn = nn.CrossEntropyLoss()
-batch_size = 64
-```
-
-The training and validation scripts for the model are taken from the  [PyTorch tutorial](https://pytorch.org/tutorials/beginner/basics/quickstart_tutorial.html) without much modification. Usually, the value of `batch` contains two parts; the first part refers to the result of `online_transform`, which generally corresponds to the `Tensor` sequence representing EEG signals. The second part refers to the result of `label_transform`, a sequence of integers representing the label.
+from torch.utils.data import DataLoader
+from torcheeg.models import CCNN
 
-```python
-def train(dataloader, model, loss_fn, optimizer):
-    size = len(dataloader.dataset)
-    model.train()
-    for batch_idx, batch in enumerate(dataloader):
-        X = batch[0].to(device)
-        y = batch[1].to(device)
-
-        # Compute prediction error
-        pred = model(X)
-        loss = loss_fn(pred, y)
-
-        # Backpropagation
-        optimizer.zero_grad()
-        loss.backward()
-        optimizer.step()
-
-        if batch_idx % 100 == 0:
-            loss, current = loss.item(), batch_idx * len(X)
-            print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")
-
-
-def valid(dataloader, model, loss_fn):
-    size = len(dataloader.dataset)
-    num_batches = len(dataloader)
-    model.eval()
-    val_loss, correct = 0, 0
-    with torch.no_grad():
-        for batch in dataloader:
-            X = batch[0].to(device)
-            y = batch[1].to(device)
-
-            pred = model(X)
-            val_loss += loss_fn(pred, y).item()
-            correct += (pred.argmax(1) == y).type(torch.float).sum().item()
-    val_loss /= num_batches
-    correct /= size
-    print(f"Test Error: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {val_loss:>8f} \n")
-```
+from torcheeg.trainers import ClassifierTrainer
 
-Traverse `k` folds and train the model separately for testing. It is worth noting that, in general, we need to specify `shuffle=True` for the `DataLoader` of the training data set to avoid the deviation of the model training caused by consecutive labels of the same category.
+import pytorch_lightning as pl
 
-```python
 for i, (train_dataset, val_dataset) in enumerate(k_fold.split(dataset)):
-
-    model = CNN().to(device)
-    optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
-
-    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
-    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)
-
-    epochs = 50
-    for t in range(epochs):
-        print(f"Epoch {t+1}\n-------------------------------")
-        train(train_loader, model, loss_fn, optimizer)
-        valid(val_loader, model, loss_fn)
-    print("Done!")
+    train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
+    val_loader = DataLoader(val_dataset, batch_size=64, shuffle=False)
+
+    model = CCNN(num_classes=2, in_channels=4, grid_size=(9, 9))
+
+    trainer = ClassifierTrainer(model=model,
+                                num_classes=2,
+                                lr=1e-4,
+                                weight_decay=1e-4,
+                                accelerator="gpu")
+    trainer.fit(train_loader,
+                val_loader,
+                max_epochs=50,
+                default_root_dir=f'./tmp_out/examples_pipeline/model/{i}',
+                callbacks=[pl.callbacks.ModelCheckpoint(save_last=True)],
+                enable_progress_bar=True,
+                enable_model_summary=True,
+                limit_val_batches=0.0)
+    score = trainer.test(val_loader,
+                         enable_progress_bar=True,
+                         enable_model_summary=True)[0]
+    print(f'Fold {i} test accuracy: {score["test_accuracy"]:.4f}')
 ```
 
 For more specific usage of each module, please refer to [the documentation]((https://torcheeg.readthedocs.io/)).
@@ -328,4 +259,4 @@ If you would like to contribute new datasets, deep learning methods, and extensi
 
 ## License
 
-TorchEEG has a MIT license, as found in the [LICENSE](https://github.com/tczhangzhi/torcheeg/blob/main/LICENSE) file.
\ No newline at end of file
+TorchEEG has a MIT license, as found in the [LICENSE](https://github.com/torcheeg/torcheeg/blob/main/LICENSE) file.
\ No newline at end of file
diff --git a/docs/source/generated/torcheeg.datasets.MNEDataset.rst b/docs/source/generated/torcheeg.datasets.MNEDataset.rst
deleted file mode 100644
index 80b8a85..0000000
--- a/docs/source/generated/torcheeg.datasets.MNEDataset.rst
+++ /dev/null
@@ -1,13 +0,0 @@
-﻿.. role:: hidden
-    :class: hidden-section
-.. currentmodule:: torcheeg.datasets
-
-
-MNEDataset
-==========
-
-.. autoclass:: MNEDataset
-
-..
-  autogenerated from source/_templates/classtemplate.rst
-  note it does not have :inherited-members:, :members:
\ No newline at end of file
diff --git a/docs/source/generated/torcheeg.datasets.NumpyDataset.rst b/docs/source/generated/torcheeg.datasets.NumpyDataset.rst
deleted file mode 100644
index 968d362..0000000
--- a/docs/source/generated/torcheeg.datasets.NumpyDataset.rst
+++ /dev/null
@@ -1,13 +0,0 @@
-﻿.. role:: hidden
-    :class: hidden-section
-.. currentmodule:: torcheeg.datasets
-
-
-NumpyDataset
-============
-
-.. autoclass:: NumpyDataset
-
-..
-  autogenerated from source/_templates/classtemplate.rst
-  note it does not have :inherited-members:, :members:
\ No newline at end of file
diff --git a/docs/source/implementing_gnns_with_pyg.rst b/docs/source/implementing_gnns_with_pyg.rst
deleted file mode 100644
index 522066b..0000000
--- a/docs/source/implementing_gnns_with_pyg.rst
+++ /dev/null
@@ -1,216 +0,0 @@
-Implementing GNNs with PyG
-==========================
-
-In this quick tour, we'll take a closer look at how to bring together
-TorchEEG and `PyG
-(pytorch_geometric) <https://github.com/pyg-team/pytorch_geometric>`__
-to implement graph convolutional networks.
-
-Define the Dataset
-~~~~~~~~~~~~~~~~~~
-
-The ``torcheeg.datasets`` module contains dataset classes for many
-real-world EEG datasets. In this tutorial, we use the ``SEED`` dataset.
-We first go to the official website to apply for data download
-permission according to the introduction of `SEED
-dataset <https://bcmi.sjtu.edu.cn/home/seed/>`__, and download the
-dataset. Next, we need to specify the download location of the dataset
-in the ``root_path`` parameter. For the SEED dataset, we specify the
-path to the ``Preprocessed_EEG`` folder,
-e.g. ``./tmp_in/Preprocessed_EEG``.
-
-.. code:: python
-
-    from torcheeg import transforms
-    from torcheeg.transforms.pyg import ToG
-    from torcheeg.datasets import SEEDDataset
-    from torcheeg.datasets.constants.emotion_recognition.seed import \
-       SEED_ADJACENCY_MATRIX
-
-    dataset = SEEDDataset(io_path=f'./tmp_out/seed',
-                         root_path='./tmp_in/Preprocessed_EEG',
-                         offline_transform=transforms.BandDifferentialEntropy(),
-                         online_transform=ToG(SEED_ADJACENCY_MATRIX),
-                         label_transform=transforms.Compose([
-                             transforms.Select('emotion'),
-                             transforms.Lambda(lambda x: int(x) + 1),
-                         ]),
-                         num_worker=4)
-
-The ``SEEDDataset`` API further contains three parameters:
-``online_transform``, ``offline_transform``, and ``label_transform``,
-which are used to modify samples and labels, respectively.
-
-Here, ``offline_transform`` will only be called once when the dataset is
-initialized to preprocess all samples in the dataset, and the processed
-dataset will be stored in ``io_path`` to avoid time-consuming repeated
-transformations in subsequent use. If offline preprocessing is a
-computationally intensive operation, we also recommend setting multi-CPU
-parallelism for offline_transform, e.g., set ``num_worker`` to 4.
-
-``online_transform`` is used to transform samples on the fly. Please use
-``online_transform`` if you don't want to wait for the preprocessing of
-the entire dataset (suitable for scenarios where new ``transform``
-algorithms are designed) or expect data transformation with randomness
-each time a sample is indexed.
-
-To convert raw data in numpy format into a graph representation
-acceptable to PyG (``torch_geometric.data.Data``), TorchEEG provides the
-``transforms.ToG`` . Here, electrodes correspond to nodes in the graph
-structure, and the associations between electrodes are defined as edges
-and weights on edges. The commonly considered associations are spatial
-adjacency and functional connection. Here, we use the adjacency matrix
-``SEED_ADJACENCY_MATRIX`` defined based on the spatial neighbor
-relationship of electrodes. Each value in the adjacency matrix indicates
-whether two corresponding electrodes are adjacent in a 10-20 system, 1
-for adjacent and 0 for non-adjacent.
-
-Define the Data Splitting Method
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-Next, we need to divide the dataset into a training set and a test set.
-In the field of EEG analysis, commonly used data partitioning methods
-include k-fold cross-validation and leave-one-out cross-validation. In
-this tutorial, we use k-fold cross-validation per subject
-(``KFoldPerSubjectGroupbyTrial``) as an example of dataset splitting.
-
-.. code:: python
-
-   from torcheeg.model_selection import KFold
-
-   k_fold = KFoldPerSubjectGroupbyTrial(n_splits=10,
-                                 split_path=f'./tmp_out/split',
-                                 shuffle=False)
-
-For more data splitting methods, please refer to
-https://torcheeg.readthedocs.io/en/latest/torcheeg.model_selection.html
-
-Define the Model
-~~~~~~~~~~~~~~~~
-
-Let's define a simple but effective GNN model based on the convolutional
-layers and operation provided by PyG:
-
-.. code:: python
-
-   from torch_geometric.nn import GATConv, global_mean_pool
-
-   class GNN(torch.nn.Module):
-       def __init__(self, in_channels=4, num_layers=3, hid_channels=64, num_classes=3):
-           super().__init__()
-           self.conv1 = GATConv(in_channels, hid_channels)
-           self.convs = torch.nn.ModuleList()
-           for _ in range(num_layers - 1):
-               self.convs.append(GATConv(hid_channels, hid_channels))
-           self.lin1 = Linear(hid_channels, hid_channels)
-           self.lin2 = Linear(hid_channels, num_classes)
-
-       def reset_parameters(self):
-           self.conv1.reset_parameters()
-           for conv in self.convs:
-               conv.reset_parameters()
-           self.lin1.reset_parameters()
-           self.lin2.reset_parameters()
-
-       def forward(self, data):
-           x, edge_index, batch = data.x, data.edge_index, data.batch
-           x = F.relu(self.conv1(x, edge_index))
-           for conv in self.convs:
-               x = F.relu(conv(x, edge_index))
-           x = global_mean_pool(x, batch)
-           x = F.relu(self.lin1(x))
-           x = F.dropout(x, p=0.5, training=self.training)
-           x = self.lin2(x)
-           return x
-
-For more models, please refer to
-https://torcheeg.readthedocs.io/en/latest/torcheeg.models.html
-
-Define the Training and Test Process
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-Specify the device and loss function used during training and test.
-
-.. code:: python
-
-   device = "cuda" if torch.cuda.is_available() else "cpu"
-   loss_fn = nn.CrossEntropyLoss()
-   batch_size = 64
-
-The training and validation scripts for the model are taken from the
-`PyTorch
-tutorial <https://pytorch.org/tutorials/beginner/basics/quickstart_tutorial.html>`__
-without much modification. Usually, the value of ``batch`` contains two
-parts; the first part refers to the result of ``online_transform``,
-which generally corresponds to the ``Data`` sequence representing EEG
-graphs. The second part refers to the result of ``label_transform``, a
-sequence of integers representing the label.
-
-.. code:: python
-
-   def train(dataloader, model, loss_fn, optimizer):
-       size = len(dataloader.dataset)
-       model.train()
-       for batch_idx, batch in enumerate(dataloader):
-           X = batch[0].to(device)
-           y = batch[1].to(device)
-
-           # Compute prediction error
-           pred = model(X)
-           loss = loss_fn(pred, y)
-
-           # Backpropagation
-           optimizer.zero_grad()
-           loss.backward()
-           optimizer.step()
-
-           if batch_idx % 100 == 0:
-               loss, current = loss.item(), batch_idx * len(X)
-               print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")
-
-
-   def valid(dataloader, model, loss_fn):
-       size = len(dataloader.dataset)
-       num_batches = len(dataloader)
-       model.eval()
-       val_loss, correct = 0, 0
-       with torch.no_grad():
-           for batch in dataloader:
-               X = batch[0].to(device)
-               y = batch[1].to(device)
-
-               pred = model(X)
-               val_loss += loss_fn(pred, y).item()
-               correct += (pred.argmax(1) == y).type(torch.float).sum().item()
-       val_loss /= num_batches
-       correct /= size
-       print(f"Test Error: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {val_loss:>8f} \n")
-
-Traverse ``k`` folds and train the model separately for testing. It
-should be noted that the ``Dataloader`` here needs to use the
-implementation in ``PyG`` instead of ``torch``, in order to organize the
-``Data`` data structure into ``Batch``.
-
-It is also worth noting that, in general, we need to specify
-``shuffle=True`` for the ``DataLoader`` of the training data set to
-avoid the deviation of the model training caused by consecutive labels
-of the same category.
-
-.. code:: python
-
-   from torch_geometric.loader import DataLoader
-
-   for i, (train_dataset, val_dataset) in enumerate(k_fold.split(dataset)):
-       
-       model = GNN().to(device)
-       optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
-
-       train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
-       val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)
-
-       epochs = 50
-       for t in range(epochs):
-           print(f"Epoch {t+1}\n-------------------------------")
-           train(train_loader, model, loss_fn, optimizer)
-           valid(val_loader, model, loss_fn)
-       print("Done!")
\ No newline at end of file
diff --git a/docs/source/index.rst b/docs/source/index.rst
index 26c47bf..610e52f 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -38,9 +38,12 @@ Transformers.
    :maxdepth: 1
    :caption: Tutorials:
 
-   introduction_by_example
-   implementing_gnns_with_pyg
-   managing_experiments_with_pytorch_lightning
+   Deep Learning Workflow <auto_examples/examples_pipeline>
+   Introduction to datasets <auto_examples/examples_datasets>
+   Introduction to transforms <auto_examples/examples_transforms>
+   Introduction to models <auto_examples/examples_models>
+   Introduction to trainers <auto_examples/examples_trainers>
+   Vanilla PyTorch Workflow <auto_examples/examples_vanilla_torch>
    Expand to See More Examples <auto_examples/index>
 
 .. toctree::
diff --git a/docs/source/introduction_by_example.rst b/docs/source/introduction_by_example.rst
deleted file mode 100644
index 3f0c3a4..0000000
--- a/docs/source/introduction_by_example.rst
+++ /dev/null
@@ -1,209 +0,0 @@
-Introduction by Example
-=======================
-
-In this quick tour, we highlight the ease of starting an EEG analysis
-research with only modifying a few lines of `PyTorch
-tutorial <https://pytorch.org/tutorials/beginner/basics/quickstart_tutorial.html>`__.
-
-Define the Dataset
-~~~~~~~~~~~~~~~~~~
-
-The ``torcheeg.datasets`` module contains dataset classes for many
-real-world EEG datasets. In this tutorial, we use the ``DEAP`` dataset.
-We first go to the official website to apply for data download
-permission according to the introduction of `DEAP
-dataset <https://www.eecs.qmul.ac.uk/mmv/datasets/deap/>`__, and
-download the dataset. Next, we need to specify the download location of
-the dataset in the ``root_path`` parameter. For the DEAP dataset, we
-specify the path to the ``data_preprocessed_python`` folder,
-e.g. ``./tmp_in/data_preprocessed_python``.
-
-.. code:: python
-
-    from torcheeg import transforms
-    from torcheeg.datasets import DEAPDataset
-    from torcheeg.datasets.constants.emotion_recognition.deap import DEAP_CHANNEL_LOCATION_DICT
-
-    dataset = DEAPDataset(io_path=f'./tmp_out/deap',
-                         root_path='./tmp_in/data_preprocessed_python',
-                         offline_transform=transforms.Compose(
-                             [transforms.BandDifferentialEntropy(apply_to_baseline=True),
-                              transforms.ToGrid(DEAP_CHANNEL_LOCATION_DICT, apply_to_baseline=True)]),
-                         online_transform=transforms.Compose([transforms.BaselineRemoval(),
-                                                              transforms.ToTensor()]),
-                         label_transform=transforms.Compose([
-                             transforms.Select('valence'),
-                             transforms.Binary(5.0),
-                         ]), num_worker=4)
-
-The ``DEAPDataset`` API further contains three parameters:
-``online_transform``, ``offline_transform``, and ``label_transform``,
-which are used to modify samples and labels, respectively.
-
-Here, ``offline_transform`` will only be called once when the dataset is
-initialized to preprocess all samples in the dataset, and the processed
-dataset will be stored in ``io_path`` to avoid time-consuming repeated
-transformations in subsequent use. If offline preprocessing is a
-computationally intensive operation, we also recommend setting multi-CPU
-parallelism for offline_transform, e.g., set ``num_worker`` to 4.
-
-``online_transform`` is used to transform samples on the fly. Please use
-``online_transform`` if you don't want to wait for the preprocessing of
-the entire dataset (suitable for scenarios where new ``transform``
-algorithms are designed) or expect data transformation with randomness
-each time a sample is indexed.
-
-For more datasets, please refer to
-https://torcheeg.readthedocs.io/en/latest/torcheeg.datasets.html.
-
-Define the Data Splitting Method
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-Next, we need to divide the dataset into a training set and a test set.
-In the field of EEG analysis, commonly used data partitioning methods
-include k-fold cross-validation and leave-one-out cross-validation. In
-this tutorial, we use k-fold cross-validation on the entire dataset
-(``KFold``) as an example of dataset splitting.
-
-.. code:: python
-
-   from torcheeg.model_selection import KFold
-
-   k_fold = KFold(n_splits=10,
-                         split_path=f'./tmp_out/split',
-                         shuffle=True,
-                         random_state=42)
-
-For more data splitting methods, please refer to
-https://torcheeg.readthedocs.io/en/latest/torcheeg.model_selection.html
-
-Define the Model
-~~~~~~~~~~~~~~~~
-
-Let's define a simple but effective CNN model according to
-`CCNN <https://link.springer.com/chapter/10.1007/978-3-030-04239-4_39>`__:
-
-.. code:: python
-
-   class CNN(torch.nn.Module):
-       def __init__(self, in_channels=4, num_classes=3):
-           super().__init__()
-           self.conv1 = nn.Sequential(
-               nn.ZeroPad2d((1, 2, 1, 2)),
-               nn.Conv2d(in_channels, 64, kernel_size=4, stride=1),
-               nn.ReLU()
-           )
-           self.conv2 = nn.Sequential(
-               nn.ZeroPad2d((1, 2, 1, 2)),
-               nn.Conv2d(64, 128, kernel_size=4, stride=1),
-               nn.ReLU()
-           )
-           self.conv3 = nn.Sequential(
-               nn.ZeroPad2d((1, 2, 1, 2)),
-               nn.Conv2d(128, 256, kernel_size=4, stride=1),
-               nn.ReLU()
-           )
-           self.conv4 = nn.Sequential(
-               nn.ZeroPad2d((1, 2, 1, 2)),
-               nn.Conv2d(256, 64, kernel_size=4, stride=1),
-               nn.ReLU()
-           )
-
-           self.lin1 = nn.Linear(9 * 9 * 64, 1024)
-           self.lin2 = nn.Linear(1024, num_classes)
-
-       def forward(self, x):
-           x = self.conv1(x)
-           x = self.conv2(x)
-           x = self.conv3(x)
-           x = self.conv4(x)
-
-           x = x.flatten(start_dim=1)
-           x = self.lin1(x)
-           x = self.lin2(x)
-           return x
-
-For more models, please refer to
-https://torcheeg.readthedocs.io/en/latest/torcheeg.models.html
-
-Define the Training and Test Process
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-Specify the device and loss function used during training and test.
-
-.. code:: python
-
-   device = "cuda" if torch.cuda.is_available() else "cpu"
-   loss_fn = nn.CrossEntropyLoss()
-   batch_size = 64
-
-The training and validation scripts for the model are taken from the
-`PyTorch
-tutorial <https://pytorch.org/tutorials/beginner/basics/quickstart_tutorial.html>`__
-without much modification. Usually, the value of ``batch``
-contains two parts; the first part refers to the result of
-``online_transform``, which generally corresponds to the ``Tensor``
-sequence representing EEG signals. The second part refers to the result
-of ``label_transform``, a sequence of integers representing the label.
-
-.. code:: python
-
-   def train(dataloader, model, loss_fn, optimizer):
-       size = len(dataloader.dataset)
-       model.train()
-       for batch_idx, batch in enumerate(dataloader):
-           X = batch[0].to(device)
-           y = batch[1].to(device)
-
-           # Compute prediction error
-           pred = model(X)
-           loss = loss_fn(pred, y)
-
-           # Backpropagation
-           optimizer.zero_grad()
-           loss.backward()
-           optimizer.step()
-
-           if batch_idx % 100 == 0:
-               loss, current = loss.item(), batch_idx * len(X)
-               print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")
-
-
-   def valid(dataloader, model, loss_fn):
-       size = len(dataloader.dataset)
-       num_batches = len(dataloader)
-       model.eval()
-       val_loss, correct = 0, 0
-       with torch.no_grad():
-           for batch in dataloader:
-               X = batch[0].to(device)
-               y = batch[1].to(device)
-
-               pred = model(X)
-               val_loss += loss_fn(pred, y).item()
-               correct += (pred.argmax(1) == y).type(torch.float).sum().item()
-       val_loss /= num_batches
-       correct /= size
-       print(f"Test Error: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {val_loss:>8f} \n")
-
-Traverse ``k`` folds and train the model separately for testing. It is
-worth noting that, in general, we need to specify ``shuffle=True`` for
-the ``DataLoader`` of the training data set to avoid the deviation of
-the model training caused by consecutive labels of the same category.
-
-.. code:: python
-
-   for i, (train_dataset, val_dataset) in enumerate(k_fold.split(dataset)):
-
-       model = CNN().to(device)
-       optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
-
-       train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
-       val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)
-
-       epochs = 50
-       for t in range(epochs):
-           print(f"Epoch {t+1}\n-------------------------------")
-           train(train_loader, model, loss_fn, optimizer)
-           valid(val_loader, model, loss_fn)
-       print("Done!")
\ No newline at end of file
diff --git a/docs/source/managing_experiments_with_pytorch_lightning.rst b/docs/source/managing_experiments_with_pytorch_lightning.rst
deleted file mode 100644
index c72571f..0000000
--- a/docs/source/managing_experiments_with_pytorch_lightning.rst
+++ /dev/null
@@ -1,212 +0,0 @@
-Managing Experiments with PyTorch Lightning
-===========================================
-
-In this quick tour, we'll take a closer look at how to bring together
-TorchEEG and `PyTorch
-Lightning <https://github.com/PyTorchLightning/pytorch-lightning>`__ to
-organize and run experiments using multiple GPUs.
-
-Define the Dataset
-~~~~~~~~~~~~~~~~~~
-
-The ``torcheeg.datasets`` module contains dataset classes for many
-real-world EEG datasets. In this tutorial, we use the ``DEAP`` dataset.
-We first go to the official website to apply for data download
-permission according to the introduction of `DEAP
-dataset <https://www.eecs.qmul.ac.uk/mmv/datasets/deap/>`__, and
-download the dataset. Next, we need to specify the download location of
-the dataset in the ``root_path`` parameter. For the DEAP dataset, we
-specify the path to the ``data_preprocessed_python`` folder,
-e.g. ``./tmp_in/data_preprocessed_python``.
-
-.. code:: python
-
-    from torcheeg import transforms
-    from torcheeg.datasets import DEAPDataset
-    from torcheeg.datasets.constants.emotion_recognition.deap import DEAP_CHANNEL_LOCATION_DICT
-
-    dataset = DEAPDataset(io_path=f'./tmp_out/deap',
-                         root_path='./tmp_in/data_preprocessed_python',
-                         offline_transform=transforms.Compose(
-                             [transforms.BandDifferentialEntropy(apply_to_baseline=True),
-                              transforms.ToGrid(DEAP_CHANNEL_LOCATION_DICT, apply_to_baseline=True)]),
-                         online_transform=transforms.Compose([transforms.BaselineRemoval(),
-                                                              transforms.ToTensor()]),
-                         label_transform=transforms.Compose([
-                             transforms.Select('valence'),
-                             transforms.Binary(5.0),
-                         ]), num_worker=4)
-
-The ``DEAPDataset`` API further contains three parameters:
-``online_transform``, ``offline_transform``, and ``label_transform``,
-which are used to modify samples and labels, respectively.
-
-Here, ``offline_transform`` will only be called once when the dataset is
-initialized to preprocess all samples in the dataset, and the processed
-dataset will be stored in ``io_path`` to avoid time-consuming repeated
-transformations in subsequent use. If offline preprocessing is a
-computationally intensive operation, we also recommend setting multi-CPU
-parallelism for offline_transform, e.g., set ``num_worker`` to 4.
-
-``online_transform`` is used to transform samples on the fly. Please use
-``online_transform`` if you don't want to wait for the preprocessing of
-the entire dataset (suitable for scenarios where new ``transform``
-algorithms are designed) or expect data transformation with randomness
-each time a sample is indexed.
-
-For more datasets, please refer to
-https://torcheeg.readthedocs.io/en/latest/torcheeg.datasets.html.
-
-Define the Data Splitting Method
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-Next, we need to divide the dataset into a training set and a test set.
-In the field of EEG analysis, commonly used data partitioning methods
-include k-fold cross-validation and leave-one-out cross-validation. In
-this tutorial, we use k-fold cross-validation on the entire dataset
-(``KFold``) as an example of dataset splitting.
-
-.. code:: python
-
-   from torcheeg.model_selection import KFold
-
-   k_fold = KFold(n_splits=10,
-                         split_path=f'./tmp_out/split',
-                         shuffle=True,
-                         random_state=42)
-
-For more data splitting methods, please refer to
-https://torcheeg.readthedocs.io/en/latest/torcheeg.model_selection.html
-
-Define the Model
-~~~~~~~~~~~~~~~~
-
-Let's define a simple but effective CNN model according to
-`CCNN <https://link.springer.com/chapter/10.1007/978-3-030-04239-4_39>`__:
-
-.. code:: python
-
-   class CNN(torch.nn.Module):
-       def __init__(self, in_channels=4, num_classes=3):
-           super().__init__()
-           self.conv1 = nn.Sequential(
-               nn.ZeroPad2d((1, 2, 1, 2)),
-               nn.Conv2d(in_channels, 64, kernel_size=4, stride=1),
-               nn.ReLU()
-           )
-           self.conv2 = nn.Sequential(
-               nn.ZeroPad2d((1, 2, 1, 2)),
-               nn.Conv2d(64, 128, kernel_size=4, stride=1),
-               nn.ReLU()
-           )
-           self.conv3 = nn.Sequential(
-               nn.ZeroPad2d((1, 2, 1, 2)),
-               nn.Conv2d(128, 256, kernel_size=4, stride=1),
-               nn.ReLU()
-           )
-           self.conv4 = nn.Sequential(
-               nn.ZeroPad2d((1, 2, 1, 2)),
-               nn.Conv2d(256, 64, kernel_size=4, stride=1),
-               nn.ReLU()
-           )
-
-           self.lin1 = nn.Linear(9 * 9 * 64, 1024)
-           self.lin2 = nn.Linear(1024, num_classes)
-
-       def forward(self, x):
-           x = self.conv1(x)
-           x = self.conv2(x)
-           x = self.conv3(x)
-           x = self.conv4(x)
-
-           x = x.flatten(start_dim=1)
-           x = self.lin1(x)
-           x = self.lin2(x)
-           return x
-
-For more models, please refer to
-https://torcheeg.readthedocs.io/en/latest/torcheeg.models.html
-
-Define the Training and Test Process
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-A simple case implemented according to the official documentation of
-`pytorch-lightning <https://github.com/PyTorchLightning/pytorch-lightning>`__.
-Here, ``__init__``, ``forward``, ``training_step``, ``validation_step``,
-and ``configure_optimizers`` need to be implemented, where ``__init__``
-is used to specify hyperparameters and initialize related dependencies,
-``forward`` is used to define the forward propagation process of the
-network, ``training_step`` and ``validation_step`` are used to define
-the training and testing process. Usually, the value of ``batch``
-contains two parts; the first part refers to the result of
-``online_transform``, which generally corresponds to the ``Tensor``
-sequence representing EEG signals. The second part refers to the result
-of ``label_transform``, a sequence of integers representing the label.
-Besides, ``configure_optimizers`` is used to define the required
-optimizers and schedulers.
-
-.. code:: python
-
-   class EEGClassifier(LightningModule):
-       def __init__(self, model, lr=1e-4):
-           super().__init__()
-           self.save_hyperparameters(ignore="model")
-           self.model = model
-           self.val_acc = Accuracy()
-
-       def forward(self, x):
-           return self.model(x)
-
-       def training_step(self, batch, batch_idx):
-           X = batch[0]
-           y = batch[1]
-
-           logits = self.forward(X)
-           loss = F.cross_entropy(logits, y.long())
-           return loss
-
-       def validation_step(self, batch, batch_idx):
-           X = batch[0]
-           y = batch[1]
-
-           logits = self.forward(X)
-           loss = F.cross_entropy(logits, y.long())
-
-           self.val_acc(logits, y)
-           self.log("val_acc", self.val_acc)
-           self.log("val_loss", loss)
-
-       def configure_optimizers(self):
-           optimizer = torch.optim.Adam(self.model.parameters(), lr=self.hparams.lr)
-
-           return [optimizer], []
-
-Traverse ``k`` folds and train the model separately for testing. It is
-worth noting that, in general, we need to specify ``shuffle=True`` for
-the ``DataLoader`` of the training data set to avoid the deviation of
-the model training caused by consecutive labels of the same category.
-
-.. code:: python
-
-   for i, (train_dataset, val_dataset) in enumerate(k_fold.split(dataset)):
-           train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
-           val_loader = DataLoader(val_dataset, batch_size=64, shuffle=False)
-           tb_logger = TensorBoardLogger(save_dir='lightning_logs', name=f'fold_{i + 1}')
-           checkpoint_callback = ModelCheckpoint(dirpath=tb_logger.log_dir,
-                                                 filename="{epoch:02d}-{val_metric:.4f}",
-                                                 monitor='val_metric',
-                                                 mode='max')
-
-           model = EEGClassifier(CNN())
-
-           trainer = Trainer(max_epochs=50,
-                             devices=2,
-                             accelerator="auto",
-                             strategy="ddp",
-                             checkpoint_callback=checkpoint_callback,
-                             logger=tb_logger)
-
-           trainer.fit(model, train_loader, val_loader)
-
-For multi-GPU parallel training, just define the number of GPUs using
-``device`` and set ``strategy="ddp"``.
\ No newline at end of file
diff --git a/docs/source/torcheeg.datasets.rst b/docs/source/torcheeg.datasets.rst
index 262b442..6770487 100644
--- a/docs/source/torcheeg.datasets.rst
+++ b/docs/source/torcheeg.datasets.rst
@@ -65,9 +65,8 @@ Customized Datasets
     :nosignatures:
     :template: datasettemplate.rst
 
-    NumpyDataset
-    MNEDataset
     FolderDataset
+    CSVFolderDataset
 
 Hooks
 ---------------------------------------------------------
diff --git a/docs/source/torcheeg.trainers.rst b/docs/source/torcheeg.trainers.rst
index 5e9b8b1..d40b2bd 100644
--- a/docs/source/torcheeg.trainers.rst
+++ b/docs/source/torcheeg.trainers.rst
@@ -59,4 +59,5 @@ As the cost of data collection decreases, the difficulty of obtaining unlabeled
    :nosignatures:
    :template: trainertemplate.rst
 
-   SimCLRTrainer
\ No newline at end of file
+   SimCLRTrainer
+   BYOLTrainer
\ No newline at end of file
diff --git a/docs/source/torcheeg.transforms.rst b/docs/source/torcheeg.transforms.rst
index 7ced0d6..36559ea 100644
--- a/docs/source/torcheeg.transforms.rst
+++ b/docs/source/torcheeg.transforms.rst
@@ -58,6 +58,8 @@ Numpy-based Transforms
     Concatenate
     MapChunk
     Downsample
+    RearrangeElectrode
+    Flatten
 
 .. currentmodule:: torcheeg.transforms
 
@@ -107,4 +109,5 @@ Label Transforms
     Binary
     BinaryOneVSRest
     BinariesToCategory
-    StringToInt
\ No newline at end of file
+    StringToInt
+    Mapping
\ No newline at end of file
diff --git a/examples/examples_amigos_arjunvit.py b/examples/examples_amigos_arjunvit.py
deleted file mode 100644
index 0fcdd22..0000000
--- a/examples/examples_amigos_arjunvit.py
+++ /dev/null
@@ -1,164 +0,0 @@
-"""Arjun et al\'s ViT with the AMIGOS Dataset
-======================================
-In this case, we introduce how to use TorchEEG to train a vision transformer proposed by Arjun et al. on the AMIGOS dataset for emotion classification.
-"""
-
-import logging
-import os
-import random
-import time
-
-import numpy as np
-import torch
-from torch.utils.data.dataloader import DataLoader
-from torcheeg import transforms
-from torcheeg.datasets import AMIGOSDataset
-from torcheeg.datasets.constants.emotion_recognition.amigos import \
-    AMIGOS_CHANNEL_LOCATION_DICT
-from torcheeg.model_selection import KFoldGroupbyTrial
-from torcheeg.models import ArjunViT
-from torcheeg.trainers import ClassificationTrainer
-
-###############################################################################
-# Pre-experiment Preparation to Ensure Reproducibility
-# -----------------------------------------
-# Use the logging module to store output in a log file for easy reference while printing it to the screen.
-
-os.makedirs('./tmp_out/examples_amigos_arjunvit/log', exist_ok=True)
-logger = logging.getLogger('Arjun et al\'s ViT with the AMIGOS Dataset')
-logger.setLevel(logging.DEBUG)
-console_handler = logging.StreamHandler()
-timeticks = time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime())
-file_handler = logging.FileHandler(
-    os.path.join('./tmp_out/examples_amigos_arjunvit/log', f'{timeticks}.log'))
-logger.addHandler(console_handler)
-logger.addHandler(file_handler)
-
-###############################################################################
-# Set the random number seed in all modules to guarantee the same result when running again.
-
-
-def seed_everything(seed):
-    random.seed(seed)
-    np.random.seed(seed)
-    os.environ["PYTHONHASHSEED"] = str(seed)
-    torch.manual_seed(seed)
-    torch.cuda.manual_seed(seed)
-    torch.backends.cudnn.deterministic = True
-    torch.backends.cudnn.benchmark = False
-
-
-seed_everything(42)
-
-###############################################################################
-# Customize Trainer
-# -----------------------------------------
-# TorchEEG provides a large number of trainers to help complete the training of classification models, generative models and cross-domain methods. Here we choose the simplest classification trainer, inherit the trainer and overload the log function to save the log using our own defined method; other hook functions can also be overloaded to meet special needs.
-#
-
-
-class MyClassificationTrainer(ClassificationTrainer):
-    def log(self, *args, **kwargs):
-        if self.is_main:
-            logger.info(*args, **kwargs)
-
-
-######################################################################
-# Building Deep Learning Pipelines Using TorchEEG
-# -----------------------------------------
-# Step 1: Initialize the Dataset
-#
-# We use the AMIGOS dataset supported by TorchEEG. Here, we set an EEG sample to 1 second long and include 128 data points. The baseline signal is 5 seconds long, cut into five, and averaged as the baseline signal for the trial. In offline preprocessing, all EEG signals are debaselined and normalized, and the preprocessed EEG signals are stored in the local IO. In online processing, all EEG signals are converted into Tensors for input into neural networks.
-#
-
-dataset = AMIGOSDataset(io_path='./tmp_out/examples_amigos_arjunvit/amigos',
-                        root_path='./tmp_in/data_preprocessed',
-                        offline_transform=transforms.Compose([
-                            transforms.BaselineRemoval(),
-                            transforms.MeanStdNormalize()
-                        ]),
-                        online_transform=transforms.ToTensor(),
-                        label_transform=transforms.Compose([
-                            transforms.Select('valence'),
-                            transforms.Binary(5.0)
-                        ]),
-                        chunk_size=128,
-                        baseline_chunk_size=128,
-                        num_baseline=5,
-                        num_worker=4,
-                        io_mode='pickle')
-
-######################################################################
-# .. warning::
-#    If you use TorchEEG under the `Windows` system and want to use multiple processes (such as in dataset or dataloader), you should check whether :obj:`__name__` is :obj:`__main__` to avoid errors caused by multiple :obj:`import`.
-#
-# That is, under the :obj:`Windows` system, you need to:
-#  .. code-block::
-#
-#    if __name__ == "__main__":
-#        dataset = AMIGOSDataset(io_path='./tmp_out/examples_amigos_arjunvit/amigos',
-#                         root_path='./tmp_in/data_preprocessed',
-#                         offline_transform=transforms.Compose([
-#                             transforms.BaselineRemoval(),
-#                             transforms.MeanStdNormalize()
-#                         ]),
-#                         online_transform=transforms.ToTensor(),
-#                         label_transform=transforms.Compose([
-#                             transforms.Select('valence'),
-#                             transforms.Binary(5.0)
-#                         ]),
-#                         io_mode='pickle',
-#                         chunk_size=128,
-#                         baseline_chunk_size=128,
-#                         num_baseline=5,
-#                         num_worker=4)
-#        # the following codes
-#
-# .. note::
-#    LMDB may not be optimized for parts of Windows systems or storage devices. If you find that the data preprocessing speed is slow, you can consider setting :obj:`io_mode` to :obj:`pickle`, which is an alternative implemented by TorchEEG based on pickle.
-
-######################################################################
-# Step 2: Divide the Training and Test samples in the Dataset
-#
-# Here, the dataset is divided using 5-fold cross-validation. In the process of division, we group according to the trial index, and every trial takes 4 folds as training samples and 1 fold as test samples. Samples across trials are aggregated to obtain training set and test set.
-
-k_fold = KFoldGroupbyTrial(
-    n_splits=5, split_path='./tmp_out/examples_amigos_arjunvit/split')
-
-######################################################################
-# Step 3: Define the Model and Start Training
-#
-# We first use a loop to get the dataset in each cross-validation.
-# In each cross-validation, we initialize the ArjunViT model, and define the hyperparameters. For example, each EEG sample contains 128 time points, we divide it into 4 patches, each patch contains 32 time points, etc.
-#
-# We then initialize the trainer and set the hyperparameters in the trained model, such as the learning rate, the equipment used, etc. The :obj:`fit` method receives the training dataset and starts training the model. The :obj:`test` method receives a test dataset and reports the test results. The :obj:`save_state_dict` method can save the trained model.
-
-for i, (train_dataset, val_dataset) in enumerate(k_fold.split(dataset)):
-    # Initialize the model
-    model = ArjunViT(chunk_size=128,
-                     t_patch_size=32,
-                     hid_channels=32,
-                     num_electrodes=len(AMIGOS_CHANNEL_LOCATION_DICT),
-                     num_classes=2)
-
-    # Initialize the trainer and use the 0-th GPU for training, or set device_ids=[] to use CPU
-    trainer = MyClassificationTrainer(model=model,
-                                      lr=1e-4,
-                                      weight_decay=1e-4,
-                                      device_ids=[0])
-
-    # Initialize several batches of training samples and test samples
-    train_loader = DataLoader(train_dataset,
-                              batch_size=256,
-                              shuffle=True,
-                              num_workers=4)
-    val_loader = DataLoader(val_dataset,
-                            batch_size=256,
-                            shuffle=False,
-                            num_workers=4)
-
-    # Do 50 rounds of training
-    trainer.fit(train_loader, val_loader, num_epochs=50)
-    trainer.test(val_loader)
-    trainer.save_state_dict(
-        f'./tmp_out/examples_amigos_arjunvit/weight/{i}.pth')
diff --git a/examples/examples_amigos_dgcnn.py b/examples/examples_amigos_dgcnn.py
new file mode 100644
index 0000000..c2d8175
--- /dev/null
+++ b/examples/examples_amigos_dgcnn.py
@@ -0,0 +1,94 @@
+"""
+Train a DGCNN Model on the AMIGOS Dataset
+=============================
+
+In this tutorial, we'll walk through how to leverage TorchEEG for training a Dynamical Graph Convolutional Neural Network (DGCNN) to classify emotions using the AMIGOS dataset. Let's walk through the essential steps—from dataset preparation to model evaluation.
+
+"""
+
+######################################################################
+# Step 1: Initialize the Dataset
+# ----------------------------------------------
+#
+# First, let's prepare the AMIGOS dataset using TorchEEG's built-in functionality. Each sample in the dataset is 1 second long, containing 128 EEG data points. We will use several preprocessing steps to prepare the data for our neural network model.
+# Here's what each preprocessing step does:
+# - Divide each EEG signal into 4 sub-bands
+# - Calculate differential entropy for each sub-band as a feature
+# - Remove the baseline from the signal
+# - Convert the preprocessed signals into Tensors
+# The code below demonstrates these steps:
+# 
+
+from torcheeg.datasets import AMIGOSDataset
+from torcheeg import transforms
+
+dataset = AMIGOSDataset(io_path=f'./examples_amigos_dgcnn/amigos',
+                        root_path='./data_preprocessed',
+                        offline_transform=transforms.BandDifferentialEntropy(
+                            sampling_rate=128, apply_to_baseline=True),
+                        online_transform=transforms.Compose([
+                            transforms.BaselineRemoval(),
+                            transforms.ToTensor()
+                        ]),
+                        label_transform=transforms.Compose([
+                            transforms.Select('valence'),
+                            transforms.Binary(5.0)
+                        ]),
+                        chunk_size=128,
+                        baseline_chunk_size=128,
+                        num_baseline=5,
+                        num_worker=4)
+
+######################################################################
+# Step 2: Divide the Training and Test samples in the Dataset
+# ----------------------------------------------
+#
+# Now, let's partition the dataset into training and test sets. We will employ 5-fold cross-validation for this purpose. During this phase, samples are grouped based on their trial index. Four folds are designated for training, and one fold serves as the test set. Subsequently, we combine these samples across all trials to finalize the training and test sets.
+#
+
+from torcheeg.model_selection import KFoldGroupbyTrial
+
+k_fold = KFoldGroupbyTrial(n_splits=5,
+                           split_path=f'./examples_amigos_dgcnn/split')
+
+######################################################################
+# Step 3: Define the Model and Initiate Training
+# ----------------------------------------------
+#
+# With our dataset prepared and divided, the next step is to build and train the DGCNN model. In this block of code, we define our DGCNN model's architecture, setting important hyperparameters like the number of channels, electrodes, and hidden layers.
+# After that, we'll train the model for 50 epochs, utilizing TorchEEG's `ClassifierTrainer`.
+#
+
+from torch.utils.data import DataLoader
+from torcheeg.models import DGCNN
+
+from torcheeg.trainers import ClassifierTrainer
+
+import pytorch_lightning as pl
+
+for i, (train_dataset, val_dataset) in enumerate(k_fold.split(dataset)):
+    train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
+    val_loader = DataLoader(val_dataset, batch_size=64, shuffle=False)
+
+    model = DGCNN(in_channels=4,
+                  num_electrodes=14,
+                  hid_channels=64,
+                  num_classes=2)
+
+    trainer = ClassifierTrainer(model=model,
+                                num_classes=2,
+                                lr=1e-4,
+                                weight_decay=1e-4,
+                                accelerator="gpu")
+    trainer.fit(train_loader,
+                val_loader,
+                max_epochs=50,
+                default_root_dir=f'./examples_amigos_dgcnn/model/{i}',
+                callbacks=[pl.callbacks.ModelCheckpoint(save_last=True)],
+                enable_progress_bar=True,
+                enable_model_summary=True,
+                limit_val_batches=0.0)
+    score = trainer.test(val_loader,
+                         enable_progress_bar=True,
+                         enable_model_summary=True)[0]
+    print(f'Fold {i} test accuracy: {score["test_accuracy"]:.4f}')
\ No newline at end of file
diff --git a/examples/examples_datasets.py b/examples/examples_datasets.py
new file mode 100644
index 0000000..0293894
--- /dev/null
+++ b/examples/examples_datasets.py
@@ -0,0 +1,196 @@
+"""
+Introduction to the datasets Module
+===================================
+
+Welcome to the guide on TorchEEG's ``datasets`` Module! This module provides you with various benchmark datasets for EEG-based emotion recognition, such as DEAP, DREAMER, SEED, MAHNOB, AMIGOS, and MPED. These datasets use a range of stimuli like music and videos to trigger emotional responses. Once the emotional experiences are recorded, participants label them using methods like valence-arousal dimensions or discrete emotion categories.
+
+This guide will help you understand how to easily load and manipulate these datasets for your training needs, including signal segmentation, applying transformations, and more.
+
+"""
+
+######################################################################
+# An Example of Emotion Recognition Dataset
+# ----------------------------------------------
+#
+# To begin, you'll need to download the DEAP dataset from DEAP download link. Once you have the dataset, use the following code to load it. This will read EEG signals and labels, apply offline transformations, and save them for easy access later on.
+#
+
+from torcheeg.datasets import DEAPDataset
+from torcheeg import transforms
+
+from torcheeg.datasets.constants.emotion_recognition.deap import \
+    DEAP_CHANNEL_LOCATION_DICT
+
+dataset = DEAPDataset(
+    io_path=f'./deap',
+    root_path='./data_preprocessed_python',
+    offline_transform=transforms.Compose([
+        transforms.BandDifferentialEntropy(apply_to_baseline=True),
+        transforms.ToGrid(DEAP_CHANNEL_LOCATION_DICT, apply_to_baseline=True)
+    ]),
+    online_transform=transforms.Compose(
+        [transforms.BaselineRemoval(),
+         transforms.ToTensor()]),
+    label_transform=transforms.Compose([
+        transforms.Select('valence'),
+        transforms.Binary(5.0),
+    ]),
+    num_worker=8)
+
+######################################################################
+#
+# You can get individual samples from the dataset like so:
+#
+
+print(dataset[0])
+
+######################################################################
+#
+# This will output a tuple, where the first element is the EEG signal and the second is its label.
+#
+
+######################################################################
+#
+# If you want to visualize the EEG signals, you can use the helper function as shown below:
+#
+
+import torch
+from torcheeg.utils import plot_2d_tensor
+
+img = plot_2d_tensor(torch.tensor(dataset[0][0]))
+
+######################################################################
+#
+# You can also refer to the document for more information on visualizing EEG signals
+# via the ``plot_3d_tensor``, ``plot_feature_topomap``, ``plot_raw_topomap``, ``plot_signal``, and ``plot_adj_connectivity`` function.
+#
+
+######################################################################
+# Important Notes
+# ----------------------------------------------
+#
+# If you make changes to the dataset's configuration, remember to clear the cache or specify a new io_path. Otherwise, only online transformations (online_transform and label_transform) will take effect:
+#
+# .. code-block:: bash
+#
+#   !rm -rf ./deap
+#
+
+######################################################################
+# If you remove all transforms and declare a dataset, you’ll find that it
+# returns segmented EEG signals that haven’t undergone any preprocessing.
+# Additionally, there is a dictionary representing all the
+# meta-information corresponding to the EEG signal samples.
+#
+
+from torcheeg.datasets import DEAPDataset
+
+dataset = DEAPDataset(io_path=f'./deap', root_path='./data_preprocessed_python')
+print(dataset[0])
+
+######################################################################
+# Advanced Usage
+# ----------------------------------------------
+#
+# Most emotion analysis datasets can be found in TorchEEG. Moreover,
+# TorchEEG provides support for ``MoABB``, allowing access to
+# motor-imagery-related datasets with the help of ``MoABB``.
+#
+
+import torcheeg.datasets.moabb as moabb_dataset
+
+from moabb.datasets import BNCI2014001
+from moabb.paradigms import LeftRightImagery
+
+dataset = BNCI2014001()
+dataset.subject_list = [1, 2, 3]
+paradigm = LeftRightImagery()
+dataset = moabb_dataset.MOABBDataset(
+    dataset=dataset,
+    paradigm=paradigm,
+    io_path='./moabb',
+    offline_transform=transforms.Compose([transforms.BandDifferentialEntropy()
+                                          ]),
+    online_transform=transforms.ToTensor(),
+    label_transform=transforms.Compose([transforms.Select('label')]))
+
+######################################################################
+# Custom Datasets
+# ----------------------------------------------
+#
+# TorchEEG also supports custom datasets. You can place recorded EEG
+# signal files in a folder following a specific rule, such as:
+#
+# ::
+#
+#    label01
+#    |- sub01.edf
+#    |- sub02.edf
+#    label02
+#    |- sub01.edf
+#    |- sub02.edf
+#
+# Then, you can use ``FolderDataset`` to automatically access the
+# corresponding EEG signal samples.
+#
+
+from torcheeg.datasets import FolderDataset
+
+label_map = {'label01': 0, 'label02': 1}
+dataset = FolderDataset(io_path='./folder',
+                        root_path='./root_folder',
+                        structure='subject_in_label',
+                        num_channel=14,
+                        online_transform=transforms.ToTensor(),
+                        label_transform=transforms.Compose([
+                            transforms.Select('label'),
+                            transforms.Lambda(lambda x: label_map[x])
+                        ]),
+                        num_worker=4)
+
+######################################################################
+# Alternatively, you can use a CSV file to specify more detailed
+# meta-information for reading:
+#
+# ::
+#
+#    | subject_id | trial_id | label | file_path                 |
+#    | ---------- | -------  | ----- | ------------------------- |
+#    | sub1       | 0        | 0     | './data/label1/sub1.fif' |
+#    | sub1       | 1        | 1     | './data/label2/sub1.fif' |
+#    | sub1       | 2        | 2     | './data/label3/sub1.fif' |
+#    | sub2       | 0        | 0     | './data/label1/sub2.fif' |
+#    | sub2       | 1        | 1     | './data/label2/sub2.fif' |
+#    | sub2       | 2        | 2     | './data/label3/sub2.fif' |
+#
+
+from torcheeg.datasets import CSVFolderDataset
+
+dataset = CSVFolderDataset(csv_path='./data.csv',
+                           online_transform=transforms.ToTensor(),
+                           label_transform=transforms.Select('label'),
+                           num_worker=4)
+
+######################################################################
+# By default, TorchEEG uses mne to read recorded EEG signals, but you can
+# also specify your own file reading logic through ``read_fn``.
+#
+
+import mne
+
+
+def default_read_fn(file_path, **kwargs):
+    # Load EEG file
+    raw = mne.io.read_raw(file_path)
+    # Convert raw to epochs
+    epochs = mne.make_fixed_length_epochs(raw, duration=1)
+    # Return EEG data
+    return epochs
+
+
+dataset = CSVFolderDataset(io_path='./csv_folder',
+                           csv_path='./data.csv',
+                           read_fn=default_read_fn,
+                           online_transform=transforms.ToTensor(),
+                           label_transform=transforms.Select('label'),
+                           num_worker=4)
\ No newline at end of file
diff --git a/examples/examples_deap_ccnn.py b/examples/examples_deap_ccnn.py
deleted file mode 100644
index 3419a8a..0000000
--- a/examples/examples_deap_ccnn.py
+++ /dev/null
@@ -1,159 +0,0 @@
-"""CCNN with the DEAP Dataset
-======================================
-In this case, we introduce how to use TorchEEG to train a Continuous Convolutional Neural Network (CCNN) on the DEAP dataset for emotion classification.
-"""
-
-import logging
-import os
-import random
-import time
-
-import numpy as np
-import torch
-from torch.utils.data.dataloader import DataLoader
-from torcheeg import transforms
-from torcheeg.datasets import DEAPDataset
-from torcheeg.datasets.constants.emotion_recognition.deap import \
-    DEAP_CHANNEL_LOCATION_DICT
-from torcheeg.model_selection import KFoldGroupbyTrial
-from torcheeg.models import CCNN
-from torcheeg.trainers import ClassifierTrainer
-
-###############################################################################
-# Pre-experiment Preparation to Ensure Reproducibility
-# -----------------------------------------
-# Use the logging module to store output in a log file for easy reference while printing it to the screen.
-
-os.makedirs('./tmp_out/examples_deap_ccnn/log', exist_ok=True)
-logger = logging.getLogger('CCNN with the DEAP Dataset')
-logger.setLevel(logging.DEBUG)
-console_handler = logging.StreamHandler()
-timeticks = time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime())
-file_handler = logging.FileHandler(
-    os.path.join('./tmp_out/examples_deap_ccnn/log', f'{timeticks}.log'))
-logger.addHandler(console_handler)
-logger.addHandler(file_handler)
-
-###############################################################################
-# Set the random number seed in all modules to guarantee the same result when running again.
-
-
-def seed_everything(seed):
-    random.seed(seed)
-    np.random.seed(seed)
-    os.environ["PYTHONHASHSEED"] = str(seed)
-    torch.manual_seed(seed)
-    torch.cuda.manual_seed(seed)
-    torch.backends.cudnn.deterministic = True
-    torch.backends.cudnn.benchmark = False
-
-
-seed_everything(42)
-
-###############################################################################
-# Customize Trainer
-# -----------------------------------------
-# TorchEEG provides a large number of trainers to help complete the training of classification models, generative models and cross-domain methods. Here we choose the simplest classification trainer, inherit the trainer and overload the log function to save the log using our own defined method; other hook functions can also be overloaded to meet special needs.
-#
-
-######################################################################
-# Building Deep Learning Pipelines Using TorchEEG
-# -----------------------------------------
-# Step 1: Initialize the Dataset
-#
-# We use the DEAP dataset supported by TorchEEG. Here, we set an EEG sample to 1 second long and include 128 data points. The baseline signal is 3 seconds long, cut into three, and averaged as the baseline signal for the trial. In offline preprocessing, we divide the EEG signal of every electrode into 4 sub-bands, and calculate the differential entropy on each sub-band as a feature, followed by debaselining and mapping on the grid. Finally, the preprocessed EEG signals are stored in the local IO. In online processing, all EEG signals are converted into Tensors for input into neural networks.
-#
-
-dataset = DEAPDataset(io_path=f'./tmp_out/examples_deap_ccnn/deap',
-                      root_path='./tmp_in/data_preprocessed_python',
-                      offline_transform=transforms.Compose([
-                          transforms.BandDifferentialEntropy(
-                              sampling_rate=128, apply_to_baseline=True),
-                          transforms.BaselineRemoval(),
-                          transforms.ToGrid(DEAP_CHANNEL_LOCATION_DICT)
-                      ]),
-                      online_transform=transforms.ToTensor(),
-                      label_transform=transforms.Compose([
-                          transforms.Select('valence'),
-                          transforms.Binary(5.0),
-                      ]),
-                      chunk_size=128,
-                      baseline_chunk_size=128,
-                      num_baseline=3,
-                      num_worker=4)
-
-######################################################################
-# .. warning::
-#    If you use TorchEEG under the `Windows` system and want to use multiple processes (such as in dataset or dataloader), you should check whether :obj:`__name__` is :obj:`__main__` to avoid errors caused by multiple :obj:`import`.
-#
-# That is, under the :obj:`Windows` system, you need to:
-#  .. code-block::
-#
-#    if __name__ == "__main__":
-#        dataset = DEAPDataset(
-#             io_path=
-#             f'./tmp_out/examples_deap_ccnn/deap',
-#             root_path='./tmp_in/data_preprocessed_python',
-#             offline_transform=transforms.Compose([
-#                 transforms.BandDifferentialEntropy(sampling_rate=128,
-#                                                    apply_to_baseline=True),
-#                 transforms.BaselineRemoval(),
-#                 transforms.ToGrid(DEAP_CHANNEL_LOCATION_DICT)
-#             ]),
-#             online_transform=transforms.ToTensor(),
-#             label_transform=transforms.Compose([
-#                 transforms.Select('valence'),
-#                 transforms.Binary(5.0),
-#             ]),
-#             io_mode='pickle',
-#             chunk_size=128,
-#             baseline_chunk_size=128,
-#             num_baseline=3,
-#             num_worker=4)
-#        # the following codes
-#
-# .. note::
-#    LMDB may not be optimized for parts of Windows systems or storage devices. If you find that the data preprocessing speed is slow, you can consider setting :obj:`io_mode` to :obj:`pickle`, which is an alternative implemented by TorchEEG based on pickle.
-
-######################################################################
-# Step 2: Divide the Training and Test samples in the Dataset
-#
-# Here, the dataset is divided using 5-fold cross-validation. In the process of division, we group according to the trial index, and every trial takes 4 folds as training samples and 1 fold as test samples. Samples across trials are aggregated to obtain training set and test set.
-#
-
-k_fold = KFoldGroupbyTrial(n_splits=5,
-                           split_path='./tmp_out/examples_deap_ccnn/split')
-
-######################################################################
-# Step 3: Define the Model and Start Training
-#
-# We first use a loop to get the dataset in each cross-validation. In each cross-validation, we initialize the CCNN model and define the hyperparameters. For example, each EEG sample contains 4-channel features from 4 sub-bands, the grid size is 9 times 9, etc.
-#
-# We then initialize the trainer and set the hyperparameters in the trained model, such as the learning rate, the equipment used, etc. The :obj:`fit` method receives the training dataset and starts training the model. The :obj:`test` method receives a test dataset and reports the test results. The :obj:`save_state_dict` method can save the trained model.
-
-for i, (train_dataset, val_dataset) in enumerate(k_fold.split(dataset)):
-    # Initialize the model
-    model = CCNN(in_channels=4, grid_size=(9, 9), num_classes=2, dropout=0.5)
-
-    # Initialize the trainer and use the 0-th GPU for training, or set device_ids=[] to use CPU
-    trainer = ClassifierTrainer(model=model,
-                                lr=1e-4,
-                                devices=1,
-                                accelerator='gpu',
-                                num_classes=2,
-                                weight_decay=1e-4)
-
-    # Initialize several batches of training samples and test samples
-    train_loader = DataLoader(train_dataset,
-                              batch_size=256,
-                              shuffle=True,
-                              num_workers=4)
-    val_loader = DataLoader(val_dataset,
-                            batch_size=256,
-                            shuffle=False,
-                            num_workers=4)
-
-    # Do 50 rounds of training
-    trainer.fit(train_loader, val_loader, max_epochs=50)
-    trainer.test(val_loader)
-    trainer.save_state_dict(f'./tmp_out/examples_deap_ccnn/weight/{i}.pth')
diff --git a/examples/examples_deap_contrastive.py b/examples/examples_deap_contrastive.py
deleted file mode 100644
index 5ed2404..0000000
--- a/examples/examples_deap_contrastive.py
+++ /dev/null
@@ -1,226 +0,0 @@
-"""CCNN with the DEAP Dataset
-======================================
-In this case, we introduce how to use TorchEEG to train a Continuous Convolutional Neural Network (CCNN) on the DEAP dataset for emotion classification.
-"""
-
-import logging
-import os
-import random
-import time
-
-import numpy as np
-import torch
-from torch.utils.data.dataloader import DataLoader
-from torcheeg import transforms
-from torcheeg.datasets import DEAPDataset
-from torcheeg.datasets.constants.emotion_recognition.deap import \
-    DEAP_CHANNEL_LOCATION_DICT
-from torcheeg.model_selection import KFoldGroupbyTrial
-from torcheeg.models import CCNN
-from torcheeg.trainers import ClassifierTrainer, ContrastiveTrainer
-
-###############################################################################
-# Pre-experiment Preparation to Ensure Reproducibility
-# -----------------------------------------
-# Use the logging module to store output in a log file for easy reference while printing it to the screen.
-
-os.makedirs('./tmp_out/examples_deap_ccnn/log', exist_ok=True)
-logger = logging.getLogger('CCNN with the DEAP Dataset')
-logger.setLevel(logging.DEBUG)
-console_handler = logging.StreamHandler()
-timeticks = time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime())
-file_handler = logging.FileHandler(
-    os.path.join('./tmp_out/examples_deap_ccnn/log', f'{timeticks}.log'))
-logger.addHandler(console_handler)
-logger.addHandler(file_handler)
-
-###############################################################################
-# Set the random number seed in all modules to guarantee the same result when running again.
-
-
-def seed_everything(seed):
-    random.seed(seed)
-    np.random.seed(seed)
-    os.environ["PYTHONHASHSEED"] = str(seed)
-    torch.manual_seed(seed)
-    torch.cuda.manual_seed(seed)
-    torch.backends.cudnn.deterministic = True
-    torch.backends.cudnn.benchmark = False
-
-
-seed_everything(42)
-
-###############################################################################
-# Customize Trainer
-# -----------------------------------------
-# TorchEEG provides a large number of trainers to help complete the training of classification models, generative models and cross-domain methods. Here we choose the simplest classification trainer, inherit the trainer and overload the log function to save the log using our own defined method; other hook functions can also be overloaded to meet special needs.
-#
-
-######################################################################
-# Building Deep Learning Pipelines Using TorchEEG
-# -----------------------------------------
-# Step 1: Initialize the Dataset
-#
-# We use the DEAP dataset supported by TorchEEG. Here, we set an EEG sample to 1 second long and include 128 data points. The baseline signal is 3 seconds long, cut into three, and averaged as the baseline signal for the trial. In offline preprocessing, we divide the EEG signal of every electrode into 4 sub-bands, and calculate the differential entropy on each sub-band as a feature, followed by debaselining and mapping on the grid. Finally, the preprocessed EEG signals are stored in the local IO. In online processing, all EEG signals are converted into Tensors for input into neural networks.
-#
-
-dataset = DEAPDataset(io_path=f'./tmp_out/examples_deap_ccnn/deap',
-                      root_path='./tmp_in/data_preprocessed_python',
-                      offline_transform=transforms.Compose([
-                          transforms.BandDifferentialEntropy(
-                              sampling_rate=128, apply_to_baseline=True),
-                          transforms.BaselineRemoval(),
-                          transforms.ToGrid(DEAP_CHANNEL_LOCATION_DICT)
-                      ]),
-                      online_transform=transforms.ToTensor(),
-                      label_transform=transforms.Compose([
-                          transforms.Select('valence'),
-                          transforms.Binary(5.0),
-                      ]),
-                      chunk_size=128,
-                      baseline_chunk_size=128,
-                      num_baseline=3,
-                      num_worker=4)
-
-contras_dataset = DEAPDataset(
-    io_path=f'./tmp_out/examples_deap_ccnn/deap',
-    root_path='./tmp_in/data_preprocessed_python',
-    offline_transform=transforms.Compose([
-        transforms.BandDifferentialEntropy(sampling_rate=128,
-                                           apply_to_baseline=True),
-        transforms.BaselineRemoval(),
-        transforms.ToGrid(DEAP_CHANNEL_LOCATION_DICT)
-    ]),
-    online_transform=transforms.Compose([
-        transforms.ToTensor(),
-        transforms.Contrastive(transforms.Compose(
-            [transforms.RandomMask(p=0.5),
-             transforms.RandomNoise(p=0.5)]),
-                               num_views=2)
-    ]),
-    label_transform=transforms.Compose([
-        transforms.Select('valence'),
-        transforms.Binary(5.0),
-    ]),
-    chunk_size=128,
-    baseline_chunk_size=128,
-    num_baseline=3,
-    num_worker=4)
-
-######################################################################
-# .. warning::
-#    If you use TorchEEG under the `Windows` system and want to use multiple processes (such as in dataset or dataloader), you should check whether :obj:`__name__` is :obj:`__main__` to avoid errors caused by multiple :obj:`import`.
-#
-# That is, under the :obj:`Windows` system, you need to:
-#  .. code-block::
-#
-#    if __name__ == "__main__":
-#        dataset = DEAPDataset(
-#             io_path=
-#             f'./tmp_out/examples_deap_ccnn/deap',
-#             root_path='./tmp_in/data_preprocessed_python',
-#             offline_transform=transforms.Compose([
-#                 transforms.BandDifferentialEntropy(sampling_rate=128,
-#                                                    apply_to_baseline=True),
-#                 transforms.BaselineRemoval(),
-#                 transforms.ToGrid(DEAP_CHANNEL_LOCATION_DICT)
-#             ]),
-#             online_transform=transforms.ToTensor(),
-#             label_transform=transforms.Compose([
-#                 transforms.Select('valence'),
-#                 transforms.Binary(5.0),
-#             ]),
-#             io_mode='pickle',
-#             chunk_size=128,
-#             baseline_chunk_size=128,
-#             num_baseline=3,
-#             num_worker=4)
-#        # the following codes
-#
-# .. note::
-#    LMDB may not be optimized for parts of Windows systems or storage devices. If you find that the data preprocessing speed is slow, you can consider setting :obj:`io_mode` to :obj:`pickle`, which is an alternative implemented by TorchEEG based on pickle.
-
-######################################################################
-# Step 2: Divide the Training and Test samples in the Dataset
-#
-# Here, the dataset is divided using 5-fold cross-validation. In the process of division, we group according to the trial index, and every trial takes 4 folds as training samples and 1 fold as test samples. Samples across trials are aggregated to obtain training set and test set.
-#
-
-k_fold = KFoldGroupbyTrial(n_splits=5,
-                           split_path='./tmp_out/examples_deap_ccnn/split')
-
-######################################################################
-# Step 3: Define the Model and Start Training
-#
-# We first use a loop to get the dataset in each cross-validation. In each cross-validation, we initialize the CCNN model and define the hyperparameters. For example, each EEG sample contains 4-channel features from 4 sub-bands, the grid size is 9 times 9, etc.
-#
-# We then initialize the trainer and set the hyperparameters in the trained model, such as the learning rate, the equipment used, etc. The :obj:`fit` method receives the training dataset and starts training the model. The :obj:`test` method receives a test dataset and reports the test results. The :obj:`save_state_dict` method can save the trained model.
-
-
-class Extractor(CCNN):
-
-    def forward(self, x):
-        x = self.conv1(x)
-        x = self.conv2(x)
-        x = self.conv3(x)
-        x = self.conv4(x)
-
-        x = x.flatten(start_dim=1)
-        return x
-
-test_accuracy_list = []
-for i, ((train_dataset, val_dataset), (contras_train_dataset,
-                                      contras_val_dataset)) in enumerate(
-                                          zip(k_fold.split(dataset),
-                                              k_fold.split(contras_dataset))):
-    # Initialize the model
-    model = Extractor(in_channels=4,
-                      grid_size=(9, 9),
-                      num_classes=2,
-                      dropout=0.5)
-
-    # Initialize the trainer and use the 0-th GPU for training, or set device_ids=[] to use CPU
-
-    train_loader = DataLoader(contras_train_dataset,
-                              batch_size=128,
-                              shuffle=True,
-                              num_workers=4)
-
-    val_loader = DataLoader(contras_val_dataset,
-                            batch_size=128,
-                            shuffle=True,
-                            num_workers=4)
-
-    trainer = ContrastiveTrainer(extractor=model,
-                                 lr=1e-4,
-                                 devices=1,
-                                 accelerator='gpu')
-    # Do 50 rounds of training
-    trainer.fit(train_loader, val_loader, max_epochs=50)
-
-    model = CCNN(in_channels=4, grid_size=(9, 9), num_classes=2, dropout=0.5)
-    # model.load_state_dict(trainer.extractor.state_dict())
-
-    # Initialize several batches of training samples and test samples
-    train_loader = DataLoader(train_dataset,
-                              batch_size=256,
-                              shuffle=True,
-                              num_workers=4)
-
-    val_loader = DataLoader(val_dataset,
-                            batch_size=256,
-                            shuffle=False,
-                            num_workers=4)
-
-    trainer = ClassifierTrainer(model=model,
-                                lr=1e-4,
-                                devices=1,
-                                accelerator='gpu',
-                                num_classes=2,
-                                weight_decay=1e-4)
-    # Do 50 rounds of training
-    trainer.fit(train_loader, val_loader, max_epochs=50)
-    test_results = trainer.test(val_loader)
-    test_accuracy_list.append(test_results[0]['test_accuracy'])
-
-print(f'Average test accuracy: {np.mean(test_accuracy_list)}')
diff --git a/examples/examples_deap_cwgangp.py b/examples/examples_deap_cwgangp.py
new file mode 100644
index 0000000..4720fe3
--- /dev/null
+++ b/examples/examples_deap_cwgangp.py
@@ -0,0 +1,140 @@
+"""
+Train a C-WGAN-GP Model on the DEAP Dataset
+===================================================
+
+In this tutorial, we'll walk through how to train a Conditional Wasserstein Generative Adversarial Network with Gradient Penalty (C-WGAN-GP) on the DEAP EEG dataset. This will allow us to generate new EEG signals. We'll also evaluate the generated EEG signals using Fréchet Inception Distance (FID) and Kernel Inception Distance (KID).
+
+"""
+
+
+######################################################################
+# Step 1: Initialize the Dataset
+# ------------------------------
+# 
+# We use the DEAP dataset supported by TorchEEG. Each EEG sample is set to
+# be 1 second long, encompassing 128 data points. The baseline signal is 3
+# seconds long, which we divide into three sections and then average to
+# obtain the trial’s baseline signal.
+# 
+# During offline preprocessing, we divide each electrode’s EEG signal into
+# 4 sub-bands, calculate the differential entropy for each sub-band as a
+# feature, perform debaselining, and map onto a grid. Finally, the
+# preprocessed EEG signals are saved locally. For online processing, we
+# convert all EEG signals into Tensors, making them suitable for neural
+# network input.
+# 
+
+from torcheeg.datasets import DEAPDataset
+from torcheeg import transforms
+
+from torcheeg.datasets.constants.emotion_recognition.deap import \
+    DEAP_CHANNEL_LOCATION_DICT
+
+dataset = DEAPDataset(
+    io_path=f'./examples_deap_cwgangp/deap',
+    root_path='./data_preprocessed_python',
+    offline_transform=transforms.Compose([
+        transforms.BandDifferentialEntropy(apply_to_baseline=True),
+        transforms.ToGrid(DEAP_CHANNEL_LOCATION_DICT, apply_to_baseline=True)
+    ]),
+    online_transform=transforms.Compose(
+        [transforms.BaselineRemoval(),
+         transforms.ToTensor()]),
+    label_transform=transforms.Compose([
+        transforms.Select('valence'),
+        transforms.Binary(5.0),
+    ]),
+    num_worker=8)
+
+
+######################################################################
+# Step 2: Split the Dataset into Training and Test Sets
+# -----------------------------------------------------
+# 
+# In this case, we use per-subject 5-fold cross-validation to split the
+# dataset. During this process, we separate each subject’s EEG samples
+# into training and test sets. We use 4 folds for training and 1 fold for
+# testing.
+# 
+
+from torcheeg.datasets import KFoldGroupbyTrial
+
+k_fold = KFoldGroupbyTrial(n_splits=5,
+                           split_path='./examples_deap_cwgangp/split',
+                           shuffle=True,
+                           random_state=42)
+
+
+######################################################################
+# Step 3: Define the Model
+# ----------------------------------------------
+#
+# We'll employ the CCNN model's feature extraction part to evaluate our CWGAN-GP model. Specifically, we'll compare the features from generated and real EEG signals.
+#
+
+from torcheeg.models import CCNN
+
+
+class Extractor(CCNN):
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.conv2(x)
+        x = self.conv3(x)
+        x = self.conv4(x)
+        x = x.flatten(start_dim=1)
+        return x
+
+######################################################################
+# Step 3: Define the Model and Initiate Training
+# ----------------------------------------------
+# 
+# Now we're ready to train the classifier and feature extractor, followed by the CWGAN-GP model.
+# 
+
+from torch.utils.data import DataLoader
+from torcheeg.trainers import ClassifierTrainer, CWGANGPTrainer
+from torcheeg.models import CCNN, BCGenerator, BCDiscriminator
+
+import pytorch_lightning as pl
+
+for i, (train_dataset, val_dataset) in enumerate(k_fold.split(dataset)):
+    train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
+    val_loader = DataLoader(val_dataset, batch_size=64, shuffle=False)
+
+    classifier = CCNN(num_classes=2, in_channels=4, grid_size=(9, 9))
+
+    # if you have a pre-trained classifier, you can just load it, instead of training it from scratch
+    trainer = ClassifierTrainer(model=classifier,
+                                num_classes=2,
+                                lr=1e-4,
+                                weight_decay=1e-4,
+                                accelerator="gpu")
+    trainer.fit(train_loader,
+                val_loader,
+                max_epochs=50,
+                default_root_dir=f'./examples_deap_cwgangp/model/{i}',
+                callbacks=[pl.callbacks.ModelCheckpoint(save_last=True)],
+                enable_progress_bar=True,
+                enable_model_summary=True,
+                limit_val_batches=0.0)
+    
+    extractor = Extractor(num_classes=2, in_channels=4, grid_size=(9, 9))
+    extractor.load_state_dict(classifier.state_dict())
+
+    g_model = BCGenerator(in_channels=128)
+    d_model = BCDiscriminator(in_channels=4)
+
+    trainer = CWGANGPTrainer(g_model,
+                              d_model,
+                              metric_extractor=extractor,
+                              metric_classifier=classifier,
+                              metric_num_features=9 * 9 * 64,
+                              metrics=['fid', 'is'],
+                             accelerator='gpu')
+    trainer.fit(train_loader, val_loader, max_epochs=1)
+    trainer.test(val_loader)
+
+######################################################################
+# 
+# That's it! You've successfully trained a C-WGAN-GP model on the DEAP EEG dataset and evaluated it using FID and KID metrics.
+#
\ No newline at end of file
diff --git a/examples/examples_deap_domain_adaption.py b/examples/examples_deap_domain_adaption.py
deleted file mode 100644
index 19d1993..0000000
--- a/examples/examples_deap_domain_adaption.py
+++ /dev/null
@@ -1,187 +0,0 @@
-"""CCNN with the DEAP Dataset
-======================================
-In this case, we introduce how to use TorchEEG to train a Continuous Convolutional Neural Network (CCNN) on the DEAP dataset for emotion classification.
-"""
-
-import logging
-import os
-import random
-import time
-
-import numpy as np
-import torch
-from torch.utils.data.dataloader import DataLoader
-from torcheeg import transforms
-from torcheeg.datasets import DEAPDataset
-from torcheeg.datasets.constants.emotion_recognition.deap import \
-    DEAP_CHANNEL_LOCATION_DICT
-from torcheeg.model_selection import LeaveOneSubjectOut
-from torcheeg.models import CCNN
-from torcheeg.trainers import ADATrainer
-
-###############################################################################
-# Pre-experiment Preparation to Ensure Reproducibility
-# -----------------------------------------
-# Use the logging module to store output in a log file for easy reference while printing it to the screen.
-
-os.makedirs('./tmp_out/examples_deap_domain_adaption/log', exist_ok=True)
-logger = logging.getLogger('CCNN with the DEAP Dataset')
-logger.setLevel(logging.DEBUG)
-console_handler = logging.StreamHandler()
-timeticks = time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime())
-file_handler = logging.FileHandler(
-    os.path.join('./tmp_out/examples_deap_domain_adaption/log', f'{timeticks}.log'))
-logger.addHandler(console_handler)
-logger.addHandler(file_handler)
-
-###############################################################################
-# Set the random number seed in all modules to guarantee the same result when running again.
-
-
-def seed_everything(seed):
-    random.seed(seed)
-    np.random.seed(seed)
-    os.environ["PYTHONHASHSEED"] = str(seed)
-    torch.manual_seed(seed)
-    torch.cuda.manual_seed(seed)
-    torch.backends.cudnn.deterministic = True
-    torch.backends.cudnn.benchmark = False
-
-
-seed_everything(42)
-
-###############################################################################
-# Customize Trainer
-# -----------------------------------------
-# TorchEEG provides a large number of trainers to help complete the training of classification models, generative models and cross-domain methods. Here we choose the simplest classification trainer, inherit the trainer and overload the log function to save the log using our own defined method; other hook functions can also be overloaded to meet special needs.
-#
-
-######################################################################
-# Building Deep Learning Pipelines Using TorchEEG
-# -----------------------------------------
-# Step 1: Initialize the Dataset
-#
-# We use the DEAP dataset supported by TorchEEG. Here, we set an EEG sample to 1 second long and include 128 data points. The baseline signal is 3 seconds long, cut into three, and averaged as the baseline signal for the trial. In offline preprocessing, we divide the EEG signal of every electrode into 4 sub-bands, and calculate the differential entropy on each sub-band as a feature, followed by debaselining and mapping on the grid. Finally, the preprocessed EEG signals are stored in the local IO. In online processing, all EEG signals are converted into Tensors for input into neural networks.
-#
-
-dataset = DEAPDataset(io_path=f'./tmp_out/examples_deap_domain_adaption/deap',
-                      root_path='./tmp_in/data_preprocessed_python',
-                      offline_transform=transforms.Compose([
-                          transforms.BandDifferentialEntropy(
-                              sampling_rate=128, apply_to_baseline=True),
-                          transforms.BaselineRemoval(),
-                          transforms.ToGrid(DEAP_CHANNEL_LOCATION_DICT)
-                      ]),
-                      online_transform=transforms.ToTensor(),
-                      label_transform=transforms.Compose([
-                          transforms.Select('valence'),
-                          transforms.Binary(5.0),
-                      ]),
-                      chunk_size=128,
-                      baseline_chunk_size=128,
-                      num_baseline=3,
-                      num_worker=4)
-
-######################################################################
-# .. warning::
-#    If you use TorchEEG under the `Windows` system and want to use multiple processes (such as in dataset or dataloader), you should check whether :obj:`__name__` is :obj:`__main__` to avoid errors caused by multiple :obj:`import`.
-#
-# That is, under the :obj:`Windows` system, you need to:
-#  .. code-block::
-#
-#    if __name__ == "__main__":
-#        dataset = DEAPDataset(
-#             io_path=
-#             f'./tmp_out/examples_deap_domain_adaption/deap',
-#             root_path='./tmp_in/data_preprocessed_python',
-#             offline_transform=transforms.Compose([
-#                 transforms.BandDifferentialEntropy(sampling_rate=128,
-#                                                    apply_to_baseline=True),
-#                 transforms.BaselineRemoval(),
-#                 transforms.ToGrid(DEAP_CHANNEL_LOCATION_DICT)
-#             ]),
-#             online_transform=transforms.ToTensor(),
-#             label_transform=transforms.Compose([
-#                 transforms.Select('valence'),
-#                 transforms.Binary(5.0),
-#             ]),
-#             io_mode='pickle',
-#             chunk_size=128,
-#             baseline_chunk_size=128,
-#             num_baseline=3,
-#             num_worker=4)
-#        # the following codes
-#
-# .. note::
-#    LMDB may not be optimized for parts of Windows systems or storage devices. If you find that the data preprocessing speed is slow, you can consider setting :obj:`io_mode` to :obj:`pickle`, which is an alternative implemented by TorchEEG based on pickle.
-
-######################################################################
-# Step 2: Divide the Training and Test samples in the Dataset
-#
-# Here, the dataset is divided using 5-fold cross-validation. In the process of division, we group according to the trial index, and every trial takes 4 folds as training samples and 1 fold as test samples. Samples across trials are aggregated to obtain training set and test set.
-#
-
-k_fold = LeaveOneSubjectOut(
-    split_path='./tmp_out/examples_deap_domain_adaption/leave_one_subject_out')
-
-######################################################################
-# Step 3: Define the Model and Start Training
-#
-# We first use a loop to get the dataset in each cross-validation. In each cross-validation, we initialize the CCNN model and define the hyperparameters. For example, each EEG sample contains 4-channel features from 4 sub-bands, the grid size is 9 times 9, etc.
-#
-# We then initialize the trainer and set the hyperparameters in the trained model, such as the learning rate, the equipment used, etc. The :obj:`fit` method receives the training dataset and starts training the model. The :obj:`test` method receives a test dataset and reports the test results. The :obj:`save_state_dict` method can save the trained model.
-
-
-class Extractor(CCNN):
-
-    def forward(self, x):
-        x = self.conv1(x)
-        x = self.conv2(x)
-        x = self.conv3(x)
-        x = self.conv4(x)
-
-        x = x.flatten(start_dim=1)
-        return x
-
-
-class Classifier(CCNN):
-
-    def forward(self, x):
-        x = self.lin1(x)
-        x = self.lin2(x)
-        return x
-
-
-for i, (train_dataset, val_dataset) in enumerate(k_fold.split(dataset)):
-    # Initialize the model
-    extractor = Extractor(in_channels=4,
-                          grid_size=(9, 9),
-                          num_classes=2,
-                          dropout=0.5)
-    classifier = Classifier(in_channels=4,
-                            grid_size=(9, 9),
-                            num_classes=2,
-                            dropout=0.5)
-
-    # Initialize the trainer and use the 0-th GPU for training, or set device_ids=[] to use CPU
-    trainer = ADATrainer(extractor=extractor,
-                         classifier=classifier,
-                         devices=1,
-                         accelerator='gpu',
-                         num_classes=2)
-
-    # Initialize several batches of training samples and test samples
-    train_loader = DataLoader(train_dataset,
-                              batch_size=1000,
-                              shuffle=True,
-                              num_workers=4)
-    val_loader = DataLoader(val_dataset,
-                            batch_size=1000,
-                            shuffle=True,
-                            num_workers=4)
-
-    # Do 50 rounds of training
-    # validation set is used as the target domain
-    trainer.fit(train_loader, val_loader, val_loader, max_epochs=300)
-    trainer.test(val_loader)
-    trainer.save_state_dict(f'./tmp_out/examples_deap_domain_adaption/weight/{i}.pth')
diff --git a/examples/examples_deap_gan.py b/examples/examples_deap_gan.py
deleted file mode 100644
index f8596cb..0000000
--- a/examples/examples_deap_gan.py
+++ /dev/null
@@ -1,263 +0,0 @@
-"""GAN with the DEAP Dataset
-======================================
-In this case, we introduce how to use TorchEEG to train Generative Adversarial Networks (GAN) on the DEAP dataset for controllable EEG augmentation (which generates EEG signals confirming the given labels).
-
-Here, a zero game is played to optimize a generator and a discriminator. The generator is used to generate EEG signal samples according to the given labels, and the discriminator is used to distinguish generated samples from real samples. By confusing the discriminator, the generator will be able to produce samples that approximate the real distribution.
-"""
-
-import logging
-import os
-import random
-import time
-
-import numpy as np
-import torch
-import torch.autograd as autograd
-from tensorboardX import SummaryWriter
-from torch.utils.data.dataloader import DataLoader
-from torcheeg import transforms
-from torcheeg.datasets import DEAPDataset
-from torcheeg.datasets.constants.emotion_recognition.deap import (
-    DEAP_CHANNEL_LIST, DEAP_CHANNEL_LOCATION_DICT)
-from torcheeg.model_selection import KFoldGroupbyTrial
-from torcheeg.models import BCDiscriminator, BCGenerator
-from torcheeg.trainers import CGANTrainer
-from torcheeg.utils import plot_feature_topomap
-
-###############################################################################
-# Pre-experiment Preparation to Ensure Reproducibility
-# -----------------------------------------
-# Use the logging module to store output in a log file for easy reference while printing it to the screen.
-
-os.makedirs('./tmp_out/examples_deap_gan/log', exist_ok=True)
-logger = logging.getLogger('GAN with the DEAP Dataset')
-logger.setLevel(logging.DEBUG)
-console_handler = logging.StreamHandler()
-timeticks = time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime())
-file_handler = logging.FileHandler(
-    os.path.join('./tmp_out/examples_deap_gan/log', f'{timeticks}.log'))
-logger.addHandler(console_handler)
-logger.addHandler(file_handler)
-
-###############################################################################
-# Set the random number seed in all modules to guarantee the same result when running again.
-
-
-def seed_everything(seed):
-    random.seed(seed)
-    np.random.seed(seed)
-    os.environ["PYTHONHASHSEED"] = str(seed)
-    torch.manual_seed(seed)
-    torch.cuda.manual_seed(seed)
-    torch.backends.cudnn.deterministic = True
-    torch.backends.cudnn.benchmark = False
-
-
-seed_everything(42)
-
-###############################################################################
-# Customize Trainer
-# -----------------------------------------
-# TorchEEG provides many trainers to help complete the training of classification models, generative models and cross-domain methods. Here we choose the conditional generative adversarial network trainer (CGANTrainer), inherit the trainer and overload the log function to save the log using our own defined method; the :obj:`on_validation_step` hook function is overloaded to visualize EEG signals generated during validation using tensorboard and the tool function :obj:`plot_feature_topomap` in TorchEEG; other hook functions can also be overloaded to meet special needs.
-#
-
-writer = SummaryWriter(log_dir='./tmp_out/examples_deap_gan/vis',
-                       comment='examples_deap_gan')
-
-
-def gradient_penalty(model, real, fake, label=None):
-    device = real.device
-    real = real.data
-    fake = fake.data
-    alpha = torch.rand(real.size(0), *([1] * (len(real.shape) - 1))).to(device)
-    inputs = alpha * real + ((1 - alpha) * fake)
-    inputs.requires_grad_()
-
-    if label is None:
-        outputs = model(inputs)
-    else:
-        outputs = model(inputs, label)
-
-    gradient = autograd.grad(outputs=outputs,
-                             inputs=inputs,
-                             grad_outputs=torch.ones_like(outputs).to(device),
-                             create_graph=True,
-                             retain_graph=True,
-                             only_inputs=True)[0]
-
-    gradient = gradient.flatten(1)
-    return ((gradient.norm(2, dim=1) - 1)**2).mean()
-
-
-class MyCGANTrainer(CGANTrainer):
-    def log(self, *args, **kwargs):
-        if self.is_main:
-            logger.info(*args, **kwargs)
-
-    def before_validation_epoch(self, epoch_id, num_epochs, **kwargs):
-        # record the current epoch
-        self.cur_epoch = epoch_id
-        self.val_g_loss.reset()
-        self.val_d_loss.reset()
-
-    def on_validation_step(self, val_batch, batch_id, num_batches, **kwargs):
-        X = val_batch[0].to(self.device)
-        y = val_batch[1].to(self.device)
-
-        # for g_loss
-        z = torch.normal(mean=0,
-                         std=1,
-                         size=(X.shape[0],
-                               self.modules['generator'].in_channels)).to(
-                                   self.device)
-        gen_X = self.modules['generator'](z, y)
-        g_loss = -torch.mean(self.modules['discriminator'](gen_X, y))
-
-        # for d_loss
-        real_loss = self.modules['discriminator'](X, y)
-        fake_loss = self.modules['discriminator'](gen_X.detach(), y)
-        gp_term = gradient_penalty(self.modules['discriminator'], X, gen_X, y)
-        d_loss = -torch.mean(real_loss) + torch.mean(
-            fake_loss) + self.lambd * gp_term
-
-        self.val_g_loss.update(g_loss)
-        self.val_d_loss.update(d_loss)
-
-        vis_batch = num_batches // 20
-        if batch_id % vis_batch == 0:
-            t = transforms.ToInterpolatedGrid(DEAP_CHANNEL_LOCATION_DICT)
-            # center should be 0.0
-            signal = t.reverse(eeg=gen_X[y == 0][0].detach().cpu().numpy() -
-                               0.5)['eeg']
-            top_img = plot_feature_topomap(
-                torch.tensor(signal),
-                channel_list=DEAP_CHANNEL_LIST,
-                feature_list=["theta", "alpha", "beta", "gamma"])
-            # generate the visualization results and record them for the current epoch
-            writer.add_image(f'top{batch_id}/eeg-0',
-                             top_img,
-                             self.cur_epoch,
-                             dataformats='HWC')
-
-            signal = t.reverse(eeg=gen_X[y == 1][0].detach().cpu().numpy() -
-                               0.5)['eeg']
-            top_img = plot_feature_topomap(
-                torch.tensor(signal),
-                channel_list=DEAP_CHANNEL_LIST,
-                feature_list=["theta", "alpha", "beta", "gamma"])
-            writer.add_image(f'top{batch_id}/eeg-1',
-                             top_img,
-                             self.cur_epoch,
-                             dataformats='HWC')
-
-
-######################################################################
-# Building Deep Learning Pipelines Using TorchEEG
-# -----------------------------------------
-# Step 1: Initialize the Dataset
-#
-# We use the DEAP dataset supported by TorchEEG. Here, we set an EEG sample to 1 second long and include 128 data points. The baseline signal is 3 seconds long, cut into three, and averaged as the baseline signal for the trial. In offline preprocessing, we divide the EEG signal of every electrode into 4 sub-bands, and calculate the differential entropy on each sub-band as a feature, followed by debaselining and mapping on the grid. Finally, the preprocessed EEG signals are stored in the local IO. In online processing, all EEG signals are converted into Tensors and normalized (in GANs normalization helps with convergence) for input into neural networks.
-#
-
-dataset = DEAPDataset(io_path=f'./tmp_out/examples_deap_gan/deap',
-                      root_path='./tmp_in/data_preprocessed_python',
-                      offline_transform=transforms.Compose([
-                          transforms.BandDifferentialEntropy(
-                              sampling_rate=128, apply_to_baseline=True),
-                          transforms.BaselineRemoval(),
-                          transforms.ToGrid(DEAP_CHANNEL_LOCATION_DICT)
-                      ]),
-                      online_transform=transforms.Compose([
-                          transforms.MinMaxNormalize(),
-                          transforms.ToTensor(),
-                      ]),
-                      label_transform=transforms.Compose([
-                          transforms.Select('valence'),
-                          transforms.Binary(5.0),
-                      ]),
-                      chunk_size=128,
-                      baseline_chunk_size=128,
-                      num_baseline=3,
-                      num_worker=4)
-
-######################################################################
-# .. warning::
-#    If you use TorchEEG under the `Windows` system and want to use multiple processes (such as in dataset or dataloader), you should check whether :obj:`__name__` is :obj:`__main__` to avoid errors caused by multiple :obj:`import`.
-#
-# That is, under the :obj:`Windows` system, you need to:
-#  .. code-block::
-#
-#    if __name__ == "__main__":
-#        dataset = DEAPDataset(
-#             io_path=
-#             f'./tmp_out/examples_deap_gan/deap',
-#             root_path='./tmp_in/data_preprocessed_python',
-#             offline_transform=transforms.Compose([
-#                 transforms.BandDifferentialEntropy(sampling_rate=128,
-#                                                    apply_to_baseline=True),
-#                 transforms.BaselineRemoval(),
-#                 transforms.ToGrid(DEAP_CHANNEL_LOCATION_DICT)
-#             ]),
-#             online_transform=transforms.Compose([
-#                 transforms.MinMaxNormalize(),
-#                 transforms.ToTensor(),
-#             ]),
-#             label_transform=transforms.Compose([
-#                 transforms.Select('valence'),
-#                 transforms.Binary(5.0),
-#             ]),
-#             io_mode='pickle',
-#             chunk_size=128,
-#             baseline_chunk_size=128,
-#             num_baseline=3,
-#             num_worker=4)
-#        # the following codes
-#
-# .. note::
-#    LMDB may not be optimized for parts of Windows systems or storage devices. If you find that the data preprocessing speed is slow, you can consider setting :obj:`io_mode` to :obj:`pickle`, which is an alternative implemented by TorchEEG based on pickle.
-
-######################################################################
-# Step 2: Divide the Training and Test samples in the Dataset
-#
-# Here, the dataset is divided using 5-fold cross-validation. In the process of division, we group according to the trial index, and every trial takes 4 folds as training samples and 1 fold as test samples. Samples across trials are aggregated to obtain training set and test set.
-#
-
-k_fold = KFoldGroupbyTrial(n_splits=5,
-                           split_path='./tmp_out/examples_deap_gan/split')
-
-######################################################################
-# Step 3: Define the Model and Start Training
-#
-# We first use a loop to get the dataset in each cross-validation. In each cross-validation, we initialize the generator and discriminator models and define the hyperparameters. For example, For example, we want to generate the differential entropy features of the 4 sub-bands of the simulated EEG signal.  The generated samples are sampled and transformed from a feature space of 128.
-#
-# We then initialize the trainer and set the hyperparameters in the trained model, such as the learning rate, the equipment used, etc. The :obj:`fit` method receives the training dataset and starts training the model. The :obj:`test` method receives a test dataset and reports the test results. The :obj:`save_state_dict` method can save the trained model.
-
-for i, (train_dataset, val_dataset) in enumerate(k_fold.split(dataset)):
-    # Initialize the model
-    generator = BCGenerator(in_channels=128, out_channels=4, num_classes=2)
-    discriminator = BCDiscriminator(hid_channels=128,
-                                    in_channels=4,
-                                    num_classes=2)
-
-    # Initialize the trainer and use the 0-th GPU for training, or set device_ids=[] to use CPU
-    trainer = MyCGANTrainer(generator=generator,
-                            discriminator=discriminator,
-                            generator_lr=0.0001,
-                            discriminator_lr=0.00001,
-                            weight_decay=0,
-                            device_ids=[0])
-
-    # Initialize several batches of training samples and test samples
-    train_loader = DataLoader(train_dataset,
-                              batch_size=256,
-                              shuffle=True,
-                              num_workers=4)
-    val_loader = DataLoader(val_dataset,
-                            batch_size=256,
-                            shuffle=False,
-                            num_workers=4)
-
-    # Do 50 rounds of training
-    trainer.fit(train_loader, val_loader, num_epochs=50)
-    trainer.test(val_loader)
-    trainer.save_state_dict(f'./tmp_out/examples_deap_gan/weight/{i}.pth')
diff --git a/examples/examples_dreamer_tsception.py b/examples/examples_dreamer_tsception.py
index c4cd136..cabd372 100644
--- a/examples/examples_dreamer_tsception.py
+++ b/examples/examples_dreamer_tsception.py
@@ -1,142 +1,77 @@
-"""TSCeption with the DREAMER Dataset
-======================================
-In this case, we introduce how to use TorchEEG to train a TSCeption on the DREAMER dataset for emotion classification.
 """
+Train a Tsception Model on the DREAMER Dataset
+==========================
 
-import logging
-import os
-import random
-import time
-
-import numpy as np
-import torch
-from torch.utils.data.dataloader import DataLoader
-from torcheeg import transforms
-from torcheeg.datasets import DREAMERDataset
-from torcheeg.datasets.constants.emotion_recognition.dreamer import DREAMER_CHANNEL_LOCATION_DICT
-from torcheeg.model_selection import KFoldGroupbyTrial
-from torcheeg.models import TSCeption
-from torcheeg.trainers import ClassificationTrainer
-
-###############################################################################
-# Pre-experiment Preparation to Ensure Reproducibility
-# -----------------------------------------
-# Use the logging module to store output in a log file for easy reference while printing it to the screen.
-
-os.makedirs('./tmp_out/examples_dreamer_tsception/log', exist_ok=True)
-logger = logging.getLogger('TSCeption with the DREAMER Dataset')
-logger.setLevel(logging.DEBUG)
-console_handler = logging.StreamHandler()
-timeticks = time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime())
-file_handler = logging.FileHandler(
-    os.path.join('./tmp_out/examples_dreamer_tsception/log',
-                 f'{timeticks}.log'))
-logger.addHandler(console_handler)
-logger.addHandler(file_handler)
-
-###############################################################################
-# Set the random number seed in all modules to guarantee the same result when running again.
-
-
-def seed_everything(seed):
-    random.seed(seed)
-    np.random.seed(seed)
-    os.environ["PYTHONHASHSEED"] = str(seed)
-    torch.manual_seed(seed)
-    torch.cuda.manual_seed(seed)
-    torch.backends.cudnn.deterministic = True
-    torch.backends.cudnn.benchmark = False
-
-
-seed_everything(42)
-
-###############################################################################
-# Customize Trainer
-# -----------------------------------------
-# TorchEEG provides a large number of trainers to help complete the training of classification models, generative models and cross-domain methods. Here we choose the simplest classification trainer, inherit the trainer and overload the log function to save the log using our own defined method; other hook functions can also be overloaded to meet special needs.
-#
-
-
-class MyClassificationTrainer(ClassificationTrainer):
-    def log(self, *args, **kwargs):
-        if self.is_main:
-            logger.info(*args, **kwargs)
+In this tutorial, we'll showcase how you can use TorchEEG to train a Tsception model on the DREAMER dataset for the task of emotion classification. Let's go through the process step-by-step, covering everything from dataset preparation to model evaluation.
 
+"""
 
 ######################################################################
-# Building Deep Learning Pipelines Using TorchEEG
-# -----------------------------------------
 # Step 1: Initialize the Dataset
+# ----------------------------------------------
 #
-# We use the AMIGOS dataset supported by TorchEEG. Here, we set an EEG sample to 1 second long and include 128 data points. The baseline signal is 61 seconds long, cut into 61, and averaged as the baseline signal for the trial. In offline preprocessing, all EEG signals are debaselined and normalized, and the preprocessed EEG signals are stored in the local IO. In online processing, all EEG signals are converted into Tensors for input into neural networks.
+# First off, we'll use the DREAMER dataset provided by TorchEEG. Each EEG sample in this dataset spans 1 second and contains 128 data points. For each trial, the baseline signal lasts 61 seconds. 
+# During offline preprocessing, we perform several steps:
+# - Divide each electrode's EEG signal into 4 frequency sub-bands
+# - Compute the differential entropy of each sub-band as a feature
+# - Eliminate the baseline from the signal
+# - Map the preprocessed signals onto a 2D grid
+# For online processing, we convert the EEG signals into PyTorch Tensors to make them compatible with neural network inputs.
+# Let's see how to accomplish these steps in code.
 #
 
-dataset = DREAMERDataset(
-    io_path=f'./tmp_out/examples_dreamer_tsception/dreamer',
-    mat_path='./tmp_in/DREAMER.mat',
-    offline_transform=transforms.Compose([
-        transforms.BaselineRemoval(),
-        transforms.MeanStdNormalize(),
-        transforms.To2d()
-    ]),
-    online_transform=transforms.ToTensor(),
-    label_transform=transforms.Compose(
-        [transforms.Select('valence'),
-         transforms.Binary(3.0)]),
-    chunk_size=128,
-    baseline_chunk_size=128,
-    num_baseline=61,
-    num_worker=4)
+from torcheeg.datasets import DREAMERDataset
+from torcheeg import transforms
 
-######################################################################
-# .. warning::
-#    If you use TorchEEG under the `Windows` system and want to use multiple processes (such as in dataset or dataloader), you should check whether :obj:`__name__` is :obj:`__main__` to avoid errors caused by multiple :obj:`import`.
-#
-# That is, under the :obj:`Windows` system, you need to:
-#  .. code-block::
-#
-#    if __name__ == "__main__":
-#        dataset = DREAMERDataset(io_path='./tmp_out/examples_dreamer_tsception/amigos',
-#                         root_path='./tmp_in/data_preprocessed',
-#                         offline_transform=transforms.Compose([
-#                             transforms.BaselineRemoval(),
-#                             transforms.MeanStdNormalize()
-#                         ]),
-#                         online_transform=transforms.ToTensor(),
-#                         label_transform=transforms.Compose([
-#                             transforms.Select('valence'),
-#                             transforms.Binary(5.0)
-#                         ]),
-#                         io_mode='pickle',
-#                         chunk_size=128,
-#                         baseline_chunk_size=128,
-#                         num_baseline=5,
-#                         num_worker=4)
-#        # the following codes
-#
-# .. note::
-#    LMDB may not be optimized for parts of Windows systems or storage devices. If you find that the data preprocessing speed is slow, you can consider setting :obj:`io_mode` to :obj:`pickle`, which is an alternative implemented by TorchEEG based on pickle.
+dataset = DREAMERDataset(io_path='./examples_dreamer_tsception/dreamer',
+                         mat_path='./DREAMER.mat',
+                         offline_transform=transforms.Compose([
+                             transforms.BaselineRemoval(),
+                             transforms.MeanStdNormalize(),
+                             transforms.To2d()
+                         ]),
+                         online_transform=transforms.ToTensor(),
+                         label_transform=transforms.Compose([
+                             transforms.Select('valence'),
+                             transforms.Binary(3.0)
+                         ]),
+                         chunk_size=128,
+                         baseline_chunk_size=128,
+                         num_baseline=61,
+                         num_worker=4)
 
 ######################################################################
-# Step 2: Divide the Training and Test samples in the Dataset
+# Step 2: Divide the Training and Test Samples in the Dataset
+# ----------------------------------------------
 #
-# Here, the dataset is divided using 5-fold cross-validation. In the process of division, we group according to the trial index, and every trial takes 4 folds as training samples and 1 fold as test samples. Samples across trials are aggregated to obtain training set and test set.
+# Next, let's partition our dataset into training and test sets using 5-fold cross-validation. We group the data based on their trial index, where each trial contributes 4 folds to the training set and 1 fold to the test set. These grouped samples are then combined to form the final training and test sets.
 #
 
-k_fold = KFoldGroupbyTrial(
-    n_splits=5, split_path='./tmp_out/examples_dreamer_tsception/split')
+from torcheeg.model_selection import KFoldGroupbyTrial
+
+k_fold = KFoldGroupbyTrial(n_splits=5,
+                           split_path=f'./examples_dreamer_tsception/split')
 
 ######################################################################
-# Step 3: Define the Model and Start Training
+# Step 3: Define the Model and Initiate Training
+# ----------------------------------------------
 #
-# We first use a loop to get the dataset in each cross-validation.
-# In each cross-validation, we initialize the TSCeption model, and define the hyperparameters. For example, the sampling rate of EEG sample is 128, there are 15 temporal modules and 15 spatial modules, etc.
+# Having prepared and divided our dataset, we can now move on to model building and training. In each iteration of the cross-validation loop, we'll initialize the Tsception model and set its hyperparameters.
+# Here, each EEG sample contains 128 data points across 14 electrodes. We'll train this model for 50 epochs using TorchEEG's `ClassifierTrainer`.
 #
-# We then initialize the trainer and set the hyperparameters in the trained model, such as the learning rate, the equipment used, etc. The :obj:`fit` method receives the training dataset and starts training the model. The :obj:`test` method receives a test dataset and reports the test results. The :obj:`save_state_dict` method can save the trained model.
+
+from torch.utils.data import DataLoader
+from torcheeg.models import TSCeption
+
+from torcheeg.trainers import ClassifierTrainer
+
+import pytorch_lightning as pl
 
 for i, (train_dataset, val_dataset) in enumerate(k_fold.split(dataset)):
-    # Initialize the model
-    model = TSCeption(num_electrodes=len(DREAMER_CHANNEL_LOCATION_DICT),
+    train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
+    val_loader = DataLoader(val_dataset, batch_size=64, shuffle=False)
+
+    model = TSCeption(num_electrodes=14,
                       num_classes=2,
                       num_T=15,
                       num_S=15,
@@ -145,24 +80,20 @@ def log(self, *args, **kwargs):
                       sampling_rate=128,
                       dropout=0.5)
 
-    # Initialize the trainer and use the 0-th GPU for training, or set device_ids=[] to use CPU
-    trainer = MyClassificationTrainer(model=model,
-                                      lr=1e-4,
-                                      weight_decay=1e-4,
-                                      device_ids=[0])
-
-    # Initialize several batches of training samples and test samples
-    train_loader = DataLoader(train_dataset,
-                              batch_size=256,
-                              shuffle=True,
-                              num_workers=4)
-    val_loader = DataLoader(val_dataset,
-                            batch_size=256,
-                            shuffle=False,
-                            num_workers=4)
-
-    # Do 50 rounds of training
-    trainer.fit(train_loader, val_loader, num_epochs=50)
-    trainer.test(val_loader)
-    trainer.save_state_dict(
-        f'./tmp_out/examples_dreamer_tsception/weight/{i}.pth')
+    trainer = ClassifierTrainer(model=model,
+                                num_classes=2,
+                                lr=1e-4,
+                                weight_decay=1e-4,
+                                accelerator="gpu")
+    trainer.fit(train_loader,
+                val_loader,
+                max_epochs=50,
+                default_root_dir=f'./examples_dreamer_tsception/model/{i}',
+                callbacks=[pl.callbacks.ModelCheckpoint(save_last=True)],
+                enable_progress_bar=True,
+                enable_model_summary=True,
+                limit_val_batches=0.0)
+    score = trainer.test(val_loader,
+                         enable_progress_bar=True,
+                         enable_model_summary=True)[0]
+    print(f'Fold {i} test accuracy: {score["test_accuracy"]:.4f}')
\ No newline at end of file
diff --git a/examples/examples_mne_dataset.py b/examples/examples_mne_dataset.py
deleted file mode 100644
index 4f9cf8f..0000000
--- a/examples/examples_mne_dataset.py
+++ /dev/null
@@ -1,200 +0,0 @@
-"""Examples of MNEDataset
-======================================
-In this case, we will introduce how to combine TorchEEG with MNE, and use deep learning algorithms to analyze the existing :obj:`mne.Epochs` format data.
-"""
-
-import logging
-import os
-import random
-import time
-
-import mne
-import numpy as np
-import torch
-from torch.utils.data.dataloader import DataLoader
-
-from torcheeg import transforms
-from torcheeg.datasets import MNEDataset
-from torcheeg.model_selection import KFold
-from torcheeg.models import TSCeption
-from torcheeg.trainers import ClassificationTrainer
-
-###############################################################################
-# Pre-experiment Preparation to Ensure Reproducibility
-# -----------------------------------------
-# Use the logging module to store output in a log file for easy reference while printing it to the screen.
-
-os.makedirs('./tmp_out/examples_mne_dataset/log', exist_ok=True)
-logger = logging.getLogger('Examples of MNEDataset')
-logger.setLevel(logging.DEBUG)
-console_handler = logging.StreamHandler()
-timeticks = time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime())
-file_handler = logging.FileHandler(
-    os.path.join('./tmp_out/examples_mne_dataset/log', f'{timeticks}.log'))
-logger.addHandler(console_handler)
-logger.addHandler(file_handler)
-
-###############################################################################
-# Set the random number seed in all modules to guarantee the same result when running again.
-
-
-def seed_everything(seed):
-    random.seed(seed)
-    np.random.seed(seed)
-    os.environ["PYTHONHASHSEED"] = str(seed)
-    torch.manual_seed(seed)
-    torch.cuda.manual_seed(seed)
-    torch.backends.cudnn.deterministic = True
-    torch.backends.cudnn.benchmark = False
-
-
-seed_everything(42)
-
-###############################################################################
-# Customize Trainer
-# -----------------------------------------
-# TorchEEG provides a large number of trainers to help complete the training of classification models, generative models and cross-domain methods. Here we choose the simplest classification trainer, inherit the trainer and overload the log function to save the log using our own defined method; other hook functions can also be overloaded to meet special needs.
-#
-
-
-class MyClassificationTrainer(ClassificationTrainer):
-    def log(self, *args, **kwargs):
-        if self.is_main:
-            logger.info(*args, **kwargs)
-
-
-###############################################################################
-# Read data using MNE and formalize as :obj:`mne.Epochs`
-# -----------------------------------------
-# We use mne's API to automatically download the motor imagery dataset in Physionet. The EEG signals of subjects 1-21 in runs 6, 10, and 14 were downloaded and filtered. We store multiple :obj:`mne.Epochs` into an array, and use a counterpart array :obj:`metadata_list` to describe the metadata corresponding to the corresponding Epochs.
-#
-metadata_list = [{
-    'subject': subject_id,
-    'run': run_id
-} for subject_id in range(1, 22)
-                 for run_id in [6, 10, 14]]  # motor imagery: hands vs feet
-
-epochs_list = []
-for metadata in metadata_list:
-    physionet_path = mne.datasets.eegbci.load_data(metadata['subject'],
-                                                   metadata['run'],
-                                                   update_path=False)[0]
-    raw = mne.io.read_raw_edf(physionet_path, preload=True, stim_channel='auto')
-    mne.datasets.eegbci.standardize(raw)
-
-    montage = mne.channels.make_standard_montage('standard_1005')
-    raw.set_montage(montage)
-
-    raw.filter(7., 30., fir_design='firwin', skip_by_annotation='edge')
-    events, _ = mne.events_from_annotations(raw, event_id=dict(T1=2, T2=3))
-    picks = mne.pick_types(raw.info,
-                           meg=False,
-                           eeg=True,
-                           stim=False,
-                           eog=False,
-                           exclude='bads')
-    # init Epochs with raw EEG signals and corresponding event annotations. Here, tmin is set to -1., and tmax is set to 4.0, to avoid classification of evoked responses by using epochs that start 1s after cue onset.
-    epochs_list.append(
-        mne.Epochs(raw,
-                   events,
-                   dict(hands=2, feet=3),
-                   tmin=-1.,
-                   tmax=4.0,
-                   proj=True,
-                   picks=picks))
-
-###############################################################################
-# Convert :obj:`mne.Epochs` into MNEDataset
-# -----------------------------------------
-# We use MNEDataset to window the Epochs. Here, we set the window size to 160 (1-second long) and the overlap to 80 to segment the EEG signal corresponding to each event. The corresponding information in :obj:`metadata_list` will be assigned to the corresponding window. At the same time, the window also includes the start position of the window :obj:`start_at`, the end position :obj:`start_at`, the epoch index :obj:`trial_id` and the corresponding event type :obj:`event`, which can be used and transformed as label.
-#
-dataset = MNEDataset(epochs_list=epochs_list,
-                     metadata_list=metadata_list,
-                     chunk_size=160,
-                     overlap=80,
-                     io_path='./tmp_out/examples_mne_dataset/physionet',
-                     offline_transform=transforms.Compose(
-                         [transforms.MeanStdNormalize(),
-                          transforms.To2d()]),
-                     online_transform=transforms.ToTensor(),
-                     label_transform=transforms.Compose([
-                         transforms.Select('event'),
-                         transforms.Lambda(lambda x: x - 2)
-                     ]),
-                     num_worker=2)
-
-######################################################################
-# .. warning::
-#    If you use TorchEEG under the `Windows` system and want to use multiple processes (such as in dataset or dataloader), you should check whether :obj:`__name__` is :obj:`__main__` to avoid errors caused by multiple :obj:`import`.
-#
-# That is, under the :obj:`Windows` system, you need to:
-#  .. code-block::
-#
-#    if __name__ == "__main__":
-#        dataset = MNEDataset(epochs_list=epochs_list,
-#                      metadata_list=metadata_list,
-#                      chunk_size=160,
-#                      overlap=80,
-#                      io_path='./tmp_out/examples_mne_dataset/physionet',
-#                      offline_transform=transforms.Compose(
-#                          [transforms.MeanStdNormalize(),
-#                           transforms.To2d()]),
-#                      online_transform=transforms.ToTensor(),
-#                      label_transform=transforms.Compose([
-#                          transforms.Select('event'),
-#                          transforms.Lambda(lambda x: x - 2)
-#                      ]),
-#                      io_mode='pickle',
-#                      num_worker=2)
-#        # the following codes
-#
-# .. note::
-#    LMDB may not be optimized for parts of Windows systems or storage devices. If you find that the data preprocessing speed is slow, you can consider setting :obj:`io_mode` to :obj:`pickle`, which is an alternative implemented by TorchEEG based on pickle.
-
-######################################################################
-# Step 2: Divide the Training and Test samples in the Dataset
-#
-# Here, the dataset is divided using 5-fold cross-validation. In the process of division, the total dataset takes 4 folds as training samples and 1 fold as test samples.
-#
-k_fold = KFold(n_splits=5, split_path='./tmp_out/examples_mne_dataset/split')
-
-######################################################################
-# Step 3: Define the Model and Start Training
-#
-# We first use a loop to get the dataset in each cross-validation.
-# In each cross-validation, we initialize the TSCeption model, and define the hyperparameters. For example, the sampling rate of EEG sample is 160, there are 15 temporal modules and 15 spatial modules, etc. In this example, the shape of EEG signals is :obj:`[60, 160]` and the number of classes is 2.
-#
-# We then initialize the trainer and set the hyperparameters in the trained model, such as the learning rate, the equipment used, etc. The :obj:`fit` method receives the training dataset and starts training the model. The :obj:`test` method receives a test dataset and reports the test results. The :obj:`save_state_dict` method can save the trained model.
-#
-
-for i, (train_dataset, val_dataset) in enumerate(k_fold.split(dataset)):
-    # Initialize the model
-    model = TSCeption(num_electrodes=64,
-                      num_classes=2,
-                      num_T=15,
-                      num_S=15,
-                      in_channels=1,
-                      hid_channels=32,
-                      sampling_rate=160,
-                      dropout=0.5)
-
-    # Initialize the trainer and use the 0-th GPU for training, or set device_ids=[] to use CPU
-    trainer = MyClassificationTrainer(model=model,
-                                      lr=1e-4,
-                                      weight_decay=1e-4,
-                                      device_ids=[0])
-
-    # Initialize several batches of training samples and test samples
-    train_loader = DataLoader(train_dataset,
-                              batch_size=256,
-                              shuffle=True,
-                              num_workers=4)
-    val_loader = DataLoader(val_dataset,
-                            batch_size=256,
-                            shuffle=False,
-                            num_workers=4)
-
-    # Do 50 rounds of training
-    trainer.fit(train_loader, val_loader, num_epochs=50)
-    trainer.test(val_loader)
-    trainer.save_state_dict(f'./tmp_out/examples_mne_dataset/weight/{i}.pth')
diff --git a/examples/examples_model_selection.py b/examples/examples_model_selection.py
new file mode 100644
index 0000000..58e5d3b
--- /dev/null
+++ b/examples/examples_model_selection.py
@@ -0,0 +1,113 @@
+"""
+Introduction to the model_selection Module
+=============================================
+
+Welcome to the guide on TorchEEG's ``model_selection`` Module! This module provides you with a rich array of dataset-splitting strategies.
+
+In EEG-based emotion recognition, the division of training and test sets
+may differ based on research objectives. The ``model_selection`` module
+addresses this need by supporting a variety of splitting strategies
+considering trials, sessions, and subjects.
+
+"""
+
+######################################################################
+# The Elementary Method, KFold
+# ----------------------------------------------
+#
+# The simplest approach is ``KFold``, where the dataset is randomly
+# divided into ‘k’ subsets. Each subset serves as the test set in one of
+# the ‘k’ experiments, while the remaining subsets form the training set.
+# The model’s performance is reported as the average result across the ‘k’
+# test sets.
+#
+
+from torcheeg.datasets import DEAPDataset
+from torcheeg.model_selection import KFold
+
+dataset = DEAPDataset(io_path='./deap',
+                      root_path='./tmp_in/data_preprocessed_python')
+
+k_fold = KFold(split_path='./examples_model_selection_1/split')
+for train_dataset, test_dataset in k_fold.split(dataset):
+    print(len(train_dataset))
+    print(len(test_dataset))
+
+######################################################################
+# Important Notes
+# ----------------------------------------------
+#
+# Since TorchEEG automatically caches the divided indices to
+# ``split_path``, there’s no need to worry about encountering different
+# partitions each time the program runs. However, when you want to
+# generate a new partition or use other parameters for partitioning,
+# please delete the old partition at ``split_path`` or define a new
+# ``split_path`` to create a new partition cache.
+#
+# .. code-block:: bash
+#
+#   !rm -rf ./deap
+#
+
+######################################################################
+# Advanced KFold Variants
+# ----------------------------------------------
+#
+# To evaluate model performance under unseen trials or periods, TorchEEG
+# provides ``KFoldGroupbyTrial`` and ``KFoldCrossTrial``. The former
+# divides each trial into ‘k’ periods; each period serves as the test set
+# in one of the ‘k’ experiments, while the remaining subsets form the
+# training set. The latter takes into account inter-trial differences and
+# splits each subject’s trials into ‘k’ folds, including only specific
+# trials’ EEG samples in the training set.
+#
+
+from torcheeg.model_selection import KFoldCrossTrial
+
+k_fold = KFoldCrossTrial(split_path='./examples_model_selection_2/split')
+for train_dataset, test_dataset in k_fold.split(dataset):
+    print(len(train_dataset))
+    print(len(test_dataset))
+
+from torcheeg.model_selection import KFoldGroupbyTrial
+
+k_fold = KFoldGroupbyTrial(split_path='./examples_model_selection_3/split')
+for train_dataset, test_dataset in k_fold.split(dataset):
+    print(len(train_dataset))
+    print(len(test_dataset))
+
+######################################################################
+# Cross-Subject Scenarios
+# ----------------------------------------------
+#
+# For the prevalent research scenario of cross-subject problems, TorchEEG
+# further provides ``LeaveOneSubjectOut`` and ``KFoldCrossSubject``.
+# ``KFoldCrossSubject`` divides subjects into ‘k’ folds, with the training
+# set comprising only ‘k-1’ folds of subjects’ EEG samples.
+# ``LeaveOneSubjectOut`` is a special case where ‘k’ equals the total
+# number of subjects.
+#
+
+from torcheeg.model_selection import LeaveOneSubjectOut
+
+loso = LeaveOneSubjectOut(split_path='./examples_model_selection_3/split')
+for train_dataset, test_dataset in loso.split(dataset):
+    print(len(train_dataset))
+    print(len(test_dataset))
+
+######################################################################
+# Per-Subject Evaluation
+# ----------------------------------------------
+#
+# Additionally, we offer ``PerSubject`` versions of ``KFold``,
+# ``KFoldGroupbyTrial``, and ``KFoldCrossTrial``, among others. These are
+# specifically designed for research focused on individual subject model
+# performance.
+#
+
+from torcheeg.model_selection import KFoldGroupbyTrial
+
+loso = KFoldGroupbyTrial(split_path='./examples_model_selection_3/split')
+for train_dataset, test_dataset in loso.split(dataset):
+    print(len(train_dataset))
+    print(len(test_dataset))
\ No newline at end of file
diff --git a/examples/examples_models.py b/examples/examples_models.py
new file mode 100644
index 0000000..6880ed3
--- /dev/null
+++ b/examples/examples_models.py
@@ -0,0 +1,200 @@
+"""
+Introduction to the models Module
+====================================
+
+Welcome to the guide on TorchEEG's ``models`` Module! This module provides you with a variety of both discriminative and generative models. Each model has been meticulously replicated from its respective research studies, offering a toolkit for EEG signal analysis.
+
+"""
+
+
+######################################################################
+# Discriminative Models
+# ----------------------------------------------
+#
+# Regarding discriminative models, the models module houses
+# implementations of Convolutional Neural Networks (CNNs), such as EEGNet
+# and TSception. These models create a variety of spatial, temporal, or
+# spectral tensor representations of EEG signals and execute local pattern
+# analysis via convolution.
+# 
+
+import torch
+from torcheeg.models.cnn import TSCeption
+
+eeg = torch.randn(1, 1, 28, 512)
+model = TSCeption(num_classes=2,
+                  num_electrodes=28,
+                  sampling_rate=128,
+                  num_T=15,
+                  num_S=15,
+                  hid_channels=32,
+                  dropout=0.5)
+pred = model(eeg)
+
+
+######################################################################
+# To understand how to preprocess datasets into formats required by the
+# models, please refer to the corresponding documentation of each model.
+# 
+
+from torcheeg.datasets import DEAPDataset
+from torcheeg.transforms import transforms
+from torcheeg.utils import DEAP_CHANNEL_LIST
+
+dataset = DEAPDataset(io_path=f'./deap',
+                    root_path='./data_preprocessed_python',
+                    chunk_size=512,
+                    num_baseline=1,
+                    baseline_chunk_size=512,
+                    offline_transform=transforms.Compose([
+                        transforms.PickElectrode(transforms.PickElectrode.to_index_list(
+                        ['FP1', 'AF3', 'F3', 'F7',
+                        'FC5', 'FC1', 'C3', 'T7',
+                        'CP5', 'CP1', 'P3', 'P7',
+                        'PO3','O1', 'FP2', 'AF4',
+                        'F4', 'F8', 'FC6', 'FC2',
+                        'C4', 'T8', 'CP6', 'CP2',
+                        'P4', 'P8', 'PO4', 'O2'], DEAP_CHANNEL_LIST)),
+                        transforms.To2d()
+                    ]),
+                    online_transform=transforms.ToTensor(),
+                    label_transform=transforms.Compose([
+                        transforms.Select('valence'),
+                        transforms.Binary(5.0),
+                    ]))
+
+model = TSCeption(num_classes=2,
+                  num_electrodes=28,
+                  sampling_rate=128,
+                  num_T=15,
+                  num_S=15,
+                  hid_channels=32,
+                  dropout=0.5)
+
+
+######################################################################
+# The module also includes Recurrent Neural Networks (RNNs) like GRU and
+# LSTM. These models view EEG signals as multivariate time series data and
+# construct recurrent modules for emotion decoding.
+# 
+
+from torcheeg.models import GRU
+
+model = GRU(num_electrodes=32, hid_channels=64, num_classes=2)
+
+eeg = torch.randn(2, 32, 128)
+pred = model(eeg)
+
+
+######################################################################
+# In the field of Graph Neural Networks (GNNs), the models module
+# incorporates renowned networks like DGCNN, RGNN, and LGGNet. These
+# models aim to analyze the functional connections between electrodes by
+# depicting them as a graph network and designing graph convolution
+# kernels.
+# 
+
+from torcheeg.models import DGCNN
+
+eeg = torch.randn(1, 62, 200)
+model = DGCNN(in_channels=200, num_electrodes=62, hid_channels=32, num_layers=2,num_classes=2)
+pred = model(eeg)
+
+
+######################################################################
+# In recent years, the increasing popularity of Transformer-based models,
+# such as EEG-ConvTransformer, has been recognized in the models module.
+# These models predominantly utilize various self-attention mechanisms to
+# analyze electrode correlations, delivering valuable insights.
+# 
+
+from torcheeg.models import SimpleViT
+
+eeg = torch.randn(1, 128, 9, 9)
+model = SimpleViT(chunk_size=128, t_patch_size=32, s_patch_size=(3, 3), num_classes=2)
+pred = model(eeg)
+
+
+######################################################################
+# Generative Models
+# ----------------------------------------------
+#
+# We also provide a variety of generative models known for their
+# impressive advancements in computer vision, natural language processing,
+# and other domains. When applied to EEG analysis, four categories of
+# generative models are offered as sturdy benchmarks for researchers. The
+# Generative Adversarial Network (GAN), for instance, WGAN-GP, includes a
+# generator and discriminator. The discriminator learns to distinguish
+# between real and generated signals, and the generator is trained to
+# approximate accurate signals through adversarial training against the
+# discriminator.
+# 
+
+from torcheeg.models import BCGenerator, BCDiscriminator
+
+g_model = BCGenerator(in_channels=128, num_classes=3)
+d_model = BCDiscriminator(in_channels=4, num_classes=3)
+z = torch.normal(mean=0, std=1, size=(1, 128))
+y = torch.randint(low=0, high=3, size=(1, ))
+fake_X = g_model(z, y)
+disc_X = d_model(fake_X, y)
+
+
+######################################################################
+# The Variational Auto-Encoder (VAE), such as Beta VAE, equipped with an
+# encoder and decoder, maps observed EEG signals into a latent space using
+# the encoder and then employs the decoder to reproduce EEG signals.
+# 
+
+from torcheeg.models import BCEncoder, BCDecoder
+
+encoder = BCEncoder(in_channels=4, num_classes=3)
+decoder = BCDecoder(in_channels=64, out_channels=4, num_classes=3)
+y = torch.randint(low=0, high=3, size=(1, ))
+mock_eeg = torch.randn(1, 4, 9, 9)
+mu, logvar = encoder(mock_eeg, y)
+std = torch.exp(0.5 * logvar)
+eps = torch.randn_like(std)
+z = eps * std + mu
+fake_X = decoder(z, y)
+
+
+######################################################################
+# Normalizing Flow, for example, Glow, offers a series of invertible
+# transformations. It learns the sequence of reversible transformations to
+# convert EEG signals into latent variables and utilizes the reverse of
+# the flow function to return to samples for generation.
+# 
+
+import torch.nn.functional as F
+
+import torch
+from torcheeg.models import BCGlow
+
+model = BCGlow(num_classes=2)
+# forward to calculate loss function
+mock_eeg = torch.randn(2, 4, 32, 32)
+y = torch.randint(0, 2, (2, ))
+
+nll_loss, y_logits = model(mock_eeg, y)
+loss = nll_loss.mean() + F.cross_entropy(y_logits, y)
+
+# sample a generated result
+fake_X = model.sample(y, temperature=1.0)
+
+
+######################################################################
+# The Diffusion Model, such as DDPM, introduces a sequential corruption of
+# observed data with increasing noise and learns to reverse this process.
+# The generation process inverts this diffusion process, starting with
+# white noise and gradually denoising it into corresponding observed EEG
+# signals.
+# 
+
+from torcheeg.models import BCUNet
+
+unet = BCUNet(num_classes=2)
+mock_eeg = torch.randn(2, 4, 9, 9)
+t = torch.randint(low=1, high=1000, size=(2, ))
+y = torch.randint(low=0, high=2, size=(1, ))
+fake_X = unet(mock_eeg, t, y)
\ No newline at end of file
diff --git a/examples/examples_multi_gpu.py b/examples/examples_multi_gpu.py
deleted file mode 100644
index 3ecd88b..0000000
--- a/examples/examples_multi_gpu.py
+++ /dev/null
@@ -1,178 +0,0 @@
-"""CCNN with the DEAP Dataset
-======================================
-In this case, we introduce how to use TorchEEG to train a Continuous Convolutional Neural Network (CCNN) on the DEAP dataset for emotion classification.
-"""
-
-import logging
-import os
-import random
-import time
-
-import numpy as np
-import torch
-from torch.utils.data.dataloader import DataLoader
-from torcheeg import transforms
-from torcheeg.datasets import DEAPDataset
-from torcheeg.datasets.constants.emotion_recognition.deap import \
-    DEAP_CHANNEL_LOCATION_DICT
-from torcheeg.model_selection import KFoldGroupbyTrial
-from torcheeg.models import CCNN
-from torcheeg.trainers import ClassificationTrainer
-
-###############################################################################
-# Pre-experiment Preparation to Ensure Reproducibility
-# -----------------------------------------
-# Use the logging module to store output in a log file for easy reference while printing it to the screen.
-
-os.makedirs('./tmp_out/examples_multi_gpu/log', exist_ok=True)
-logger = logging.getLogger('CCNN with the DEAP Dataset')
-logger.setLevel(logging.DEBUG)
-console_handler = logging.StreamHandler()
-timeticks = time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime())
-file_handler = logging.FileHandler(
-    os.path.join('./tmp_out/examples_multi_gpu/log', f'{timeticks}.log'))
-logger.addHandler(console_handler)
-logger.addHandler(file_handler)
-
-###############################################################################
-# Set the random number seed in all modules to guarantee the same result when running again.
-
-
-def seed_everything(seed):
-    random.seed(seed)
-    np.random.seed(seed)
-    os.environ["PYTHONHASHSEED"] = str(seed)
-    torch.manual_seed(seed)
-    torch.cuda.manual_seed(seed)
-    torch.backends.cudnn.deterministic = True
-    torch.backends.cudnn.benchmark = False
-
-
-seed_everything(42)
-
-###############################################################################
-# Customize Trainer
-# -----------------------------------------
-# TorchEEG provides a large number of trainers to help complete the training of classification models, generative models and cross-domain methods. Here we choose the simplest classification trainer, inherit the trainer and overload the log function to save the log using our own defined method; other hook functions can also be overloaded to meet special needs.
-#
-
-
-class MyClassificationTrainer(ClassificationTrainer):
-    def log(self, *args, **kwargs):
-        if self.is_main:
-            logger.info(*args, **kwargs)
-
-
-######################################################################
-# Building Deep Learning Pipelines Using TorchEEG
-# -----------------------------------------
-# Step 1: Initialize the Dataset
-#
-# We use the DEAP dataset supported by TorchEEG. Here, we set an EEG sample to 1 second long and include 128 data points. The baseline signal is 3 seconds long, cut into three, and averaged as the baseline signal for the trial. In offline preprocessing, we divide the EEG signal of every electrode into 4 sub-bands, and calculate the differential entropy on each sub-band as a feature, followed by debaselining and mapping on the grid. Finally, the preprocessed EEG signals are stored in the local IO. In online processing, all EEG signals are converted into Tensors for input into neural networks.
-#
-
-dataset = DEAPDataset(io_path=f'./tmp_out/examples_multi_gpu/deap',
-                      root_path='./tmp_in/data_preprocessed_python',
-                      offline_transform=transforms.Compose([
-                          transforms.BandDifferentialEntropy(
-                              sampling_rate=128, apply_to_baseline=True),
-                          transforms.BaselineRemoval(),
-                          transforms.ToGrid(DEAP_CHANNEL_LOCATION_DICT)
-                      ]),
-                      online_transform=transforms.ToTensor(),
-                      label_transform=transforms.Compose([
-                          transforms.Select('valence'),
-                          transforms.Binary(5.0),
-                      ]),
-                      chunk_size=128,
-                      baseline_chunk_size=128,
-                      num_baseline=3,
-                      num_worker=4)
-
-######################################################################
-# .. warning::
-#    If you use TorchEEG under the `Windows` system and want to use multiple processes (such as in dataset or dataloader), you should check whether :obj:`__name__` is :obj:`__main__` to avoid errors caused by multiple :obj:`import`.
-#
-# That is, under the :obj:`Windows` system, you need to:
-#  .. code-block::
-#
-#    if __name__ == "__main__":
-#        dataset = DEAPDataset(
-#             io_path=
-#             f'./tmp_out/examples_multi_gpu/deap',
-#             root_path='./tmp_in/data_preprocessed_python',
-#             offline_transform=transforms.Compose([
-#                 transforms.BandDifferentialEntropy(sampling_rate=128,
-#                                                    apply_to_baseline=True),
-#                 transforms.BaselineRemoval(),
-#                 transforms.ToGrid(DEAP_CHANNEL_LOCATION_DICT)
-#             ]),
-#             online_transform=transforms.ToTensor(),
-#             label_transform=transforms.Compose([
-#                 transforms.Select('valence'),
-#                 transforms.Binary(5.0),
-#             ]),
-#             io_mode='pickle',
-#             chunk_size=128,
-#             baseline_chunk_size=128,
-#             num_baseline=3,
-#             num_worker=4)
-#        # the following codes
-#
-# .. note::
-#    LMDB may not be optimized for parts of Windows systems or storage devices. If you find that the data preprocessing speed is slow, you can consider setting :obj:`io_mode` to :obj:`pickle`, which is an alternative implemented by TorchEEG based on pickle.
-
-######################################################################
-# Step 2: Divide the Training and Test samples in the Dataset
-#
-# Here, the dataset is divided using 5-fold cross-validation. In the process of division, we group according to the trial index, and every trial takes 4 folds as training samples and 1 fold as test samples. Samples across trials are aggregated to obtain training set and test set.
-#
-
-k_fold = KFoldGroupbyTrial(n_splits=5,
-                           split_path='./tmp_out/examples_multi_gpu/split')
-
-######################################################################
-# Step 3: Define the Model and Start Training
-#
-# We first use a loop to get the dataset in each cross-validation. In each cross-validation, we initialize the CCNN model and define the hyperparameters. For example, each EEG sample contains 4-channel features from 4 sub-bands, the grid size is 9 times 9, etc.
-#
-# We then initialize the trainer and set the hyperparameters in the trained model, such as the learning rate, the equipment used, etc. The :obj:`fit` method receives the training dataset and starts training the model. The :obj:`test` method receives a test dataset and reports the test results. The :obj:`save_state_dict` method can save the trained model.
-
-for i, (train_dataset, val_dataset) in enumerate(k_fold.split(dataset)):
-    # Initialize the model
-    model = CCNN(in_channels=4, grid_size=(9, 9), num_classes=2, dropout=0.5)
-
-    # Initialize the trainer and use the 0, 1, 2-th GPUs for training
-    trainer = MyClassificationTrainer(model=model,
-                                      lr=1e-4,
-                                      weight_decay=1e-4,
-                                      device_ids=[0, 1, 2])
-
-    # Initialize several batches of training samples and test samples
-    train_loader = DataLoader(train_dataset,
-                              batch_size=256,
-                              shuffle=True,
-                              num_workers=4)
-    val_loader = DataLoader(val_dataset,
-                            batch_size=256,
-                            shuffle=False,
-                            num_workers=4)
-
-    # Do 50 rounds of training
-    trainer.fit(train_loader, val_loader, num_epochs=50)
-    trainer.test(val_loader)
-    trainer.save_state_dict(f'./tmp_out/examples_multi_gpu/weight/{i}.pth')
-
-######################################################################
-# Start Multi-GPU Training
-# -----------------------------------------
-# Instead of running python :obj:`examples/examples_multi_gpu_trainer.py` directly from the command line, you need use :obj:`torch.distributed.launch` to start multi-GPU training:
-# .. code-block::
-#
-#    python -m torch.distributed.launch \
-#        --nproc_per_node=3 \
-#        --nnodes=1 \
-#        --node_rank=0 \
-#        --master_addr="localhost" \
-#        --master_port=2345 \
-#        examples/examples_multi_gpu_trainer.py
\ No newline at end of file
diff --git a/examples/examples_pipeline.py b/examples/examples_pipeline.py
new file mode 100644
index 0000000..08b52c3
--- /dev/null
+++ b/examples/examples_pipeline.py
@@ -0,0 +1,113 @@
+"""
+Using TorchEEG to Complete a Deep Learning Workflow
+===================================================
+
+In this tutorial, we demonstrate a complete deep learning workflow using TorchEEG. We will cover the following aspects:
+
+1. Utilizing Datasets and Transformers in TorchEEG
+2. Data Partitioning Strategies in TorchEEG
+3. Leveraging Models and Trainers in TorchEEG
+
+"""
+
+
+######################################################################
+# Step 1: Initialize the Dataset
+# ------------------------------
+# 
+# We use the DEAP dataset supported by TorchEEG. Each EEG sample is set to
+# be 1 second long, encompassing 128 data points. The baseline signal is 3
+# seconds long, which we divide into three sections and then average to
+# obtain the trial’s baseline signal.
+# 
+# During offline preprocessing, we divide each electrode’s EEG signal into
+# 4 sub-bands, calculate the differential entropy for each sub-band as a
+# feature, perform debaselining, and map onto a grid. Finally, the
+# preprocessed EEG signals are saved locally. For online processing, we
+# convert all EEG signals into Tensors, making them suitable for neural
+# network input.
+# 
+
+from torcheeg.datasets import DEAPDataset
+from torcheeg import transforms
+
+from torcheeg.datasets.constants.emotion_recognition.deap import \
+    DEAP_CHANNEL_LOCATION_DICT
+
+dataset = DEAPDataset(
+    io_path=f'./examples_pipeline/deap',
+    root_path='./data_preprocessed_python',
+    offline_transform=transforms.Compose([
+        transforms.BandDifferentialEntropy(apply_to_baseline=True),
+        transforms.ToGrid(DEAP_CHANNEL_LOCATION_DICT, apply_to_baseline=True)
+    ]),
+    online_transform=transforms.Compose(
+        [transforms.BaselineRemoval(),
+         transforms.ToTensor()]),
+    label_transform=transforms.Compose([
+        transforms.Select('valence'),
+        transforms.Binary(5.0),
+    ]),
+    num_worker=8)
+
+
+######################################################################
+# Step 2: Split the Dataset into Training and Test Sets
+# -----------------------------------------------------
+# 
+# In this case, we use per-subject 5-fold cross-validation to split the
+# dataset. During this process, we separate each subject’s EEG samples
+# into training and test sets. We use 4 folds for training and 1 fold for
+# testing.
+# 
+
+from torcheeg.datasets import KFoldGroupbyTrial
+
+k_fold = KFoldGroupbyTrial(n_splits=10,
+                           split_path='./examples_pipeline/split',
+                           shuffle=True,
+                           random_state=42)
+
+
+######################################################################
+# Step 3: Define the Model and Initiate Training
+# ----------------------------------------------
+# 
+# We loop through each cross-validation set, and for each one, we
+# initialize the CCNN model and define its hyperparameters. For instance,
+# each EEG sample contains 4-channel features from 4 sub-bands, and the
+# grid size is 9x9.
+# 
+# We then train the model for 50 epochs using the ``ClassifierTrainer``.
+# 
+
+from torch.utils.data import DataLoader
+from torcheeg.models import CCNN
+
+from torcheeg.trainers import ClassifierTrainer
+
+import pytorch_lightning as pl
+
+for i, (train_dataset, val_dataset) in enumerate(k_fold.split(dataset)):
+    train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
+    val_loader = DataLoader(val_dataset, batch_size=64, shuffle=False)
+
+    model = CCNN(num_classes=2, in_channels=4, grid_size=(9, 9))
+
+    trainer = ClassifierTrainer(model=model,
+                                num_classes=2,
+                                lr=1e-4,
+                                weight_decay=1e-4,
+                                accelerator="gpu")
+    trainer.fit(train_loader,
+                val_loader,
+                max_epochs=50,
+                default_root_dir=f'./examples_pipeline/model/{i}',
+                callbacks=[pl.callbacks.ModelCheckpoint(save_last=True)],
+                enable_progress_bar=True,
+                enable_model_summary=True,
+                limit_val_batches=0.0)
+    score = trainer.test(val_loader,
+                         enable_progress_bar=True,
+                         enable_model_summary=True)[0]
+    print(f'Fold {i} test accuracy: {score["test_accuracy"]:.4f}')
\ No newline at end of file
diff --git a/examples/examples_seed_ccnn.py b/examples/examples_seed_ccnn.py
index 1c62448..82341e9 100644
--- a/examples/examples_seed_ccnn.py
+++ b/examples/examples_seed_ccnn.py
@@ -1,140 +1,92 @@
-"""CCNN with the SEED Dataset
-======================================
-In this case, we introduce how to use TorchEEG to train a Continuous Convolutional Neural Network (CCNN) on the SEED dataset for emotion classification.
 """
+Train a CCNN Model on the SEED Dataset
+==========================
 
-import logging
-import os
-import random
-import time
-
-import numpy as np
-import torch
-from torch.utils.data.dataloader import DataLoader
-from torcheeg import transforms
-from torcheeg.datasets import SEEDDataset
-from torcheeg.datasets.constants.emotion_recognition.seed import \
-    SEED_CHANNEL_LOCATION_DICT
-from torcheeg.model_selection import KFoldGroupbyTrial
-from torcheeg.models import CCNN
-from torcheeg.trainers import ClassificationTrainer
-
-###############################################################################
-# Pre-experiment Preparation to Ensure Reproducibility
-# -----------------------------------------
-# Use the logging module to store output in a log file for easy reference while printing it to the screen.
-
-os.makedirs('./tmp_out/examples_seed_ccnn/log', exist_ok=True)
-logger = logging.getLogger('CCNN with the SEED Dataset')
-logger.setLevel(logging.DEBUG)
-console_handler = logging.StreamHandler()
-timeticks = time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime())
-file_handler = logging.FileHandler(os.path.join('./tmp_out/examples_seed_ccnn/log', f'{timeticks}.log'))
-logger.addHandler(console_handler)
-logger.addHandler(file_handler)
-
-###############################################################################
-# Set the random number seed in all modules to guarantee the same result when running again.
-
-
-def seed_everything(seed):
-    random.seed(seed)
-    np.random.seed(seed)
-    os.environ["PYTHONHASHSEED"] = str(seed)
-    torch.manual_seed(seed)
-    torch.cuda.manual_seed(seed)
-    torch.backends.cudnn.deterministic = True
-    torch.backends.cudnn.benchmark = False
-
-
-seed_everything(42)
-
-###############################################################################
-# Customize Trainer
-# -----------------------------------------
-# TorchEEG provides a large number of trainers to help complete the training of classification models, generative models and cross-domain methods. Here we choose the simplest classification trainer, inherit the trainer and overload the log function to save the log using our own defined method; other hook functions can also be overloaded to meet special needs.
-#
-
-
-class MyClassificationTrainer(ClassificationTrainer):
-    def log(self, *args, **kwargs):
-        if self.is_main:
-            logger.info(*args, **kwargs)
+In this tutorial, we'll demonstrate how you can leverage TorchEEG to train a Continuous Convolutional Neural Network (CCNN) on the SEED dataset for emotion classification. Let's navigate through this process step-by-step, covering everything from dataset initialization to model evaluation.
 
+"""
 
 ######################################################################
-# Building Deep Learning Pipelines Using TorchEEG
-# -----------------------------------------
 # Step 1: Initialize the Dataset
+# ----------------------------------------------
 #
-# We use the SEED dataset supported by TorchEEG. Here, we set an EEG sample to 1 second long and include 200 data points. In offline preprocessing, we divide the EEG signal of every electrode into 4 sub-bands, and calculate the differential entropy on each sub-band as a feature, followed by debaselining and mapping on the grid. Finally, the preprocessed EEG signals are stored in the local IO. In online processing, all EEG signals are converted into Tensors for input into neural networks.
-#
+# To begin with, we employ the SEED dataset provided by TorchEEG. Each EEG sample in this dataset lasts for 1 second and contains 200 data points.
+# In terms of offline preprocessing, we perform the following tasks:
+# - Divide each electrode's EEG signal into 4 frequency sub-bands
+# - Compute the differential entropy of each sub-band to serve as a feature
+# - Remove the baseline from the signal
+# - Map the preprocessed signals onto a 2D grid
+# For online processing, we convert the EEG signals into PyTorch Tensors, ensuring they are compatible with neural network inputs.
+# 
+
+from torcheeg.datasets import SEEDDataset
+from torcheeg import transforms
+from torcheeg.datasets.constants.emotion_recognition.seed import SEED_CHANNEL_LOCATION_DICT
 
-dataset = SEEDDataset(io_path='./tmp_out/examples_seed_ccnn/seed',
-                      root_path='./tmp_in/Preprocessed_EEG',
+dataset = SEEDDataset(io_path=f'./examples_seed_ccnn/seed',
+                      root_path='./Preprocessed_EEG',
                       offline_transform=transforms.Compose([
                           transforms.BandDifferentialEntropy(sampling_rate=200),
                           transforms.ToGrid(SEED_CHANNEL_LOCATION_DICT)
                       ]),
                       online_transform=transforms.ToTensor(),
-                      label_transform=transforms.Compose(
-                          [transforms.Select('emotion'),
-                           transforms.Lambda(lambda x: x + 1)]),
+                      label_transform=transforms.Compose([
+                          transforms.Select('emotion'),
+                          transforms.Lambda(lambda x: x + 1)
+                      ]),
                       chunk_size=200,
-                      num_worker=16)
+                      num_worker=4)
 
 ######################################################################
-# .. warning::
-#    If you use TorchEEG under the `Windows` system and want to use multiple processes (such as in dataset or dataloader), you should check whether :obj:`__name__` is :obj:`__main__` to avoid errors caused by multiple :obj:`import`.
+# Step 2: Divide the Training and Test Samples in the Dataset
+# ----------------------------------------------
 #
-# That is, under the :obj:`Windows` system, you need to:
-#  .. code-block::
+# Next, let's partition our dataset into training and test sets using 5-fold cross-validation. We group the data based on their trial index, where each trial contributes 4 folds to the training set and 1 fold to the test set. These grouped samples are then combined to form the final training and test sets.
 #
-#    if __name__ == "__main__":
-#        dataset = SEEDDataset(io_path='./tmp_out/examples_seed_ccnn/seed',
-#                      root_path='./tmp_in/data_preprocessed_python',
-#                      offline_transform=transforms.Compose([
-#                          transforms.BandDifferentialEntropy(sampling_rate=200),
-#                          transforms.ToGrid(SEED_CHANNEL_LOCATION_DICT)
-#                      ]),
-#                      online_transform=transforms.ToTensor(),
-#                      label_transform=transforms.Compose(
-#                          [transforms.Select('emotion'),
-#                           transforms.Lambda(lambda x: x + 1)]),
-#                      chunk_size=200,
-#                      num_worker=4)
-#        # the following codes
-#
-# .. warning::
-#    LMDB may not be optimized for parts of Windows systems or storage devices. If you find that the data preprocessing speed is slow, you can consider setting :obj:`io_mode` to :obj:`pickle`, which is an alternative implemented by TorchEEG based on pickle.
 
-######################################################################
-# Step 2: Divide the Training and Test samples in the Dataset
-#
-# Here, the dataset is divided using 5-fold cross-validation. In the process of division, we group according to the trial index, and every trial takes 4 folds as training samples and 1 fold as test samples. Samples across trials are aggregated to obtain training set and test set.
-#
+from torcheeg.model_selection import KFoldGroupbyTrial
 
-k_fold = KFoldGroupbyTrial(n_splits=5, split_path='./tmp_out/examples_seed_ccnn/split')
+k_fold = KFoldGroupbyTrial(n_splits=5, split_path=f'./examples_seed_ccnn/split')
 
 ######################################################################
-# Step 3: Define the Model and Start Training
+# Step 3: Define the Model and Initiate Training
+# ----------------------------------------------
+#
+# We loop through each cross-validation set, and for each one, we
+# initialize the CCNN model and define its hyperparameters. For instance,
+# each EEG sample contains 4-channel features from 4 sub-bands, and the
+# grid size is 9x9.
 #
-# We first use a loop to get the dataset in each cross-validation. In each cross-validation, we initialize the CCNN model and define the hyperparameters. For example, each EEG sample contains 4-channel features from 4 sub-bands, the grid size is 9 times 9, etc.
+# We then train the model for 50 epochs using the ``ClassifierTrainer``.
 #
-# We then initialize the trainer and set the hyperparameters in the trained model, such as the learning rate, the equipment used, etc. The :obj:`fit` method receives the training dataset and starts training the model. The :obj:`test` method receives a test dataset and reports the test results. The :obj:`save_state_dict` method can save the trained model.
 
-for i, (train_dataset, val_dataset) in enumerate(k_fold.split(dataset)):
-    # Initialize the model
-    model = CCNN(in_channels=4, grid_size=(9, 9), num_classes=3, dropout=0.5)
+from torch.utils.data import DataLoader
+from torcheeg.models import CCNN
 
-    # Initialize the trainer and use the 0-th GPU for training
-    trainer = MyClassificationTrainer(model=model, lr=1e-4, weight_decay=1e-4, device_ids=[0])
+from torcheeg.trainers import ClassifierTrainer
 
-    # Initialize several batches of training samples and test samples
-    train_loader = DataLoader(train_dataset, batch_size=256, shuffle=True, num_workers=4)
-    val_loader = DataLoader(val_dataset, batch_size=256, shuffle=False, num_workers=4)
+import pytorch_lightning as pl
 
-    # Do 50 rounds of training
-    trainer.fit(train_loader, val_loader, num_epochs=50)
-    trainer.test(val_loader)
-    trainer.save_state_dict(f'./tmp_out/examples_seed_ccnn/weight/{i}.pth')
+for i, (train_dataset, val_dataset) in enumerate(k_fold.split(dataset)):
+    train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
+    val_loader = DataLoader(val_dataset, batch_size=64, shuffle=False)
+
+    model = CCNN(num_classes=3, in_channels=4, grid_size=(9, 9))
+
+    trainer = ClassifierTrainer(model=model,
+                                num_classes=2,
+                                lr=1e-4,
+                                weight_decay=1e-4,
+                                accelerator="gpu")
+    trainer.fit(train_loader,
+                val_loader,
+                max_epochs=50,
+                default_root_dir=f'./examples_seed_ccnn/model/{i}',
+                callbacks=[pl.callbacks.ModelCheckpoint(save_last=True)],
+                enable_progress_bar=True,
+                enable_model_summary=True,
+                limit_val_batches=0.0)
+    score = trainer.test(val_loader,
+                         enable_progress_bar=True,
+                         enable_model_summary=True)[0]
+    print(f'Fold {i} test accuracy: {score["test_accuracy"]:.4f}')
\ No newline at end of file
diff --git a/examples/examples_seed_domain_adaption.py b/examples/examples_seed_domain_adaption.py
new file mode 100644
index 0000000..0fb7214
--- /dev/null
+++ b/examples/examples_seed_domain_adaption.py
@@ -0,0 +1,145 @@
+"""
+Apply the Domain Adaption Algorithm on the SEED Dataset
+===================================================
+
+In this tutorial, you'll learn how to utilize TorchEEG's Associative Domain Adaptation (ADA) algorithm to address the cross-subject emotion recognition challenge using the SEED dataset.
+
+"""
+
+######################################################################
+# Step 1: Initialize the Dataset
+# ------------------------------
+#
+# We leverage the SEED dataset, supported by TorchEEG, where each EEG sample lasts 1 second and consists of 200 data points. 
+#
+# In offline preprocessing, each electrode's EEG signal is segmented into 4 sub-bands. We then calculate the differential entropy for each sub-band as a feature. The signals are debaselined, mapped onto a grid, and finally saved locally. 
+#
+# For online processing, we convert all EEG signals into Tensors for neural network input. Additionally, we apply before_trial_normalize to reduce trial variance effects on cross-subject emotion recognition.
+#
+
+from torcheeg.datasets import SEEDFeatureDataset
+from torcheeg import transforms
+from torcheeg.datasets.constants.emotion_recognition.seed import SEED_CHANNEL_LOCATION_DICT
+from torcheeg.datasets.functional import before_trial_normalize
+from pytorch_lightning.callbacks import ModelCheckpoint
+import numpy as np
+
+dataset = SEEDFeatureDataset(
+    io_path=f'./examples_seed_domain_adaption/seed',
+    root_path='./ExtractedFeatures',
+    after_session=before_trial_normalize,
+    offline_transform=transforms.Compose(
+        [transforms.ToGrid(SEED_CHANNEL_LOCATION_DICT)]),
+    online_transform=transforms.ToTensor(),  # seed do not have baseline signals
+    label_transform=transforms.Compose(
+        [transforms.Select('emotion'),
+         transforms.Lambda(lambda x: x + 1)]),
+    feature=['de_LDS'],
+    num_worker=4)
+
+######################################################################
+# Step 2: Define the Model
+# ----------------------------------------------
+#
+# We need to split the CCNN model into two parts: the feature extractor, and
+# the classifier. The feature extractor is trained to extract domain-invariant
+# features, while the classifier is trained to classify the emotion.
+#
+
+from torcheeg.models import CCNN
+
+
+class Extractor(CCNN):
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.conv2(x)
+        x = self.conv3(x)
+        x = self.conv4(x)
+        x = x.flatten(start_dim=1)
+        return x
+
+
+class Classifier(CCNN):
+    def forward(self, x):
+        x = self.lin1(x)
+        x = self.lin2(x)
+        return x
+
+
+######################################################################
+# Step 3: Split the Dataset and Train the Model
+# -----------------------------------------------------
+#
+# In this case, we do not consider the cross-session problem. We use the
+# subset of each session to train the model, and use the subset of the
+# same session to test the model. We use the leave-one-subject-out method for evaluation. Each subject is iteratively left out from the training set and used as the test set.
+#
+
+from torcheeg.model_selection import Subcategory
+
+subset = Subcategory(
+    criteria='session_id',
+    split_path=f'./examples_seed_domain_adaption/split/session')
+
+
+######################################################################
+# Step 4: Train the Model
+# ----------------------------------------------
+#
+# For training, we make use of the ADA algorithm to adapt domain-invariant features across subjects.
+#
+
+from torch.utils.data import DataLoader
+from torcheeg.model_selection import LeaveOneSubjectOut
+from torcheeg.trainers import ADATrainer
+
+scores = []
+for j, sub_dataset in enumerate(subset.split(dataset)):
+    loo = LeaveOneSubjectOut(
+        split_path=f'./examples_seed_domain_adaption/split/loo_{j}')
+    for i, (train_dataset, test_dataset) in enumerate(loo.split(sub_dataset)):
+
+        extractor = Extractor(in_channels=5, num_classes=3)
+        classifier = Classifier(in_channels=5, num_classes=3)
+
+        source_loader = DataLoader(train_dataset,
+                                   batch_size=128,
+                                   shuffle=True,
+                                   num_workers=4,
+                                   drop_last=True)
+
+        target_loader = DataLoader(test_dataset,
+                                   batch_size=128,
+                                   shuffle=True,
+                                   num_workers=4,
+                                   drop_last=True)
+
+        test_loader = DataLoader(test_dataset,
+                                 batch_size=128,
+                                 shuffle=True,
+                                 num_workers=4)
+
+        trainer = ADATrainer(extractor=extractor,
+                             classifier=classifier,
+                             num_classes=2,
+                             lr=1e-4,
+                             weight_decay=0.0,
+                             accelerator='gpu')
+
+        trainer.fit(
+            source_loader,
+            target_loader,
+            test_loader,
+            max_epochs=50,
+            default_root_dir=
+            f'./examples_seed_domain_adaption/model/ses_{i}_sub_{j}.pth',
+            callbacks=[ModelCheckpoint(save_last=True)],
+            enable_progress_bar=True,
+            enable_model_summary=True,
+            limit_val_batches=0.0)
+        score = trainer.test(test_loader,
+                             enable_progress_bar=True,
+                             enable_model_summary=True)[0]
+        scores.append(score['test_accuracy'])
+
+    print(f'final accuracy: {np.array(scores).mean():>0.1f}%')
\ No newline at end of file
diff --git a/examples/examples_seed_feature_ada.py b/examples/examples_seed_feature_ada.py
deleted file mode 100644
index 2535ece..0000000
--- a/examples/examples_seed_feature_ada.py
+++ /dev/null
@@ -1,203 +0,0 @@
-"""Cross-domain Emotion Recognition with ADA
-======================================
-In this case, we introduce how to use TorchEEG and a simple CCNN model to complete emotion recognition across subjects. Here, the EEG signals of different subjects have a distribution gap, so the model trained on some subjects has a performance drop on unknown subjects. We use a cross-domain algorithm, Associative Domain Adaptation (ADA), to solve this problem.
-"""
-
-import logging
-import os
-import random
-import time
-
-import numpy as np
-import pandas as pd
-import torch
-from torch.utils.data.dataloader import DataLoader
-
-from torcheeg import transforms
-from torcheeg.datasets import SEEDFeatureDataset
-from torcheeg.datasets.constants.emotion_recognition.seed import \
-    SEED_CHANNEL_LOCATION_DICT
-from torcheeg.model_selection import LeaveOneSubjectOut, Subcategory
-from torcheeg.models import CCNN
-from torcheeg.trainers import ADATrainer
-
-###############################################################################
-# Pre-experiment Preparation to Ensure Reproducibility
-# -----------------------------------------
-# Use the logging module to store output in a log file for easy reference while printing it to the screen.
-
-os.makedirs('./tmp_out/examples_seed_feature_ada/log', exist_ok=True)
-logger = logging.getLogger('Cross-domain Emotion Recognition with DAD')
-logger.setLevel(logging.DEBUG)
-console_handler = logging.StreamHandler()
-timeticks = time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime())
-file_handler = logging.FileHandler(
-    os.path.join('./tmp_out/examples_seed_feature_ada/log', f'{timeticks}.log'))
-logger.addHandler(console_handler)
-logger.addHandler(file_handler)
-
-###############################################################################
-# Set the random number seed in all modules to guarantee the same result when running again.
-
-
-def seed_everything(seed):
-    random.seed(seed)
-    np.random.seed(seed)
-    os.environ["PYTHONHASHSEED"] = str(seed)
-    torch.manual_seed(seed)
-    torch.cuda.manual_seed(seed)
-    torch.backends.cudnn.deterministic = True
-    torch.backends.cudnn.benchmark = False
-
-
-seed_everything(42)
-
-###############################################################################
-# Define feature extractor and classifier
-# -----------------------------------------
-# Different from the vanilla classification model, cross-domain algorithms usually need to constrain the features extracted by the model (such as domain invariant characteristics, etc.), so we need to split the classification model into feature extraction part and classification part.
-#
-
-
-class Extractor(CCNN):
-    def forward(self, x):
-        x = self.conv1(x)
-        x = self.conv2(x)
-        x = self.conv3(x)
-        x = self.conv4(x)
-        x = x.flatten(start_dim=1)
-        return x
-
-
-class Classifier(CCNN):
-    def forward(self, x):
-        x = self.lin1(x)
-        x = self.lin2(x)
-        return x
-
-
-###############################################################################
-# Customize Trainer
-# -----------------------------------------
-# TorchEEG provides a large number of domain adaption trainers to help complete the training of cross-subject/session emotion classification models. Here we choose the associative domain adaptation trainer, inherit the trainer and overload the log function to save the log using our own defined method; other hook functions can also be overloaded to meet special needs.
-#
-
-
-class MyClassificationTrainer(ADATrainer):
-    def log(self, *args, **kwargs):
-        if self.is_main:
-            logger.info(*args, **kwargs)
-
-
-######################################################################
-# Building Deep Learning Pipelines Using TorchEEG
-# -----------------------------------------
-# Step 1: Initialize the Dataset
-#
-# We use the SEED dataset supported by TorchEEG. Here we are using extracted EEG features. In the feature dataset, EEG signals (200 data points) per second are pre-computed with differential entropy in five sub-bands and smoothed using a linear dynamical system. We then normalize it and map it onto the grid. Finally, EEG signals are converted into Tensors for input into neural networks.
-#
-
-dataset = SEEDFeatureDataset(io_path='./tmp_out/examples_seed_feature_ada/seed',
-                             root_path='./tmp_in/ExtractedFeatures',
-                             feature=['de_movingAve'],
-                             online_transform=transforms.Compose([
-                                 transforms.MinMaxNormalize(axis=-1),
-                                 transforms.ToGrid(SEED_CHANNEL_LOCATION_DICT),
-                                 transforms.ToTensor()
-                             ]),
-                             label_transform=transforms.Compose([
-                                 transforms.Select('emotion'),
-                                 transforms.Lambda(lambda x: int(x) + 1),
-                             ]),
-                             num_worker=8)
-
-######################################################################
-# .. warning::
-#    If you use TorchEEG under the `Windows` system and want to use multiple processes (such as in dataset or dataloader), you should check whether :obj:`__name__` is :obj:`__main__` to avoid errors caused by multiple :obj:`import`.
-#
-# That is, under the :obj:`Windows` system, you need to:
-#  .. code-block::
-#
-#    if __name__ == "__main__":
-#        dataset = SEEDFeatureDataset(io_path='./tmp_out/examples_seed_feature_ada/seed',
-#                              root_path='./tmp_in/ExtractedFeatures',
-#                              feature=['de_movingAve'],
-#                              online_transform=transforms.Compose([
-#                                  transforms.MinMaxNormalize(axis=-1),
-#                                  transforms.ToGrid(SEED_CHANNEL_LOCATION_DICT),
-#                                  transforms.ToTensor()
-#                              ]),
-#                              label_transform=transforms.Compose([
-#                                  transforms.Select('emotion'),
-#                                  transforms.Lambda(lambda x: int(x) + 1),
-#                              ]),
-#                              io_mode='pickle',
-#                              num_worker=8)
-#        # the following codes
-#
-# .. note::
-#    LMDB may not be optimized for parts of Windows systems or storage devices. If you find that the data preprocessing speed is slow, you can consider setting :obj:`io_mode` to :obj:`pickle`, which is an alternative implemented by TorchEEG based on pickle.
-
-######################################################################
-# Step 2: Divide the Training and Test samples in the Dataset
-#
-# Here we do not consider the impact of cross-session on the test results. Therefore, we first mark the session index on the sample according to the collection date. Next, we use :obj:`Subcategory` to divide the data set to obtain the sub-data set of the first session, the second session and the third session.
-#
-
-subject_info_list = []
-for subject_id in dataset.info['subject_id'].unique().tolist():
-    subject_info = dataset.info[dataset.info['subject_id'] == subject_id]
-    session_id_set = subject_info['date'].unique().tolist()
-    session_id_map = {
-        session_id: i
-        for i, session_id in enumerate(session_id_set)
-    }
-    subject_info['session_id'] = subject_info['date'].apply(
-        lambda x: session_id_map[x])
-    subject_info_list.append(subject_info)
-dataset.info = pd.concat(subject_info_list)
-
-subset = Subcategory(criteria='session_id',
-                     split_path='./tmp_out/examples_seed_feature_ada/split')
-
-######################################################################
-# Step 3: Define the Model and Start Training
-#
-# For the dataset of each session, we use the leave-one-out method to conduct experiments, and use one subject as the test dataset and the other subjects as the training dataset to train the model.
-#
-# We define the training set as the source domain and the test set as the target domain, hoping that the model trained on the source domain can be generalized to the target domain. We then initialize the trainer and set the hyperparameters in the trained model, such as the learning rate, the equipment used, etc. The :obj:`fit` method receives the source domain data samples and the target domain samples (without labels, as they are from test dataset) and starts training the model. The :obj:`test` method receives a test dataset and reports the test results. The :obj:`save_state_dict` method can save the trained model.
-
-for j, sub_dataset in enumerate(subset.split(dataset)):
-    k_fold = LeaveOneSubjectOut(
-        split_path=f'./tmp_out/examples_seed_feature_ada/split_{j}')
-    for i, (train_dataset,
-            test_dataset) in enumerate(k_fold.split(sub_dataset)):
-        extractor = Extractor(in_channels=5, num_classes=3)
-        classifier = Classifier(in_channels=5, num_classes=3)
-        trainer = MyClassificationTrainer(extractor=extractor,
-                                          classifier=classifier,
-                                          lr=1e-4,
-                                          lambd=1.0,
-                                          weight_decay=0.0,
-                                          device_ids=[0])
-
-        source_loader = DataLoader(train_dataset,
-                                   batch_size=128,
-                                   shuffle=True,
-                                   num_workers=4,
-                                   drop_last=True)
-
-        target_loader = DataLoader(test_dataset,
-                                   batch_size=128,
-                                   shuffle=True,
-                                   num_workers=4,
-                                   drop_last=True)
-
-        test_loader = DataLoader(test_dataset,
-                                 batch_size=128,
-                                 shuffle=True,
-                                 num_workers=4)
-        trainer.fit(source_loader, target_loader, test_loader, num_epochs=50)
-        trainer.test(test_loader)
-        trainer.save_state_dict(
-            f'./tmp_out/examples_seed_feature_ada/weight/{j}-{i}.pth')
diff --git a/examples/examples_seed_tsception.py b/examples/examples_seed_tsception.py
deleted file mode 100644
index ea7fcfb..0000000
--- a/examples/examples_seed_tsception.py
+++ /dev/null
@@ -1,143 +0,0 @@
-"""TSCeption with the SEED Dataset
-======================================
-In this case, we introduce how to use TorchEEG to train a TSCeption on the SEED dataset for emotion classification.
-"""
-
-import logging
-import os
-import random
-import time
-
-import numpy as np
-import torch
-from torch.utils.data.dataloader import DataLoader
-from torcheeg import transforms
-from torcheeg.datasets import SEEDDataset
-from torcheeg.datasets.constants.emotion_recognition.seed import SEED_CHANNEL_LOCATION_DICT
-from torcheeg.model_selection import KFoldGroupbyTrial
-from torcheeg.models import TSCeption
-from torcheeg.trainers import ClassificationTrainer
-
-###############################################################################
-# Pre-experiment Preparation to Ensure Reproducibility
-# -----------------------------------------
-# Use the logging module to store output in a log file for easy reference while printing it to the screen.
-
-os.makedirs('./tmp_out/examples_seed_tsception/log', exist_ok=True)
-logger = logging.getLogger('TSCeption with the SEED Dataset')
-logger.setLevel(logging.DEBUG)
-console_handler = logging.StreamHandler()
-timeticks = time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime())
-file_handler = logging.FileHandler(os.path.join('./tmp_out/examples_seed_tsception/log', f'{timeticks}.log'))
-logger.addHandler(console_handler)
-logger.addHandler(file_handler)
-
-###############################################################################
-# Set the random number seed in all modules to guarantee the same result when running again.
-
-
-def seed_everything(seed):
-    random.seed(seed)
-    np.random.seed(seed)
-    os.environ["PYTHONHASHSEED"] = str(seed)
-    torch.manual_seed(seed)
-    torch.cuda.manual_seed(seed)
-    torch.backends.cudnn.deterministic = True
-    torch.backends.cudnn.benchmark = False
-
-
-seed_everything(42)
-
-###############################################################################
-# Customize Trainer
-# -----------------------------------------
-# TorchEEG provides a large number of trainers to help complete the training of classification models, generative models and cross-domain methods. Here we choose the simplest classification trainer, inherit the trainer and overload the log function to save the log using our own defined method; other hook functions can also be overloaded to meet special needs.
-#
-
-
-class MyClassificationTrainer(ClassificationTrainer):
-    def log(self, *args, **kwargs):
-        if self.is_main:
-            logger.info(*args, **kwargs)
-
-
-######################################################################
-# Building Deep Learning Pipelines Using TorchEEG
-# -----------------------------------------
-# Step 1: Initialize the Dataset
-#
-# We use the SEED dataset supported by TorchEEG. Here, we set an EEG sample to 1 second long and include 200 data points. In offline preprocessing, all EEG signals are debaselined and normalized, and the preprocessed EEG signals are stored in the local IO. In online processing, all EEG signals are converted into Tensors for input into neural networks.
-#
-
-dataset = SEEDDataset(io_path=f'./tmp_out/examples_seed_tsception/seed',
-                      root_path='./tmp_in/Preprocessed_EEG',
-                      offline_transform=transforms.Compose([transforms.MeanStdNormalize(),
-                                                            transforms.To2d()]),
-                      online_transform=transforms.ToTensor(),
-                      label_transform=transforms.Compose(
-                          [transforms.Select('emotion'),
-                           transforms.Lambda(lambda x: x + 1)]),
-                      chunk_size=200,
-                      num_worker=4)
-
-######################################################################
-# .. warning::
-#    If you use TorchEEG under the `Windows` system and want to use multiple processes (such as in dataset or dataloader), you should check whether :obj:`__name__` is :obj:`__main__` to avoid errors caused by multiple :obj:`import`.
-#
-# That is, under the :obj:`Windows` system, you need to:
-#  .. code-block::
-#
-#    if __name__ == "__main__":
-#        dataset = SEEDDataset(io_path=f'./tmp_out/examples_seed_tsception/seed',
-#                      root_path='./tmp_in/Preprocessed_EEG',
-#                      offline_transform=transforms.Compose([transforms.MeanStdNormalize(),
-#                                                            transforms.To2d()]),
-#                      online_transform=transforms.ToTensor(),
-#                      label_transform=transforms.Compose(
-#                          [transforms.Select('emotion'),
-#                           transforms.Lambda(lambda x: x + 1)]),
-#                      chunk_size=200,
-#                      num_worker=4)
-#        # the following codes
-#
-# .. warning::
-#    LMDB may not be optimized for parts of Windows systems or storage devices. If you find that the data preprocessing speed is slow, you can consider setting :obj:`io_mode` to :obj:`pickle`, which is an alternative implemented by TorchEEG based on pickle.
-
-######################################################################
-# Step 2: Divide the Training and Test samples in the Dataset
-#
-# Here, the dataset is divided using 5-fold cross-validation. In the process of division, we group according to the trial index, and every trial takes 4 folds as training samples and 1 fold as test samples. Samples across trials are aggregated to obtain training set and test set.
-#
-
-k_fold = KFoldGroupbyTrial(n_splits=5, split_path='./tmp_out/examples_seed_tsception/split')
-
-######################################################################
-# Step 3: Define the Model and Start Training
-#
-# We first use a loop to get the dataset in each cross-validation.
-# In each cross-validation, we initialize the TSCeption model, and define the hyperparameters. For example, the sampling rate of EEG sample is 128, there are 15 temporal modules and 15 spatial modules, etc.
-#
-# We then initialize the trainer and set the hyperparameters in the trained model, such as the learning rate, the equipment used, etc. The :obj:`fit` method receives the training dataset and starts training the model. The :obj:`test` method receives a test dataset and reports the test results. The :obj:`save_state_dict` method can save the trained model.
-
-for i, (train_dataset, val_dataset) in enumerate(k_fold.split(dataset)):
-    # Initialize the model
-    model = TSCeption(num_electrodes=len(SEED_CHANNEL_LOCATION_DICT),
-                      num_classes=3,
-                      num_T=15,
-                      num_S=15,
-                      in_channels=1,
-                      hid_channels=32,
-                      sampling_rate=200,
-                      dropout=0.5)
-
-    # Initialize the trainer and use the 0-th GPU for training
-    trainer = MyClassificationTrainer(model=model, lr=1e-4, weight_decay=1e-4, device_ids=[0])
-
-    # Initialize several batches of training samples and test samples
-    train_loader = DataLoader(train_dataset, batch_size=256, shuffle=True, num_workers=4)
-    val_loader = DataLoader(val_dataset, batch_size=256, shuffle=False, num_workers=4)
-
-    # Do 50 rounds of training
-    trainer.fit(train_loader, val_loader, num_epochs=50)
-    trainer.test(val_loader)
-    trainer.save_state_dict(f'./tmp_out/examples_seed_tsception/weight/{i}.pth')
diff --git a/examples/examples_torch_geometric.py b/examples/examples_torch_geometric.py
deleted file mode 100644
index 828d7e2..0000000
--- a/examples/examples_torch_geometric.py
+++ /dev/null
@@ -1,239 +0,0 @@
-"""Training models with PyTorch Geometric
-======================================
-In this case, we introduce how to use TorchEEG and a customized training process based on vanilla PyTorch to train a PyTorch Geometric-based graph convolutional network on the SEED dataset for emotion classification.
-"""
-
-import logging
-import os
-import random
-import time
-
-import numpy as np
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from torch.nn import Linear
-from torch_geometric.loader import DataLoader
-from torch_geometric.nn import GATConv, global_mean_pool
-
-from torcheeg import transforms
-from torcheeg.datasets import SEEDFeatureDataset
-from torcheeg.datasets.constants.emotion_recognition.seed import \
-    SEED_ADJACENCY_MATRIX
-from torcheeg.model_selection import KFoldPerSubjectGroupbyTrial
-from torcheeg.transforms.pyg import ToG
-
-###############################################################################
-# Pre-experiment Preparation to Ensure Reproducibility
-# -----------------------------------------
-# Use the logging module to store output in a log file for easy reference while printing it to the screen.
-
-os.makedirs('./tmp_out/examples_torch/log', exist_ok=True)
-logger = logging.getLogger('Training models with vanilla PyTorch')
-logger.setLevel(logging.DEBUG)
-console_handler = logging.StreamHandler()
-timeticks = time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime())
-file_handler = logging.FileHandler(
-    os.path.join('./tmp_out/examples_torch/log', f'{timeticks}.log'))
-logger.addHandler(console_handler)
-logger.addHandler(file_handler)
-
-###############################################################################
-# Set the random number seed in all modules to guarantee the same result when running again.
-
-
-def seed_everything(seed):
-    random.seed(seed)
-    np.random.seed(seed)
-    os.environ["PYTHONHASHSEED"] = str(seed)
-    torch.manual_seed(seed)
-    torch.cuda.manual_seed(seed)
-    torch.backends.cudnn.deterministic = True
-    torch.backends.cudnn.benchmark = False
-
-
-seed_everything(42)
-
-###############################################################################
-# Defining a graph convolutional network
-# -----------------------------------------
-# Use the API provided by PyTorch Geometric to define graph convolutional networks. Here, the EEG signal or feature of the electrode corresponds to the input :obj:`data.x`, and the relationship between electrodes corresponds to :obj:`data.edge_index`. Depending on the definition of the adjacency matrix, the relationship may represent spatial adjacency, etc.
-#
-
-
-class GNN(torch.nn.Module):
-    def __init__(self,
-                 in_channels=5,
-                 num_layers=3,
-                 hid_channels=64,
-                 num_classes=3):
-        super().__init__()
-        self.conv1 = GATConv(in_channels, hid_channels)
-        self.convs = torch.nn.ModuleList()
-        for _ in range(num_layers - 1):
-            self.convs.append(GATConv(hid_channels, hid_channels))
-        self.lin1 = Linear(hid_channels, hid_channels)
-        self.lin2 = Linear(hid_channels, num_classes)
-
-    def reset_parameters(self):
-        self.conv1.reset_parameters()
-        for conv in self.convs:
-            conv.reset_parameters()
-        self.lin1.reset_parameters()
-        self.lin2.reset_parameters()
-
-    def forward(self, data):
-        x, edge_index, batch = data.x, data.edge_index, data.batch
-        x = F.relu(self.conv1(x, edge_index))
-        for conv in self.convs:
-            x = F.relu(conv(x, edge_index))
-        x = global_mean_pool(x, batch)
-        x = F.relu(self.lin1(x))
-        x = F.dropout(x, p=0.5, training=self.training)
-        x = self.lin2(x)
-        return x
-
-
-###############################################################################
-# Customize the training process
-# -----------------------------------------
-# TorchEEG provides a large number of trainers to help complete the training of classification models, however, you can also define the training functions to complete the training and testing of the model. Here is a simple example:
-#
-
-
-# training process
-def train(dataloader, model, loss_fn, optimizer):
-    size = len(dataloader.dataset)
-    model.train()
-    for batch_idx, batch in enumerate(dataloader):
-        X = batch[0].to(device)
-        y = batch[1].to(device)
-
-        # Compute prediction error
-        pred = model(X)
-        loss = loss_fn(pred, y)
-
-        # Backpropagation
-        optimizer.zero_grad()
-        loss.backward()
-        optimizer.step()
-
-        if batch_idx % 20 == 0:
-            loss, current = loss.item(), batch_idx * batch_size
-            logger.info(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")
-
-
-# validation process
-def valid(dataloader, model, loss_fn):
-    size = len(dataloader.dataset)
-    num_batches = len(dataloader)
-    model.eval()
-    val_loss, correct = 0, 0
-    with torch.no_grad():
-        for batch in dataloader:
-            X = batch[0].to(device)
-            y = batch[1].to(device)
-
-            pred = model(X)
-            val_loss += loss_fn(pred, y).item()
-            correct += (pred.argmax(1) == y).type(torch.float).sum().item()
-    val_loss /= num_batches
-    correct /= size
-    logger.info(
-        f"Test Error: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {val_loss:>8f} \n"
-    )
-
-
-######################################################################
-# Building Deep Learning Pipelines Using TorchEEG
-# -----------------------------------------
-# Step 1: Initialize the Dataset
-#
-# We use the SEED dataset supported by TorchEEG. Here we are using extracted EEG features. In the feature dataset, EEG signals (200 data points) per second are pre-computed with differential entropy in five sub-bands and smoothed using a linear dynamical system.
-# .. note::
-#   In online processing, All EEG signals are converted into a graph structure that can be processed by the PyTorch Geometric model, i.e., :obj:`torch_geometric.data.Data`, according to the adjacency matrix and the signals or features from electrodes. Here, electrodes represent nodes, and the adjacency matrix defines the association between electrodes.
-#
-
-dataset = SEEDFeatureDataset(io_path='./tmp_out/examples_torch_geometric/seed',
-                             root_path='./tmp_in/ExtractedFeatures',
-                             feature=['de_movingAve'],
-                             online_transform=transforms.Compose([
-                                 transforms.MinMaxNormalize(axis=-1),
-                                 ToG(SEED_ADJACENCY_MATRIX)
-                             ]),
-                             label_transform=transforms.Compose([
-                                 transforms.Select('emotion'),
-                                 transforms.Lambda(lambda x: int(x) + 1),
-                             ]),
-                             num_worker=8)
-
-######################################################################
-# .. warning::
-#    If you use TorchEEG under the `Windows` system and want to use multiple processes (such as in dataset or dataloader), you should check whether :obj:`__name__` is :obj:`__main__` to avoid errors caused by multiple :obj:`import`.
-#
-# That is, under the :obj:`Windows` system, you need to:
-#  .. code-block::
-#
-#    if __name__ == "__main__":
-#        dataset = SEEDFeatureDataset(io_path='./tmp_out/examples_torch_geometric/seed',
-#                              root_path='./tmp_in/ExtractedFeatures',
-#                              feature=['de_movingAve'],
-#                              online_transform=transforms.Compose([
-#                                  transforms.MinMaxNormalize(axis=-1),
-#                                  ToG(SEED_ADJACENCY_MATRIX)
-#                              ]),
-#                              label_transform=transforms.Compose([
-#                                  transforms.Select('emotion'),
-#                                  transforms.Lambda(lambda x: int(x) + 1),
-#                              ]),
-#                              io_mode='pickle',
-#                              num_worker=8)
-#        # the following codes
-#
-# .. note::
-#    LMDB may not be optimized for parts of Windows systems or storage devices. If you find that the data preprocessing speed is slow, you can consider setting :obj:`io_mode` to :obj:`pickle`, which is an alternative implemented by TorchEEG based on pickle.
-
-######################################################################
-# Step 2: Divide the Training and Test samples in the Dataset
-#
-# Here, the dataset is divided using per-subject 5-fold cross-validation. In the process of division, we split the training and test sets separately on each subject's EEG samples. Here, we take 4 folds as training samples and 1 fold as test samples.
-#
-
-k_fold = KFoldPerSubjectGroupbyTrial(
-    n_splits=10,
-    split_path='./tmp_out/examples_torch_geometric/split',
-    shuffle=False)
-
-######################################################################
-# Step 3: Define the Model and Start Training
-#
-# We first use a loop to get the dataset in each cross-validation. In each cross-validation, we initialize the above-mentioned GNN model.
-#
-# Next, we train the model for 50 epochs using the training function defined above and report the model performance on the validation set at each epoch with the validation function defined above.
-#
-
-######################################################################
-# .. note::
-#    Please note that since the EEG signal sample returned by ToG is of type :obj:`torch_geometric.data.Data`, which represents a graph structure, the DataLoader provided by PyTorch cannot be used here to form a batch. Instead, :obj:`torch_geometric.loader.DataLoader` should be used to batch the adjacency matrix of the graph structure and the node features.
-#
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-loss_fn = nn.CrossEntropyLoss()
-batch_size = 64
-
-for i, (train_dataset, val_dataset) in enumerate(k_fold.split(dataset)):
-    # initialize model
-    model = GNN().to(device)
-    # initialize optimizer
-    optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
-
-    train_loader = DataLoader(train_dataset,
-                              batch_size=batch_size,
-                              shuffle=True)
-    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)
-
-    epochs = 50
-    for t in range(epochs):
-        logger.info(f"Epoch {t+1}\n-------------------------------")
-        train(train_loader, model, loss_fn, optimizer)
-        valid(val_loader, model, loss_fn)
diff --git a/examples/examples_torch_lightning.py b/examples/examples_torch_lightning.py
deleted file mode 100644
index a3f7d0a..0000000
--- a/examples/examples_torch_lightning.py
+++ /dev/null
@@ -1,163 +0,0 @@
-"""Training models with Pytorch-Lightning
-======================================
-In this case, we introduce how to use TorchEEG and Pytorch-Lightning to train a Continuous Convolutional Neural Network (CCNN) on the DEAP dataset for emotion classification.
-"""
-
-import os
-
-import torch
-import torch.nn as nn
-from pytorch_lightning import Trainer, seed_everything
-from pytorch_lightning.callbacks import ModelCheckpoint
-from pytorch_lightning.core import LightningModule
-from pytorch_lightning.loggers import TensorBoardLogger
-from torch.nn import functional as F
-from torch.utils.data.dataloader import DataLoader
-from torchmetrics import Accuracy
-
-from torcheeg import transforms
-from torcheeg.datasets import DEAPDataset
-from torcheeg.datasets.constants.emotion_recognition.deap import \
-    DEAP_CHANNEL_LOCATION_DICT
-from torcheeg.model_selection import KFold
-from torcheeg.models import CCNN
-
-###############################################################################
-# Pre-experiment Preparation to Ensure Reproducibility
-# -----------------------------------------
-# Set the random number seed in all modules to guarantee the same result when running again.
-
-seed_everything(42)
-
-###############################################################################
-# Building Deep Learning Pipelines Using Pytorch-Lightning
-# -----------------------------------------
-# Step 1: Define the Pytorch-Lightning Module with training process, validation process, and optimizer configuration.
-
-
-class EEGClassifier(LightningModule):
-    def __init__(self, model, lr=1e-4):
-        super().__init__()
-        self.save_hyperparameters(ignore="model")
-        self.model = model
-        self.val_acc = Accuracy()
-
-    def forward(self, x):
-        return self.model(x)
-
-    def training_step(self, batch, batch_idx):
-        X = batch[0]
-        y = batch[1]
-
-        logits = self.forward(X)
-        loss = F.cross_entropy(logits, y.long())
-        return loss
-
-    def validation_step(self, batch, batch_idx):
-        X = batch[0]
-        y = batch[1]
-
-        logits = self.forward(X)
-        loss = F.cross_entropy(logits, y.long())
-
-        self.val_acc(logits, y)
-        self.log("val_acc", self.val_acc)
-        self.log("val_loss", loss)
-
-    def configure_optimizers(self):
-        optimizer = torch.optim.Adam(self.model.parameters(),
-                                     lr=self.hparams.lr)
-
-        return [optimizer], []
-
-
-######################################################################
-# Step 2: Initialize the Dataset
-#
-# We use the DEAP dataset supported by TorchEEG. Here, we set an EEG sample to 1 second long and include 128 data points. The baseline signal is 3 seconds long, cut into three, and averaged as the baseline signal for the trial. In offline preprocessing, we divide the EEG signal of every electrode into 4 sub-bands, and calculate the differential entropy on each sub-band as a feature, followed by debaselining and mapping on the grid. Finally, the preprocessed EEG signals are stored in the local IO. In online processing, all EEG signals are converted into Tensors for input into neural networks.
-#
-
-dataset = DEAPDataset(
-    io_path=f'./tmp_out/examples_torch_lightning/deap',
-    root_path='./tmp_in/data_preprocessed_python',
-    offline_transform=transforms.Compose([
-        transforms.BandDifferentialEntropy(apply_to_baseline=True),
-        transforms.ToGrid(DEAP_CHANNEL_LOCATION_DICT, apply_to_baseline=True)
-    ]),
-    online_transform=transforms.Compose(
-        [transforms.BaselineRemoval(),
-         transforms.ToTensor()]),
-    label_transform=transforms.Compose([
-        transforms.Select('valence'),
-        transforms.Binary(5.0),
-    ]))
-
-######################################################################
-# .. warning::
-#    If you use TorchEEG under the `Windows` system and want to use multiple processes (such as in dataset or dataloader), you should check whether :obj:`__name__` is :obj:`__main__` to avoid errors caused by multiple :obj:`import`.
-#
-# That is, under the :obj:`Windows` system, you need to:
-#  .. code-block::
-#
-#    if __name__ == "__main__":
-#        dataset = DEAPDataset(
-#                       io_path=f'./tmp_out/examples_torch_lightning/deap',
-#                       root_path='./tmp_in/data_preprocessed_python',
-#                       offline_transform=transforms.Compose([
-#                           transforms.BandDifferentialEntropy(apply_to_baseline=True),
-#                           transforms.ToGrid(DEAP_CHANNEL_LOCATION_DICT, apply_to_baseline=True)
-#                       ]),
-#                       io_mode='pickle',
-#                       online_transform=transforms.Compose(
-#                           [transforms.BaselineRemoval(),
-#                            transforms.ToTensor()]),
-#                       label_transform=transforms.Compose([
-#                           transforms.Select('valence'),
-#                           transforms.Binary(5.0),
-#                       ]))
-#        # the following codes
-#
-# .. note::
-#    LMDB may not be optimized for parts of Windows systems or storage devices. If you find that the data preprocessing speed is slow, you can consider setting :obj:`io_mode` to :obj:`pickle`, which is an alternative implemented by TorchEEG based on pickle.
-
-######################################################################
-# Step 3: Divide the Training and Test samples in the Dataset
-#
-# Here, the dataset is divided using per-subject 5-fold cross-validation. In the process of division, we split the training and test sets separately on each subject's EEG samples. Here, we take 4 folds as training samples and 1 fold as test samples.
-#
-
-k_fold = KFold(n_splits=10,
-               split_path='./tmp_out/examples_torch_lightning/split',
-               shuffle=True,
-               random_state=42)
-
-######################################################################
-# Step 4: Define the Model and Start Training
-#
-# We first use a loop to get the dataset in each cross-validation. In each cross-validation, we initialize the CCNN model and define the hyperparameters. For example, each EEG sample contains 4-channel features from 4 sub-bands, the grid size is 9 times 9, etc.
-#
-# Next, we train the model for 50 epochs, with the Pytorch-Lightning module defined above wrapped in the :obj:`Trainer`. We use the :obj:`TensorBoardLogger` to record the training process and the :obj:`ModelCheckpoint` to save the model with the highest validation accuracy.
-#
-
-for i, (train_dataset, val_dataset) in enumerate(k_fold.split(dataset)):
-    train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
-    val_loader = DataLoader(val_dataset, batch_size=64, shuffle=False)
-
-    tb_logger = TensorBoardLogger(save_dir='lightning_logs',
-                                  name=f'fold_{i + 1}')
-    checkpoint_callback = ModelCheckpoint(
-        dirpath=tb_logger.log_dir,
-        filename="{epoch:02d}-{val_metric:.4f}",
-        monitor='val_metric',
-        mode='max')
-
-    model = EEGClassifier(CCNN(num_classes=2, in_channels=4, grid_size=(9, 9)))
-
-    trainer = Trainer(max_epochs=50,
-                      devices=2,
-                      accelerator="auto",
-                      strategy="ddp",
-                      checkpoint_callback=checkpoint_callback,
-                      logger=tb_logger)
-
-    trainer.fit(model, train_loader, val_loader)
diff --git a/examples/examples_trainers.py b/examples/examples_trainers.py
new file mode 100644
index 0000000..0d3515c
--- /dev/null
+++ b/examples/examples_trainers.py
@@ -0,0 +1,350 @@
+"""
+Introduction to the trainers Module
+===================================
+
+Welcome to the guide on TorchEEG's ``trainers`` Module! This module provides you with a suite of trainers, built on Pytorch-lightning, for
+model training. These trainers are designed to handle a wide array of
+models, ranging from discriminative to generative ones, and are equipped
+with capabilities for contrast learning and fine-tuning. They are also
+capable of running on various hardware configurations, from single CPUs,
+single GPUs, to multiple GPUs. You also have the flexibility to tailor
+these trainers according to your specific needs by extending the
+existing ones.
+
+"""
+
+
+######################################################################
+# Discriminative Models with ClassifierTrainer
+# ----------------------------------------------
+#
+# The simplest yet highly effective method for discriminative models
+# involves utilizing a classification loss function, such as cross-entropy
+# for EEG signal recognition training, encapsulated in the
+# ClassifierTrainer algorithm.
+# 
+
+from torcheeg.datasets import DEAPDataset
+from torcheeg import transforms
+
+from torcheeg.datasets import KFoldGroupbyTrial
+from torcheeg.datasets.constants.emotion_recognition.deap import \
+    DEAP_CHANNEL_LOCATION_DICT
+from torch.utils.data import DataLoader
+from torcheeg.models import CCNN
+
+from torcheeg.trainers import ClassifierTrainer
+
+import pytorch_lightning as pl
+
+dataset = DEAPDataset(
+    io_path=f'./examples_trainers_1/deap',
+    root_path='./data_preprocessed_python',
+    offline_transform=transforms.Compose([
+        transforms.BandDifferentialEntropy(apply_to_baseline=True),
+        transforms.ToGrid(DEAP_CHANNEL_LOCATION_DICT, apply_to_baseline=True)
+    ]),
+    online_transform=transforms.Compose(
+        [transforms.BaselineRemoval(),
+         transforms.ToTensor()]),
+    label_transform=transforms.Compose([
+        transforms.Select('valence'),
+        transforms.Binary(5.0),
+    ]),
+    num_worker=8)
+
+k_fold = KFoldGroupbyTrial(n_splits=10,
+                           split_path='./examples_trainers_1/split',
+                           shuffle=True,
+                           random_state=42)
+
+for i, (train_dataset, val_dataset) in enumerate(k_fold.split(dataset)):
+    train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
+    val_loader = DataLoader(val_dataset, batch_size=64, shuffle=False)
+
+    model = CCNN(num_classes=2, in_channels=4, grid_size=(9, 9))
+
+    trainer = ClassifierTrainer(model=model,
+                                num_classes=2,
+                                lr=1e-4,
+                                weight_decay=1e-4,
+                                devices=1,
+                                accelerator="gpu")
+    trainer.fit(train_loader,
+                val_loader,
+                max_epochs=50,
+                default_root_dir=f'./examples_trainers_1/model/{i}',
+                callbacks=[pl.callbacks.ModelCheckpoint(save_last=True)],
+                enable_progress_bar=True,
+                enable_model_summary=True,
+                limit_val_batches=0.0)
+    score = trainer.test(val_loader,
+                         enable_progress_bar=True,
+                         enable_model_summary=True)[0]
+    print(f'Fold {i} test accuracy: {score["test_accuracy"]:.4f}')
+
+
+######################################################################
+# Within ClassifierTrainer, you can specify the device and the number of
+# devices to use. In the trainer.fit function, you can pass any parameters
+# supported by the Trainer class in pytorch_lightning. If you wish to
+# modify the training procedure, you can do so by extending the trainers:
+# 
+
+class MyClassifierTrainer(ClassifierTrainer):
+    def training_step(self, batch, batch_idx):
+        x, y = batch
+        y_hat = self(x)
+        # your loss function
+        loss = self.ce_fn(y_hat, y)
+
+        # log to prog_bar
+        self.log("train_loss",
+                 self.train_loss(loss),
+                 prog_bar=True,
+                 on_epoch=False,
+                 logger=False,
+                 on_step=True)
+
+        for i, metric_value in enumerate(self.train_metrics.values()):
+            self.log(f"train_{self.metrics[i]}",
+                     metric_value(y_hat, y),
+                     prog_bar=True,
+                     on_epoch=False,
+                     logger=False,
+                     on_step=True)
+
+        return loss
+
+
+######################################################################
+# By referring to the related documents, you can also learn about the
+# evaluation metrics supported in the trainers module. By making a few
+# adjustments, you can get the trainer to report metrics like accuracy and
+# f1 score.
+# 
+
+trainer = ClassifierTrainer(model=model,
+                                num_classes=2,
+                                lr=1e-4,
+                                weight_decay=1e-4,
+                                metrics=['accuracy', 'recall', 'precision', 'f1score'],
+                                accelerator="gpu")
+
+
+######################################################################
+# Domain Adaptation Methods
+# ----------------------------------------------
+#
+# A challenging aspect of EEG-based emotion recognition is the
+# cross-subject problem. Even when evoked by the same stimulus, the
+# distribution of EEG signal patterns among different individuals may
+# undergo distributional shifts. This phenomenon negatively impacts the
+# performance of the trained models when applied to unknown subjects. The
+# trainers module addresses this issue by providing a host of domain
+# adaptation algorithms. These algorithms use EEG samples from known
+# subjects (source domain) and unknown subjects for testing (target
+# domain). They optimize models with specific loss functions and training
+# strategies to extract domain-invariant features or transfer knowledge
+# from the source to the target domain. A range of cross-domain trainers
+# like CORALTrainer, ADATrainer, DANNTrainer, DDCTrainer, and DANTrainer
+# are available to handle different application scenarios and assist users
+# in dealing with cross-domain problems across diverse models and
+# datasets.
+# 
+
+from torcheeg.datasets import DEAPDataset
+from torcheeg import transforms
+
+from torcheeg.datasets import LeaveOneSubjectOut
+from torcheeg.datasets.constants.emotion_recognition.deap import \
+    DEAP_CHANNEL_LOCATION_DICT
+from torch.utils.data import DataLoader
+from torcheeg.models import CCNN
+
+from torcheeg.trainers import CORALTrainer
+
+import pytorch_lightning as pl
+
+dataset = DEAPDataset(
+    io_path=f'./examples_trainers_2/deap',
+    root_path='./data_preprocessed_python',
+    offline_transform=transforms.Compose([
+        transforms.BandDifferentialEntropy(apply_to_baseline=True),
+        transforms.ToGrid(DEAP_CHANNEL_LOCATION_DICT, apply_to_baseline=True)
+    ]),
+    online_transform=transforms.Compose(
+        [transforms.BaselineRemoval(),
+         transforms.ToTensor()]),
+    label_transform=transforms.Compose([
+        transforms.Select('valence'),
+        transforms.Binary(5.0),
+    ]),
+    num_worker=8)
+
+k_fold = LeaveOneSubjectOut(split_path='./examples_trainers_2/split',
+                            shuffle=True,
+                            random_state=42)
+
+
+class Extractor(CCNN):
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.conv2(x)
+        x = self.conv3(x)
+        x = self.conv4(x)
+        x = x.flatten(start_dim=1)
+        return x
+
+
+class Classifier(CCNN):
+    def forward(self, x):
+        x = self.lin1(x)
+        x = self.lin2(x)
+        return x
+
+
+for i, (train_dataset, val_dataset) in enumerate(k_fold.split(dataset)):
+    source_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
+    target_loader = DataLoader(val_dataset, batch_size=64, shuffle=False)
+
+    extractor = Extractor(in_channels=5, num_classes=2)
+    classifier = Classifier(in_channels=5, num_classes=2)
+
+    trainer = CORALTrainer(extractor=extractor,
+                                classifier=classifier,
+                                num_classes=2,
+                                lr=1e-4,
+                                weight_decay=0.0,
+                                accelerator='gpu')
+    trainer.fit(source_loader,
+                target_loader,
+                target_loader,
+                max_epochs=50,
+                default_root_dir=f'./examples_trainers_2/model/{i}',
+                callbacks=[pl.callbacks.ModelCheckpoint(save_last=True)],
+                enable_progress_bar=True,
+                enable_model_summary=True,
+                limit_val_batches=0.0)
+    score = trainer.test(target_loader,
+                         enable_progress_bar=True,
+                         enable_model_summary=True)[0]
+    print(f'Fold {i} test accuracy: {score["test_accuracy"]:.4f}')
+
+
+######################################################################
+# Generative Models
+# ----------------------------------------------
+#
+# For generative models, we provide trainers for models including VAE,
+# GAN, Normalizing Flow, and Diffusion Model. These trainers aim to train
+# models to generate EEG signals that closely mimic the real distribution.
+# We also provide conditional versions of these trainers, enabling the use
+# of categories as conditions for category-driven EEG sample generation.
+# 
+
+from torcheeg.trainers import CDDPMTrainer
+from torcheeg.models import BCUNet
+
+model = BCUNet(in_channels=4)
+trainer = CDDPMTrainer(model, accelerator='gpu')
+
+from torcheeg.trainers import CWGANGPTrainer
+from torcheeg.models import BCGenerator, BCDiscriminator
+
+g_model = BCGenerator(in_channels=128)
+d_model = BCDiscriminator(in_channels=4)
+
+trainer = CWGANGPTrainer(g_model,
+                              d_model,
+                             accelerator='gpu')
+
+from torcheeg.trainers import CGlowTrainer
+from torcheeg.models import BCGlow
+
+model = BCGlow(in_channels=4)
+trainer = CGlowTrainer(model, accelerator='gpu')
+
+
+######################################################################
+# TorchEEG supports common evaluation metrics for generative models, such
+# as FID. To use these metrics, you need to provide additional parameters
+# to the Trainer like metric_extractor, metric_classifier, and
+# metric_num_features. For details, please refer to the related documents.
+# Here is an example:
+# 
+
+import torch.nn as nn
+
+class Extractor(nn.Module):
+
+    def __init__(self, in_channels=4):
+        super().__init__()
+        self.conv1 = nn.Sequential(
+            nn.ZeroPad2d((1, 2, 1, 2)),
+            nn.Conv2d(in_channels, 64, kernel_size=4, stride=1), nn.ReLU())
+        self.conv2 = nn.Sequential(nn.ZeroPad2d((1, 2, 1, 2)),
+                                   nn.Conv2d(64, 128, kernel_size=4, stride=1),
+                                   nn.ReLU())
+        self.conv3 = nn.Sequential(nn.ZeroPad2d((1, 2, 1, 2)),
+                                   nn.Conv2d(128, 256, kernel_size=4, stride=1),
+                                   nn.ReLU())
+        self.conv4 = nn.Sequential(nn.ZeroPad2d((1, 2, 1, 2)),
+                                   nn.Conv2d(256, 64, kernel_size=4, stride=1),
+                                   nn.ReLU())
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.conv2(x)
+        x = self.conv3(x)
+        x = self.conv4(x)
+
+        x = x.flatten(start_dim=1)
+        return x
+
+
+class Classifier(nn.Module):
+
+    def __init__(self, in_channels=4, num_classes=2):
+        super().__init__()
+        self.conv1 = nn.Sequential(
+            nn.ZeroPad2d((1, 2, 1, 2)),
+            nn.Conv2d(in_channels, 64, kernel_size=4, stride=1), nn.ReLU())
+        self.conv2 = nn.Sequential(nn.ZeroPad2d((1, 2, 1, 2)),
+                                   nn.Conv2d(64, 128, kernel_size=4, stride=1),
+                                   nn.ReLU())
+        self.conv3 = nn.Sequential(nn.ZeroPad2d((1, 2, 1, 2)),
+                                   nn.Conv2d(128, 256, kernel_size=4, stride=1),
+                                   nn.ReLU())
+        self.conv4 = nn.Sequential(nn.ZeroPad2d((1, 2, 1, 2)),
+                                   nn.Conv2d(256, 64, kernel_size=4, stride=1),
+                                   nn.ReLU())
+
+        self.lin1 = nn.Linear(9 * 9 * 64, 1024)
+        self.lin2 = nn.Linear(1024, num_classes)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.conv2(x)
+        x = self.conv3(x)
+        x = self.conv4(x)
+
+        x = x.flatten(start_dim=1)
+        x = self.lin1(x)
+        x = self.lin2(x)
+        return x
+    
+g_model = BCGenerator(in_channels=128)
+d_model = BCDiscriminator(in_channels=4)
+
+extractor = Extractor()
+classifier = Classifier()
+# you may need to load state dict from your trained extractor, classifier
+
+trainer = CWGANGPTrainer(g_model,
+                              d_model,
+                              metric_extractor=extractor,
+                              metric_classifier=classifier,
+                              metric_num_features=9 * 9 * 64,
+                              metrics=['fid'],
+                             accelerator='gpu')
\ No newline at end of file
diff --git a/examples/examples_transfer_mahnob_2_deap.py b/examples/examples_transfer_mahnob_2_deap.py
new file mode 100644
index 0000000..2ba13d3
--- /dev/null
+++ b/examples/examples_transfer_mahnob_2_deap.py
@@ -0,0 +1,205 @@
+"""
+Transfer Learning from MAHNOB-HCI to DEAP
+==========================
+
+In this tutorial, we demonstrate how to use TorchEEG to implement transfer learning with a Continuous Convolutional Neural Network (CCNN) model. We'll train our model on the MAHNOB-HCI dataset and then apply it to the DEAP dataset.
+
+"""
+
+######################################################################
+# Step 1: Initialize the Dataset
+# ------------------------------
+#
+# We use the MAHNOB dataset supported by TorchEEG. Each EEG sample is set to
+# be 1 second long, encompassing 128 data points. The baseline signal is 30
+# seconds long, which we divide into three sections and then average to
+# obtain the trial’s baseline signal.
+#
+# During offline preprocessing, we divide each electrode’s EEG signal into
+# 4 sub-bands, calculate the differential entropy for each sub-band as a
+# feature, perform debaselining, and map onto a grid. Finally, the
+# preprocessed EEG signals are saved locally. For online processing, we
+# convert all EEG signals into Tensors, making them suitable for neural
+# network input.
+#
+# We use the `before_trial_normalize` function to normalize each EEG trial. This helps in reducing the variance between trials, thereby aiding the transfer generalizable knowledge from MAHNOB to DEAP.
+#
+
+from torcheeg.datasets import MAHNOBDataset
+from torcheeg import transforms
+from torcheeg.datasets.constants.emotion_recognition.mahnob import MAHNOB_CHANNEL_LOCATION_DICT
+
+from torcheeg.datasets.functional import before_trial_normalize
+
+dataset = MAHNOBDataset(
+    io_path='./examples_transfer_mahnob_2_deap/mahnob',
+    root_path='./Sessions',
+    offline_transform=transforms.Compose([
+        transforms.BandDifferentialEntropy(sampling_rate=128,
+                                           apply_to_baseline=True),
+        transforms.BaselineRemoval(),
+        transforms.ToGrid(MAHNOB_CHANNEL_LOCATION_DICT)
+    ]),
+    online_transform=transforms.ToTensor(),
+    after_subject=before_trial_normalize,
+    label_transform=transforms.Compose(
+        [transforms.Select('feltVlnc'),
+         transforms.Binary(5.0)]),
+    chunk_size=128,
+    baseline_chunk_size=128,
+    num_baseline=30,
+    num_worker=4)
+
+######################################################################
+# Step 2: Divide the Training and Test samples in the Dataset
+# ===========================================================
+#
+# For this example, we'll use a simple train-test split to partition the MAHNOB-HCI dataset.
+#
+
+from torcheeg.model_selection import train_test_split_groupby_trial
+
+train_dataset, val_dataset = train_test_split_groupby_trial(
+    dataset,
+    test_size=0.2,
+    split_path=f'./examples_transfer_mahnob_2_deap/split/mahnob',
+    shuffle=True)
+
+######################################################################
+# Step 3: Define the Model and Initiate Training
+# ----------------------------------------------
+#
+# We loop through each cross-validation set, and for each one, we
+# initialize the CCNN model and define its hyperparameters. For instance,
+# each EEG sample contains 4-channel features from 4 sub-bands, and the
+# grid size is 9x9.
+#
+# We then train the model for 50 epochs using the ``ClassifierTrainer``.
+#
+# Here, we observe that different datasets have different level of class
+# imbalance. To address this issue, we use the ImbalancedDatasetSampler
+# from torchsampler to sample the training data.
+#
+
+import pytorch_lightning as pl
+
+from torcheeg.models import CCNN
+from torcheeg.trainers import ClassifierTrainer
+
+from torch.utils.data import DataLoader
+from torchsampler import ImbalancedDatasetSampler
+
+train_loader = DataLoader(train_dataset,
+                          batch_size=64,
+                          sampler=ImbalancedDatasetSampler(train_dataset))
+val_loader = DataLoader(val_dataset, batch_size=64, shuffle=False)
+
+model = CCNN(num_classes=3, in_channels=4, grid_size=(9, 9))
+
+trainer = ClassifierTrainer(model=model,
+                            num_classes=2,
+                            lr=1e-4,
+                            weight_decay=1e-4,
+                            accelerator="gpu")
+
+trainer.fit(
+    train_loader,
+    val_loader,
+    max_epochs=50,
+    default_root_dir=f'./examples_transfer_mahnob_2_deap/pretrained_model',
+    callbacks=[pl.callbacks.ModelCheckpoint(save_last=True)],
+    enable_progress_bar=True,
+    enable_model_summary=True,
+    limit_val_batches=0.0)
+
+state_dict = model.state_dict()
+
+######################################################################
+# Step 4: Initialize the Dataset
+# ------------------------------
+#
+# We then switch to the DEAP dataset, which is also supported by TorchEEG. Similar to the MAHNOB-HCI dataset, each EEG sample is set to
+# be 1 second long, encompassing 128 data points. The baseline signal is 3
+# seconds long, which we divide into three sections and then average to
+# obtain the trial’s baseline signal.
+#
+# During offline preprocessing, we divide each electrode’s EEG signal into
+# 4 sub-bands, calculate the differential entropy for each sub-band as a
+# feature, perform debaselining, and map onto a grid. Finally, the
+# preprocessed EEG signals are saved locally. For online processing, we
+# convert all EEG signals into Tensors, making them suitable for neural
+# network input.
+#
+
+from torcheeg.datasets import DEAPDataset
+
+from torcheeg.datasets.constants.emotion_recognition.deap import \
+    DEAP_CHANNEL_LOCATION_DICT
+
+dataset = DEAPDataset(
+    io_path=f'./examples_transfer_mahnob_2_deap/deap',
+    root_path='./data_preprocessed_python',
+    offline_transform=transforms.Compose([
+        transforms.BandDifferentialEntropy(apply_to_baseline=True),
+        transforms.ToGrid(DEAP_CHANNEL_LOCATION_DICT, apply_to_baseline=True)
+    ]),
+    online_transform=transforms.Compose(
+        [transforms.BaselineRemoval(),
+         transforms.ToTensor()]),
+    label_transform=transforms.Compose([
+        transforms.Select('valence'),
+        transforms.Binary(5.0),
+    ]),
+    num_worker=8)
+
+######################################################################
+# Step 5: Split the Dataset into Training and Test Sets
+# -----------------------------------------------------
+#
+# In this case, we use 5-fold cross-validation to split the
+# dataset. During this process, we separate EEG samples of each trial
+# into training and test sets. We use 4 folds for training and 1 fold for
+# testing.
+#
+
+from torcheeg.model_selection import KFoldGroupbyTrial
+
+k_fold = KFoldGroupbyTrial(
+    n_splits=10,
+    split_path='./examples_transfer_mahnob_2_deap/split/deap',
+    shuffle=True,
+    random_state=42)
+
+######################################################################
+# Step 6: Transfer Learning on the DEAP Dataset
+# ----------------------------------------------
+#
+# Finally, we load the pre-trained CCNN model and fine-tune it on the DEAP dataset. We maintain the same set of hyperparameters for consistency.
+#
+
+for i, (train_dataset, val_dataset) in enumerate(k_fold.split(dataset)):
+    train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
+    val_loader = DataLoader(val_dataset, batch_size=64, shuffle=False)
+
+    model = CCNN(num_classes=2, in_channels=4, grid_size=(9, 9))
+    model.load_state_dict(state_dict)
+
+    trainer = ClassifierTrainer(model=model,
+                                num_classes=2,
+                                lr=1e-4,
+                                weight_decay=1e-4,
+                                accelerator="gpu")
+    trainer.fit(
+        train_loader,
+        val_loader,
+        max_epochs=50,
+        default_root_dir=f'./examples_transfer_mahnob_2_deap/pretrained_model',
+        callbacks=[pl.callbacks.ModelCheckpoint(save_last=True)],
+        enable_progress_bar=True,
+        enable_model_summary=True,
+        limit_val_batches=0.0)
+
+    score = trainer.test(val_loader,
+                         enable_progress_bar=True,
+                         enable_model_summary=True)[0]
+    print(f'Fold {i} test accuracy: {score["test_accuracy"]:.4f}')
\ No newline at end of file
diff --git a/examples/examples_transforms.py b/examples/examples_transforms.py
new file mode 100644
index 0000000..4e2e798
--- /dev/null
+++ b/examples/examples_transforms.py
@@ -0,0 +1,122 @@
+"""
+Introduction to the transforms Module
+=====================================
+
+Welcome to the guide on TorchEEG's ``transforms`` Module! This module provides you with toolset designed for the transformation and preprocessing of EEG signals and related metadata.
+
+EEG transforms concentrate on EEG preprocessing, feature extraction,
+data conversion, and data augmentation for deep learning models. They
+typically accept various parameters for instantiation, and the
+instantiated transformers can be used as functions to process EEG
+signals. They accept 'eeg' and 'baseline' as keyword arguments and
+return a dictionary, in which 'eeg' corresponds to the processed EEG
+signal, while 'baseline' corresponds to the processed baseline signal.
+
+"""
+
+import numpy as np
+from torcheeg import transforms
+
+from torcheeg.datasets.constants.emotion_recognition.deap import \
+    DEAP_CHANNEL_LOCATION_DICT
+
+t = transforms.ToGrid(DEAP_CHANNEL_LOCATION_DICT)
+eeg = t(eeg=np.random.randn(32, 128))['eeg']
+print(eeg.shape)
+
+eeg = t(eeg=np.random.randn(32, 128), baseline=np.random.randn(32, 128))['eeg']
+print(eeg.shape)
+
+######################################################################
+# When it comes to preprocessing, EEG transforms aid in the elimination of
+# noise from EEG signals. A prime example of this is the BaselineRemoval,
+# which employs the baseline signal prior to the stimulus to remove
+# stimulus-independent fluctuations from the processed signal.
+#
+
+t = transforms.BaselineRemoval()
+
+eeg = t(eeg=np.random.randn(32, 128), baseline=np.random.randn(32, 128))['eeg']
+print(eeg.shape)
+
+######################################################################
+# For feature extraction, EEG transforms take charge of creating
+# emotion-discriminative features usually employed in emotion recognition.
+# One such example is the BandDifferentialEntropy, which gauges the
+# differential entropy of the frequency subbands within the EEG signal.
+#
+
+eeg = np.random.randn(32, 128)
+transformed_eeg = transforms.BandDifferentialEntropy()(eeg=eeg)['eeg']
+
+######################################################################
+# In terms of data conversion, EEG transforms play a pivotal role in
+# shaping EEG signals into time series, 3D grids, or graph networks.
+# Examples of these include ToGrid, ToInterpolatedGrid, ToG, and
+# ToDynamicG.
+#
+
+from torcheeg.datasets.constants import DEAP_ADJACENCY_MATRIX
+from torcheeg.transforms.pyg import ToG
+
+eeg = np.random.randn(32, 128)
+transformed_eeg = ToG(DEAP_ADJACENCY_MATRIX)(eeg=eeg)['eeg']
+print(transformed_eeg)
+
+######################################################################
+# Regarding data augmentation, transforms such as RandomMask apply random
+# transformations to the EEG signals, effectively augmenting the data
+# volume.
+#
+
+import torch
+
+eeg = torch.randn(32, 128)
+transformed_eeg = transforms.RandomMask(p=1.0)(eeg=eeg)['eeg']
+print(transformed_eeg.shape)
+
+######################################################################
+# Metadata transforms interact with the metadata associated with EEG
+# signals, generating ancillary information or labels needed for models.
+# An example is Binary, which transforms continuous emotional valence or
+# arousal into binary labels.
+#
+
+info = {'valence': 4.5, 'arousal': 5.5, 'subject_id': 7}
+transformed_label = transforms.Select(key='valence')(y=info)['y']
+print(transformed_label)
+
+transforms.Binary(threshold=5.0)(y=transformed_label)['y']
+print(transformed_label)
+
+######################################################################
+# Process transforms bring several transformers together, crafting the
+# overall processes. For instance, Compose allows the output of one
+# transformation to be the input of the next, thus chaining multiple
+# transforms together.
+#
+
+from torcheeg import transforms
+
+t = transforms.Compose([
+    transforms.BandDifferentialEntropy(),
+    transforms.ToGrid(DEAP_CHANNEL_LOCATION_DICT)
+])
+eeg = t(eeg=np.random.randn(32, 128))['eeg']
+print(eeg.shape)
+
+######################################################################
+# It's worth noting that unless specifically required, transformers before
+# BaselineRemoval should be set with apply_to_baseline=True, to ensure the
+# baseline signal and the experiment signal undergo the same
+# transformations. Of course, you don't need to do this if BaselineRemoval
+# is not needed.
+#
+
+t = transforms.Compose([
+    transforms.BandDifferentialEntropy(apply_to_baseline=True),
+    transforms.ToGrid(DEAP_CHANNEL_LOCATION_DICT, apply_to_baseline=True),
+    transforms.BaselineRemoval()
+])
+eeg = t(eeg=np.random.randn(32, 128), baseline=np.random.randn(32, 128))['eeg']
+print(eeg.shape)
diff --git a/examples/examples_torch.py b/examples/examples_vanilla_torch.py
similarity index 52%
rename from examples/examples_torch.py
rename to examples/examples_vanilla_torch.py
index 682bac8..ff2633c 100644
--- a/examples/examples_torch.py
+++ b/examples/examples_vanilla_torch.py
@@ -1,42 +1,45 @@
-"""Training models with vanilla PyTorch
+"""
+Training Models Using Vanilla PyTorch
 ======================================
-In this case, we introduce how to use TorchEEG and a customized training process based on vanilla PyTorch to train a Continuous Convolutional Neural Network (CCNN) on the DEAP dataset for emotion classification.
+
+For those who prefer working directly with PyTorch's native functionalities instead of using third-party trainers, this tutorial provides a complete guide on how to train models using standard PyTorch features.
+
 """
 
-import logging
-import os
-import random
-import time
 
-import numpy as np
-import torch
-import torch.nn as nn
-from torch.utils.data.dataloader import DataLoader
-from torcheeg import transforms
-from torcheeg.datasets import DEAPDataset
-from torcheeg.datasets.constants.emotion_recognition.deap import \
-    DEAP_CHANNEL_LOCATION_DICT
-from torcheeg.model_selection import KFoldPerSubject, train_test_split
-from torcheeg.models import CCNN
+######################################################################
+# Pre-Experiment Preparation to Guarantee Reproducibility
+# -------------------------------------------------------
+# 
+# It is recommended to utilize the logging module to record output in a
+# log file, enabling easy referencing while simultaneously displaying it
+# on the screen.
+# 
 
-###############################################################################
-# Pre-experiment Preparation to Ensure Reproducibility
-# -----------------------------------------
-# Use the logging module to store output in a log file for easy reference while printing it to the screen.
+import os
+import time
+import logging
 
-os.makedirs('./tmp_out/examples_torch/log', exist_ok=True)
+os.makedirs('./examples_vanilla_torch/log', exist_ok=True)
 logger = logging.getLogger('Training models with vanilla PyTorch')
 logger.setLevel(logging.DEBUG)
 console_handler = logging.StreamHandler()
 timeticks = time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime())
 file_handler = logging.FileHandler(
-    os.path.join('./tmp_out/examples_torch/log', f'{timeticks}.log'))
+    os.path.join('./examples_vanilla_torch/log', f'{timeticks}.log'))
 logger.addHandler(console_handler)
 logger.addHandler(file_handler)
 
-###############################################################################
-# Set the random number seed in all modules to guarantee the same result when running again.
 
+######################################################################
+# Setting the random number seed in all modules is essential to ensure
+# consistent results upon repeated runs.
+# 
+
+import random
+import numpy as np
+
+import torch
 
 def seed_everything(seed):
     random.seed(seed)
@@ -50,12 +53,71 @@ def seed_everything(seed):
 
 seed_everything(42)
 
-###############################################################################
-# Customize the training process
-# -----------------------------------------
-# TorchEEG provides a large number of trainers to help complete the training of classification models, however, you can also define the training functions to complete the training and testing of the model. Here is a simple example:
-#
 
+######################################################################
+# Step 1: Initialize the Dataset
+# ------------------------------
+# 
+# For this example, we utilize the DEAP dataset, as supported by TorchEEG.
+# We configure an EEG sample to be 1 second long, encompassing 128 data
+# points. A 3-second-long baseline signal is segmented into three parts
+# and averaged to establish the baseline signal for the trial. During
+# offline preprocessing, we divide each electrode’s EEG signal into 4
+# sub-bands and compute the differential entropy on each sub-band as a
+# feature, which is then subjected to debaselining and grid mapping. The
+# preprocessed EEG signals are then stored in the local IO. In the case of
+# online processing, all EEG signals are transformed into Tensors for
+# input into the neural networks.
+# 
+
+from torcheeg.datasets import DEAPDataset
+from torcheeg import transforms
+
+from torcheeg.datasets.constants.emotion_recognition.deap import \
+    DEAP_CHANNEL_LOCATION_DICT
+
+dataset = DEAPDataset(
+    io_path=f'./examples_vanilla_torch/deap',
+    root_path='./tmp_in/data_preprocessed_python',
+    offline_transform=transforms.Compose([
+        transforms.BandDifferentialEntropy(apply_to_baseline=True),
+        transforms.BaselineRemoval(),
+        transforms.ToGrid(DEAP_CHANNEL_LOCATION_DICT)
+    ]),
+    online_transform=transforms.ToTensor(),
+    label_transform=transforms.Compose([
+        transforms.Select('valence'),
+        transforms.Binary(5.0),
+    ]),
+    num_worker=8)
+
+
+######################################################################
+# Step 2: Partition the Dataset into Training and Test Samples
+# ------------------------------------------------------------
+# 
+# The dataset is divided using per-subject 5-fold cross-validation. During
+# the division process, we separate the training and test sets based on
+# each subject’s EEG samples. We allocate 4 folds to the training samples
+# and 1 fold to the test samples.
+# 
+
+from torcheeg.model_selection import KFoldPerSubject
+
+k_fold = KFoldPerSubject(n_splits=10,
+                         split_path='./examples_vanilla_torch/split',
+                         shuffle=True)
+
+
+######################################################################
+# The training process is customizable. You have the flexibility to define
+# the training functions to facilitate the model’s training and testing.
+# Here is a basic example:
+# 
+
+import torch
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # training process
 def train(dataloader, model, loss_fn, optimizer):
@@ -105,74 +167,26 @@ def valid(dataloader, model, loss_fn):
 
 
 ######################################################################
-# Building Deep Learning Pipelines Using TorchEEG
-# -----------------------------------------
-# Step 1: Initialize the Dataset
-#
-# We use the DEAP dataset supported by TorchEEG. Here, we set an EEG sample to 1 second long and include 128 data points. The baseline signal is 3 seconds long, cut into three, and averaged as the baseline signal for the trial. In offline preprocessing, we divide the EEG signal of every electrode into 4 sub-bands, and calculate the differential entropy on each sub-band as a feature, followed by debaselining and mapping on the grid. Finally, the preprocessed EEG signals are stored in the local IO. In online processing, all EEG signals are converted into Tensors for input into neural networks.
-#
-
-dataset = DEAPDataset(
-    io_path=f'./tmp_out/examples_torch/deap',
-    root_path='./tmp_in/data_preprocessed_python',
-    offline_transform=transforms.Compose([
-        transforms.BandDifferentialEntropy(apply_to_baseline=True),
-        transforms.BaselineRemoval(),
-        transforms.ToGrid(DEAP_CHANNEL_LOCATION_DICT)
-    ]),
-    online_transform=transforms.ToTensor(),
-    label_transform=transforms.Compose([
-        transforms.Select('valence'),
-        transforms.Binary(5.0),
-    ]),
-    num_worker=8)
-
-######################################################################
-# .. warning::
-#    If you use TorchEEG under the `Windows` system and want to use multiple processes (such as in dataset or dataloader), you should check whether :obj:`__name__` is :obj:`__main__` to avoid errors caused by multiple :obj:`import`.
-#
-# That is, under the :obj:`Windows` system, you need to:
-#  .. code-block::
-#
-#    if __name__ == "__main__":
-#        dataset = DEAPDataset(io_path=f'./tmp_out/examples_ccnn/deap',
-#                              root_path='./tmp_in/data_preprocessed_python',
-#                              offline_transform=transforms.Compose([
-#                                  transforms.BandDifferentialEntropy(apply_to_baseline=True),
-#                                  transforms.BaselineRemoval(),
-#                                  transforms.ToGrid(DEAP_CHANNEL_LOCATION_DICT)
-#                              ]),
-#                              online_transform=transforms.ToTensor(),
-#                              label_transform=transforms.Compose([
-#                                  transforms.Select('valence'),
-#                                  transforms.Binary(5.0),
-#                              ]),
-#                              io_mode='pickle',
-#                              num_worker=8)
-#        # the following codes
-#
-# .. note::
-#    LMDB may not be optimized for parts of Windows systems or storage devices. If you find that the data preprocessing speed is slow, you can consider setting :obj:`io_mode` to :obj:`pickle`, which is an alternative implemented by TorchEEG based on pickle.
-
-######################################################################
-# Step 2: Divide the Training and Test samples in the Dataset
-#
-# Here, the dataset is divided using per-subject 5-fold cross-validation. In the process of division, we split the training and test sets separately on each subject's EEG samples. Here, we take 4 folds as training samples and 1 fold as test samples.
-#
+# Step 3: Define the Model and Initiate Training
+# ----------------------------------------------
+# 
+# Initially, we use a loop to access the dataset in each cross-validation.
+# During each cross-validation, we initialize the CCNN model and establish
+# the hyperparameters. For instance, each EEG sample includes 4-channel
+# features from 4 sub-bands, the grid size is set to 9x9, and so on.
+# 
+# Subsequently, we train the model for 50 epochs using the previously
+# defined training function and monitor the model’s performance on the
+# validation set at each epoch using the validation function defined
+# earlier.
+# 
 
-k_fold = KFoldPerSubject(n_splits=10,
-                         split_path='./tmp_out/examples_ccnn/split',
-                         shuffle=True)
+import torch.nn as nn
+from torcheeg.models import CCNN
 
-######################################################################
-# Step 3: Define the Model and Start Training
-#
-# We first use a loop to get the dataset in each cross-validation. In each cross-validation, we initialize the CCNN model and define the hyperparameters. For example, each EEG sample contains 4-channel features from 4 sub-bands, the grid size is 9 times 9, etc.
-#
-# Next, we train the model for 50 epochs using the training function defined above and report the model performance on the validation set at each epoch with the validation function defined above.
-#
+from torcheeg.model_selection import train_test_split
+from torch.utils.data.dataloader import DataLoader
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
 loss_fn = nn.CrossEntropyLoss()
 batch_size = 64
 
@@ -189,7 +203,7 @@ def valid(dataloader, model, loss_fn):
     train_dataset, val_dataset = train_test_split(
         train_dataset,
         test_size=0.2,
-        split_path=f'./tmp_out/examples_ccnn/split{i}',
+        split_path=f'./examples_vanilla_torch/split{i}',
         shuffle=True)
     train_loader = DataLoader(train_dataset,
                               batch_size=batch_size,
@@ -205,12 +219,12 @@ def valid(dataloader, model, loss_fn):
         if val_acc > best_val_acc:
             best_val_acc = val_acc
             torch.save(model.state_dict(),
-                       f'./tmp_out/examples_ccnn/model{i}.pt')
+                       f'./examples_vanilla_torch/model{i}.pt')
 
     test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
 
     # load the best model to test on test set
-    model.load_state_dict(torch.load(f'./tmp_out/examples_ccnn/model{i}.pt'))
+    model.load_state_dict(torch.load(f'./examples_vanilla_torch/model{i}.pt'))
     test_acc, test_loss = valid(test_loader, model, loss_fn)
 
     # log the test result
@@ -224,4 +238,4 @@ def valid(dataloader, model, loss_fn):
 # log the average test result on cross-validation datasets
 logger.info(
     f"Test Error: \n Accuracy: {100*np.mean(test_accs):>0.1f}%, Avg loss: {np.mean(test_losses):>8f}"
-)
+)
\ No newline at end of file
diff --git a/torcheeg/__init__.py b/torcheeg/__init__.py
index b19d9ca..c1b0749 100644
--- a/torcheeg/__init__.py
+++ b/torcheeg/__init__.py
@@ -1,4 +1,40 @@
 __version__ = '1.0.11'
 
-import warnings
-warnings.filterwarnings("ignore")
\ No newline at end of file
+import logging
+
+from mne import set_log_level as _set_mne_log_level
+
+
+def _set_lightning_log_level(level: str = "INFO"):
+
+    # name to level
+    level = getattr(logging, level)
+
+    # configure logging at the root level of Lightning
+    logging.getLogger("lightning.pytorch").setLevel(level)
+
+
+def set_log_level(level: str = "DEBUG", third_party: str = "ERROR"):
+    """
+    Set log level for torcheeg.
+
+    Args:
+        level (str): Log level. Choose one of ["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"].
+        third_party (str): Log level for third party libraries. Choose one of ["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"].
+    """
+    VALID_LEVELS = ["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"]
+    level = level.upper()
+    if level not in VALID_LEVELS:
+        raise ValueError(
+            f"Invalid level {level}. Choose one of {VALID_LEVELS}.")
+
+    _set_mne_log_level(third_party)
+    _set_lightning_log_level(third_party)
+
+    logging.basicConfig(
+        level=level,
+        format="%(asctime)s %(levelname)s %(threadName)s %(name)s %(message)s",
+    )
+
+
+set_log_level()
\ No newline at end of file
diff --git a/torcheeg/datasets/constants/clinical/__init__.py b/torcheeg/datasets/constants/clinical/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/torcheeg/datasets/constants/clinical/tuh_tueg.py b/torcheeg/datasets/constants/clinical/tuh_tueg.py
new file mode 100644
index 0000000..d65d353
--- /dev/null
+++ b/torcheeg/datasets/constants/clinical/tuh_tueg.py
@@ -0,0 +1,63 @@
+from ..region_1020 import (FRONTAL_REGION_LIST, GENERAL_REGION_LIST,
+                           HEMISPHERE_REGION_LIST)
+from ..standard_1005 import STANDARD_1005_CHANNEL_LOCATION_DICT
+from ..utils import (format_adj_matrix_from_adj_list,
+                     format_adj_matrix_from_standard,
+                     format_channel_location_dict, format_region_channel_list)
+
+TUHTUEG_CHANNEL_LIST = [
+    'FP1', 'FP2', 'F3', 'F4', 'C3', 'C4', 'P3', 'P4', 'O1', 'O2', 'F7', 'F8',
+    'T3', 'T4', 'T5', 'T6', 'T1', 'T2', 'FZ', 'CZ', 'PZ'
+]
+
+TUHTUEG_LOCATION_LIST = [
+    ['-',  '-', '-',  'FP1', '-',  'FP2', '-',  '-', '-'],
+    ['-',  '-', '-',  '-',  '-',  '-',   '-',  '-', '-'],
+    ['F7', '-', 'F3', '-',  'FZ', '-',   'F4', '-', 'F8'],
+    ['T1', '-', '-',  '-',  '-',  '-',   '-',  '-', 'T2'],
+    ['T3', '-', 'C3', '-',  'CZ', '-',   'C4', '-', 'T4'],
+    ['-',  '-', '-',  '-',  '-',  '-',   '-',  '-', '-'],
+    ['T5', '-', 'P3', '-',  'PZ', '-',   'P4', '-', 'T6'],
+    ['-',  '-', '-',  '-',  '-',  '-',   '-',  '-', '-'],
+    ['-',  '-', '-',  'O1', '-',  'O2',  '-',  '-', '-']
+]
+
+TUHTUEG_CHANNEL_LOCATION_DICT = format_channel_location_dict(
+    TUHTUEG_CHANNEL_LIST, TUHTUEG_LOCATION_LIST)
+
+TUHTUEG_ADJACENCY_LIST = {
+    'FP1': ['FP2', 'F3', 'FZ'],
+    'FP2': ['FP1', 'FZ', 'F4'],
+    'F7': ['T1', 'T3'],
+    'F3': ['FP1', 'FZ', 'C3', 'CZ'],
+    'FZ': ['FP1', 'FP2', 'F3', 'F4', 'C3', 'CZ', 'C4'],
+    'F4': ['FP2', 'FZ', 'CZ', 'C4'],
+    'F8': ['T2', 'T4'],
+    'T1': ['F7', 'T3'],
+    'T2': ['F8', 'T4'],
+    'T3': ['F7', 'T1', 'T5'],
+    'C3': ['F3', 'FZ', 'CZ', 'P3', 'PZ'],
+    'CZ': ['F3', 'FZ', 'F4', 'C3', 'C4', 'P3', 'PZ', 'P4'],
+    'C4': ['FZ', 'F4', 'CZ', 'PZ', 'P4'],
+    'T4': ['F8', 'T2', 'T6'],
+    'T5': ['T3'],
+    'P3': ['C3', 'CZ', 'PZ', 'O1'],
+    'PZ': ['C3', 'CZ', 'C4', 'P3', 'P4', 'O1', 'O2'],
+    'P4': ['CZ', 'C4', 'PZ', 'O2'],
+    'T6': ['T4'],
+    'O1': ['P3', 'PZ', 'O2'],
+    'O2': ['PZ', 'P4', 'O1']
+}
+
+TUHTUEG_ADJACENCY_MATRIX = format_adj_matrix_from_adj_list(
+    TUHTUEG_CHANNEL_LIST, TUHTUEG_ADJACENCY_LIST)
+
+TUHTUEG_STANDARD_ADJACENCY_MATRIX = format_adj_matrix_from_standard(
+    TUHTUEG_CHANNEL_LIST, STANDARD_1005_CHANNEL_LOCATION_DICT)
+
+TUHTUEG_GENERAL_REGION_LIST = format_region_channel_list(
+    TUHTUEG_CHANNEL_LIST, GENERAL_REGION_LIST)
+TUHTUEG_FRONTAL_REGION_LIST = format_region_channel_list(
+    TUHTUEG_CHANNEL_LIST, FRONTAL_REGION_LIST)
+TUHTUEG_HEMISPHERE_REGION_LIST = format_region_channel_list(
+    TUHTUEG_CHANNEL_LIST, HEMISPHERE_REGION_LIST)
\ No newline at end of file
diff --git a/torcheeg/datasets/constants/motor_imagery/__init__.py b/torcheeg/datasets/constants/motor_imagery/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/torcheeg/datasets/constants/motor_imagery/physionet_mi.py b/torcheeg/datasets/constants/motor_imagery/physionet_mi.py
new file mode 100644
index 0000000..ec876f2
--- /dev/null
+++ b/torcheeg/datasets/constants/motor_imagery/physionet_mi.py
@@ -0,0 +1,110 @@
+from ..region_1020 import (FRONTAL_REGION_LIST, GENERAL_REGION_LIST,
+                           HEMISPHERE_REGION_LIST)
+from ..standard_1020 import STANDARD_1020_CHANNEL_LOCATION_DICT
+from ..utils import (format_adj_matrix_from_adj_list,
+                     format_adj_matrix_from_standard,
+                     format_channel_location_dict, format_region_channel_list)
+
+PHYSIONETMI_CHANNEL_LIST = [
+    'FC5', 'FC3', 'FC1', 'FCZ', 'FC2', 'FC4', 'FC6', 'C5', 'C3', 'C1', 'CZ',
+    'C2', 'C4', 'C6', 'CP5', 'CP3', 'CP1', 'CPZ', 'CP2', 'CP4', 'CP6', 'FP1',
+    'FPZ', 'FP2', 'AF7', 'AF3', 'AFZ', 'AF4', 'AF8', 'F7', 'F5', 'F3', 'F1',
+    'FZ', 'F2', 'F4', 'F6', 'F8', 'FT7', 'FT8', 'T7', 'T8', 'T9', 'T10', 'TP7',
+    'TP8', 'P7', 'P5', 'P3', 'P1', 'PZ', 'P2', 'P4', 'P6', 'P8', 'PO7', 'PO3',
+    'POZ', 'PO4', 'PO8', 'O1', 'OZ', 'O2', 'IZ'
+]
+
+PHYSIONETMI_LOCATION_LIST = [
+    ['-', '-', '-', 'FP1', 'FPZ', 'FP2', '-', '-', '-'],
+    ['-', 'AF7', '-', 'AF3', 'AFZ', 'AF4', '-', 'AF8', '-'],
+    ['F7', 'F5', 'F3', 'F1', 'FZ', 'F2', 'F4', 'F6', 'F8'],
+    ['FT7', 'FC5', 'FC3', 'FC1', 'FCZ', 'FC2', 'FC4', 'FC6', 'FT8'],
+    ['T7', 'C5', 'C3', 'C1', 'CZ', 'C2', 'C4', 'C6', 'T8'],
+    ['TP7', 'CP5', 'CP3', 'CP1', 'CPZ', 'CP2', 'CP4', 'CP6', 'TP8'],
+    ['P7', 'P5', 'P3', 'P1', 'PZ', 'P2', 'P4', 'P6', 'P8'],
+    ['-', 'PO7', '-', 'PO3', 'POZ', 'PO4', '-', 'PO8', '-'],
+    ['-', '-', '-', 'O1', 'OZ', 'O2', '-', '-', '-']
+]
+# without {'T10', 'T9', 'IZ'}
+
+PHYSIONETMI_CHANNEL_LOCATION_DICT = format_channel_location_dict(
+    PHYSIONETMI_CHANNEL_LIST, PHYSIONETMI_LOCATION_LIST)
+
+PHYSIONETMI_ADJACENCY_LIST = {
+    'FP1': ['AF3', 'FPZ', 'AFZ'],
+    'FPZ': ['AFZ', 'FP1', 'FP2', 'AF4', 'AF3'],
+    'FP2': ['AF4', 'FPZ', 'AFZ'],
+    'AF7': ['F5', 'F3', 'F7'],
+    'AF3': ['FP1', 'F1', 'AFZ', 'FPZ', 'FZ', 'F3'],
+    'AFZ': ['FPZ', 'FZ', 'AF3', 'AF4', 'FP2', 'F2', 'FP1', 'F1'],
+    'AF4': ['FP2', 'F2', 'AFZ', 'F4', 'FPZ', 'FZ'],
+    'AF8': ['F6', 'F8', 'F4'],
+    'F7': ['FT7', 'F5', 'AF7', 'FC5'],
+    'F5': ['AF7', 'FC5', 'F7', 'F3', 'FC3', 'FT7'],
+    'F3': ['FC3', 'F5', 'F1', 'AF3', 'FC1', 'AF7', 'FC5'],
+    'F1': ['AF3', 'FC1', 'F3', 'FZ', 'AFZ', 'FCZ', 'FC3'],
+    'FZ': ['AFZ', 'FCZ', 'F1', 'F2', 'AF4', 'FC2', 'AF3', 'FC1'],
+    'F2': ['AF4', 'FC2', 'FZ', 'F4', 'FC4', 'AFZ', 'FCZ'],
+    'F4': ['FC4', 'F2', 'F6', 'AF8', 'FC6', 'AF4', 'FC2'],
+    'F6': ['AF8', 'FC6', 'F4', 'F8', 'FT8', 'FC4'],
+    'F8': ['FT8', 'F6', 'AF8', 'FC6'],
+    'FT7': ['F7', 'T7', 'FC5', 'F5', 'C5'],
+    'FC5': ['F5', 'C5', 'FT7', 'FC3', 'F3', 'C3', 'F7', 'T7'],
+    'FC3': ['F3', 'C3', 'FC5', 'FC1', 'F1', 'C1', 'F5', 'C5'],
+    'FC1': ['F1', 'C1', 'FC3', 'FCZ', 'FZ', 'CZ', 'F3', 'C3'],
+    'FCZ': ['FZ', 'CZ', 'FC1', 'FC2', 'F2', 'C2', 'F1', 'C1'],
+    'FC2': ['F2', 'C2', 'FCZ', 'FC4', 'F4', 'C4', 'FZ', 'CZ'],
+    'FC4': ['F4', 'C4', 'FC2', 'FC6', 'F6', 'C6', 'F2', 'C2'],
+    'FC6': ['F6', 'C6', 'FC4', 'FT8', 'F8', 'T8', 'F4', 'C4'],
+    'FT8': ['F8', 'T8', 'FC6', 'F6', 'C6'],
+    'T7': ['FT7', 'TP7', 'C5', 'FC5', 'CP5'],
+    'C5': ['FC5', 'CP5', 'T7', 'C3', 'FC3', 'CP3', 'FT7', 'TP7'],
+    'C3': ['FC3', 'CP3', 'C5', 'C1', 'FC1', 'CP1', 'FC5', 'CP5'],
+    'C1': ['FC1', 'CP1', 'C3', 'CZ', 'FCZ', 'CPZ', 'FC3', 'CP3'],
+    'CZ': ['FCZ', 'CPZ', 'C1', 'C2', 'FC2', 'CP2', 'FC1', 'CP1'],
+    'C2': ['FC2', 'CP2', 'CZ', 'C4', 'FC4', 'CP4', 'FCZ', 'CPZ'],
+    'C4': ['FC4', 'CP4', 'C2', 'C6', 'FC6', 'CP6', 'FC2', 'CP2'],
+    'C6': ['FC6', 'CP6', 'C4', 'T8', 'FT8', 'TP8', 'FC4', 'CP4'],
+    'T8': ['FT8', 'TP8', 'C6', 'FC6', 'CP6'],
+    'TP7': ['T7', 'P7', 'CP5', 'C5', 'P5'],
+    'CP5': ['C5', 'P5', 'TP7', 'CP3', 'C3', 'P3', 'T7', 'P7'],
+    'CP3': ['C3', 'P3', 'CP5', 'CP1', 'C1', 'P1', 'C5', 'P5'],
+    'CP1': ['C1', 'P1', 'CP3', 'CPZ', 'CZ', 'PZ', 'C3', 'P3'],
+    'CPZ': ['CZ', 'PZ', 'CP1', 'CP2', 'C2', 'P2', 'C1', 'P1'],
+    'CP2': ['C2', 'P2', 'CPZ', 'CP4', 'C4', 'P4', 'CZ', 'PZ'],
+    'CP4': ['C4', 'P4', 'CP2', 'CP6', 'C6', 'P6', 'C2', 'P2'],
+    'CP6': ['C6', 'P6', 'CP4', 'TP8', 'T8', 'P8', 'C4', 'P4'],
+    'TP8': ['T8', 'P8', 'CP6', 'C6', 'P6'],
+    'P7': ['TP7', 'P5', 'CP5', 'PO7'],
+    'P5': ['CP5', 'PO7', 'P7', 'P3', 'CP3', 'TP7'],
+    'P3': ['CP3', 'P5', 'P1', 'CP1', 'PO3', 'CP5', 'PO7'],
+    'P1': ['CP1', 'PO3', 'P3', 'PZ', 'CPZ', 'POZ', 'CP3'],
+    'PZ': ['CPZ', 'POZ', 'P1', 'P2', 'CP2', 'PO4', 'CP1', 'PO3'],
+    'P2': ['CP2', 'PO4', 'PZ', 'P4', 'CP4', 'CPZ', 'POZ'],
+    'P4': ['CP4', 'P2', 'P6', 'CP6', 'PO8', 'CP2', 'PO4'],
+    'P6': ['CP6', 'PO8', 'P4', 'P8', 'TP8', 'CP4'],
+    'P8': ['TP8', 'P6', 'CP6', 'PO8'],
+    'PO7': ['P5', 'P3', 'P7'],
+    'PO3': ['P1', 'O1', 'POZ', 'PZ', 'OZ', 'P3'],
+    'POZ': ['PZ', 'OZ', 'PO3', 'PO4', 'P2', 'O2', 'P1', 'O1'],
+    'PO4': ['P2', 'O2', 'POZ', 'P4', 'PZ', 'OZ'],
+    'PO8': ['P6', 'P8', 'P4'],
+    'O1': ['PO3', 'OZ', 'POZ'],
+    'OZ': ['POZ', 'O1', 'O2', 'PO4', 'PO3'],
+    'O2': ['PO4', 'OZ', 'POZ']
+}
+
+PHYSIONETMI_ADJACENCY_MATRIX = format_adj_matrix_from_adj_list(
+    PHYSIONETMI_CHANNEL_LIST, PHYSIONETMI_ADJACENCY_LIST)
+
+PHYSIONETMI_STANDARD_ADJACENCY_MATRIX = format_adj_matrix_from_standard(
+    PHYSIONETMI_CHANNEL_LIST,
+    STANDARD_1020_CHANNEL_LOCATION_DICT,
+    delta=0.00035)
+
+PHYSIONETMI_GENERAL_REGION_LIST = format_region_channel_list(
+    PHYSIONETMI_CHANNEL_LIST, GENERAL_REGION_LIST)
+PHYSIONETMI_FRONTAL_REGION_LIST = format_region_channel_list(
+    PHYSIONETMI_CHANNEL_LIST, FRONTAL_REGION_LIST)
+PHYSIONETMI_HEMISPHERE_REGION_LIST = format_region_channel_list(
+    PHYSIONETMI_CHANNEL_LIST, HEMISPHERE_REGION_LIST)
diff --git a/torcheeg/datasets/constants/utils.py b/torcheeg/datasets/constants/utils.py
index 7d675d7..6ce96a4 100644
--- a/torcheeg/datasets/constants/utils.py
+++ b/torcheeg/datasets/constants/utils.py
@@ -25,6 +25,8 @@ def format_channel_location_dict(channel_list, location_list):
         if len(np.argwhere(location_list == channel)):
             location = (np.argwhere(location_list == channel)[0]).tolist()
             output[channel] = location
+        else:
+            output[channel] = None
     return output
 
 
diff --git a/torcheeg/datasets/functional/hooks.py b/torcheeg/datasets/functional/hooks.py
index bae0221..d271398 100644
--- a/torcheeg/datasets/functional/hooks.py
+++ b/torcheeg/datasets/functional/hooks.py
@@ -1,7 +1,7 @@
 import numpy as np
 
 
-def before_trial_normalize(data: np.ndarray, eps: float = 1e-6):
+def before_trial_normalize(data: np.ndarray, eps: float = 1e-6, axis=0):
     r'''
     A common hook function used to normalize the signal of the whole trial before dividing it into chunks.
 
@@ -32,8 +32,8 @@ def before_trial_normalize(data: np.ndarray, eps: float = 1e-6):
     Returns:
         np.ndarray: The normalized results of a trial.
     '''
-    min_v = data.min(axis=-1, keepdims=True)
-    max_v = data.max(axis=-1, keepdims=True)
+    min_v = data.min(axis=axis, keepdims=True)
+    max_v = data.max(axis=axis, keepdims=True)
     return (data - min_v) / (max_v - min_v + eps)
 
 
diff --git a/torcheeg/datasets/functional/operators.py b/torcheeg/datasets/functional/operators.py
index 58319e9..c07f307 100644
--- a/torcheeg/datasets/functional/operators.py
+++ b/torcheeg/datasets/functional/operators.py
@@ -1,3 +1,4 @@
+import logging
 import math
 import os
 import shutil
@@ -9,8 +10,10 @@
 
 from torcheeg.io import EEGSignalIO, MetaInfoIO
 
+log = logging.getLogger(__name__)
 
-def _set_files(**kwargs):
+
+def set_records(self, **kwargs):
     root_path = kwargs.pop('root_path', '.')  # str
     num_samples_per_worker = kwargs.pop('num_samples_per_worker', '.')  # str
 
@@ -26,12 +29,12 @@ def _set_files(**kwargs):
     return block_list
 
 
-def _load_data(block,
-           io_path: str = None,
-           io_size: int = 10485760,
-           io_mode: str = 'lmdb',
-           lock: Any = None,
-           **kwargs):
+def process_record(block,
+                   io_path: str = None,
+                   io_size: int = 10485760,
+                   io_mode: str = 'lmdb',
+                   lock: Any = None,
+                   **kwargs):
     start_id, end_id = block
 
     transform = kwargs.pop('transform', None)
@@ -162,21 +165,23 @@ def from_existing(dataset: Any,
 
         if num_worker == 0:
             lock = MockLock()  # do nothing, just for compatibility
-            for block in tqdm(_set_files(**params),
+            for block in tqdm(set_records(**params),
                               disable=not verbose,
                               desc="[PROCESS]"):
-                _load_data(block=block, lock=lock, **params)
+                process_record(block=block, lock=lock, **params)
         else:
             # lock for lmdb writter, LMDB only allows single-process writes
             manager = Manager()
             lock = manager.Lock()
 
             Parallel(n_jobs=num_worker)(
-                delayed(_load_data)(block=block, lock=lock, **params)
-                for block in tqdm(
-                    _set_files(**params), disable=not verbose, desc="[PROCESS]"))
+                delayed(process_record)(block=block, lock=lock, **params)
+                for block in tqdm(set_records(**params),
+                                  disable=not verbose,
+                                  desc="[PROCESS]"))
     else:
-        print(f'dataset already exists at path {io_path}, reading from path...')
+        print(
+            f'dataset already exists at path {io_path}, reading from path...')
 
     return type(dataset)(io_path=io_path,
                          offline_transform=transform,
diff --git a/torcheeg/datasets/moabb/moabb_dataset.py b/torcheeg/datasets/moabb/moabb_dataset.py
index 350c50c..e39321e 100644
--- a/torcheeg/datasets/moabb/moabb_dataset.py
+++ b/torcheeg/datasets/moabb/moabb_dataset.py
@@ -1,12 +1,17 @@
 import os
+import mne
+import logging
+import numpy as np
 import itertools
 from typing import Any, Callable, Dict, Tuple, Union
+import pandas as pd
 
 from ..module.base_dataset import BaseDataset
 
 from moabb.datasets.base import BaseDataset as _MOABBDataset
 from moabb.paradigms.base import BaseParadigm as _MOABBParadigm
 
+log = logging.getLogger(__name__)
 
 class MOABBDataset(BaseDataset):
     '''
@@ -34,24 +39,36 @@ class MOABBDataset(BaseDataset):
     Args:
         dataset (MOABBDataset): an instance of :obj:`BaseDataset` defined by moabb.
         paradigm (MOABBParadigm): an instance of :obj:`BaseParadigm` defined by moabb.
+        chunk_size (int): The length of each EEG sample. If set to :obj:`-1`, the length of each EEG sample is the same as the length of the EEG signal in event. (default: :obj:`-1`)
+        overlap (int): The overlap between two adjacent EEG samples. (default: :obj:`0`)
         online_transform (Callable, optional): The transformation of the EEG signals and baseline EEG signals. The input is a :obj:`np.ndarray`, and the ouput is used as the first and second value of each element in the dataset. (default: :obj:`None`)
         offline_transform (Callable, optional): The usage is the same as :obj:`online_transform`, but executed before generating IO intermediate results. (default: :obj:`None`)
         label_transform (Callable, optional): The transformation of the label. The input is an information dictionary, and the ouput is used as the third value of each element in the dataset. (default: :obj:`None`)
+        before_trial (Callable, optional): The hook performed on the trial to which the sample belongs. It is performed before the offline transformation and thus typically used to implement context-dependent sample transformations, such as filter, reference, moving average and etc. Its input is a :obj:`mne.io.Raw` object, and its output should also be a :obj:`mne.io.Raw` object. (default: :obj:`None`)
+        after_trial (Callable, optional): The hook performed on the trial to which the sample belongs. It is performed after the offline transformation and thus typically used to implement context-dependent sample transformations, such as moving averages, etc. The input and output of this hook function should be a sequence of dictionaries representing a sequence of EEG samples. Each dictionary contains two key-value pairs, indexed by :obj:`eeg` (the EEG signal matrix) and :obj:`key` (the index in the database) respectively.
         io_path (str): The path to generated unified data IO, cached as an intermediate result. (default: :obj:`./io/moabb`)
         io_size (int): Maximum size database may grow to; used to size the memory mapping. If database grows larger than ``map_size``, an exception will be raised and the user must close and reopen. (default: :obj:`10485760`)
         io_mode (str): Storage mode of EEG signal. When io_mode is set to :obj:`lmdb`, TorchEEG provides an efficient database (LMDB) for storing EEG signals. LMDB may not perform well on limited operating systems, where a file system based EEG signal storage is also provided. When io_mode is set to :obj:`pickle`, pickle-based persistence files are used. (default: :obj:`lmdb`)
+        download_path (str): The path to download raw data. If set to :obj:`None`, the raw data will be downloaded to :obj:`f'{self.io_path}/raw'`. If the path already exists, the download will be skipped. (default: :obj:`None`)
         num_worker (int): Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process. (default: :obj:`0`)
         in_memory (bool): Whether to load the entire dataset into memory. If :obj:`in_memory` is set to True, then the first time an EEG sample is read, the entire dataset is loaded into memory for subsequent retrieval. Otherwise, the dataset is stored on disk to avoid the out-of-memory problem. (default: :obj:`False`)
     '''
     def __init__(self,
                  dataset: _MOABBDataset,
                  paradigm: _MOABBParadigm,
+                 chunk_size: int = -1,
+                 overlap: int = 0,
                  online_transform: Union[None, Callable] = None,
                  offline_transform: Union[None, Callable] = None,
                  label_transform: Union[None, Callable] = None,
+                 before_trial: Union[None, Callable] = None,
+                 after_trial: Union[Callable, None] = None,
+                 after_session: Union[Callable, None] = None,
+                 after_subject: Union[Callable, None] = None,
                  io_path: str = './io/moabb',
                  io_size: int = 10485760,
                  io_mode: str = 'lmdb',
+                 download_path: str = None,
                  num_worker: int = 0,
                  verbose: bool = True,
                  in_memory: bool = False,
@@ -60,31 +77,43 @@ def __init__(self,
         params = {
             'dataset': dataset,
             'paradigm': paradigm,
+            'chunk_size': chunk_size,
+            'overlap': overlap,
             'online_transform': online_transform,
             'offline_transform': offline_transform,
             'label_transform': label_transform,
+            'before_trial': before_trial,
+            'after_trial': after_trial,
+            'after_session': after_session,
+            'after_subject': after_subject,
             'io_path': io_path,
             'io_size': io_size,
             'io_mode': io_mode,
+            'download_path': download_path,
             'num_worker': num_worker,
             'verbose': verbose,
             'in_memory': in_memory
         }
         params.update(kwargs)
+
+        if not paradigm.is_valid(dataset):
+            message = f"Dataset {dataset.code} is not valid for paradigm"
+            raise AssertionError(message)
+        
         super().__init__(**params)
         # save all arguments to __dict__
         self.__dict__.update(params)
 
     @staticmethod
-    def _load_data(file: Any = None,
-                   offline_transform: Union[None, Callable] = None,
-                   dataset: _MOABBDataset = None,
-                   paradigm: _MOABBParadigm = None,
-                   **kwargs):
-
-        if not paradigm.is_valid(dataset):
-            message = f"Dataset {dataset.code} is not valid for paradigm"
-            raise AssertionError(message)
+    def process_record(file: Any = None,
+                       chunk_size: int = -1,
+                       overlap: int = 0,
+                       offline_transform: Union[None, Callable] = None,
+                       dataset: _MOABBDataset = None,
+                       paradigm: _MOABBParadigm = None,
+                       before_trial: Union[None, Callable] = None,
+                       after_trial: Union[None, Callable] = None,
+                       **kwargs):
 
         subject_id, session_id = file
 
@@ -93,6 +122,8 @@ def _load_data(file: Any = None,
             subjects=[subject_id])[subject_id][session_id]
 
         for run_id, run_signal in session_signal.items():
+            if before_trial is not None:
+                run_signal = before_trial(run_signal)
 
             proc = paradigm.process_raw(run_signal,
                                         dataset,
@@ -102,36 +133,75 @@ def _load_data(file: Any = None,
                 # go to next
                 continue
 
-            clip_signals, labels, metadatas = proc
-            for clip_id, (clip_signal, label, metadata) in enumerate(
-                    zip(clip_signals, labels, metadatas.iterrows())):
-                t_eeg = clip_signal
-                if not offline_transform is None:
-                    t = offline_transform(eeg=clip_signal)
-                    t_eeg = t['eeg']
-
-                clip_id = f'{subject_id}_{session_id}_{write_pointer}'
-                write_pointer += 1
-
-                record_info = {
-                    'subject_id': subject_id,
-                    'session_id': session_id,
-                    'run_id': run_id,
-                    'clip_id': clip_id,
-                    'label': label,
-                    'metadata': metadata[1].to_dict()
-                }
-
-                yield {'eeg': t_eeg, 'key': clip_id, 'info': record_info}
-
-    @staticmethod
-    def _set_files(dataset: _MOABBDataset,
-                   io_path: str = './io/moabb',
-                   **kwargs):
+            trial_queue = []
+            roi_signals, labels, metadatas = proc
+            for roi_id, (roi_signal, label, metadata) in enumerate(
+                    zip(roi_signals, labels, metadatas.iterrows())):
+
+                start_at = 0
+                if chunk_size <= 0:
+                    chunk_size = roi_signal.shape[1] - start_at
+                end_at = start_at + chunk_size
+
+                step = chunk_size - overlap
+
+                while end_at <= roi_signal.shape[1]:
+                    clip_sample = roi_signal[:, start_at:end_at]
+
+                    t_eeg = clip_sample
+                    if not offline_transform is None:
+                        t = offline_transform(eeg=t_eeg)
+                        t_eeg = t['eeg']
+
+                    clip_id = f'{subject_id}_{session_id}_{write_pointer}'
+                    write_pointer += 1
+
+                    record_info = {
+                        'subject_id': subject_id,
+                        'session_id': session_id,
+                        'trial_id': run_id,
+                        'roi_id': roi_id,
+                        'clip_id': clip_id,
+                        'label': label,
+                        'start_at': start_at,
+                        'end_at': end_at,
+                        'metadata': metadata[1].to_dict()
+                    }
+                    write_pointer += 1
+                    start_at = start_at + step
+                    end_at = start_at + chunk_size
+
+                    if after_trial is not None:
+                        trial_queue.append({
+                            'eeg': t_eeg,
+                            'key': clip_id,
+                            'info': record_info
+                        })
+                    else:
+                        yield {
+                            'eeg': t_eeg,
+                            'key': clip_id,
+                            'info': record_info
+                        }
+
+            if len(trial_queue) and after_trial is not None:
+                trial_queue = after_trial(trial_queue)
+                for obj in trial_queue:
+                    assert 'eeg' in obj and 'key' in obj and 'info' in obj, 'after_trial must return a list of dictionaries, where each dictionary corresponds to an EEG sample, containing `eeg`, `key` and `info` as keys.'
+                    yield obj
+
+    def set_records(self,
+                    dataset: _MOABBDataset,
+                    io_path: str = './io/moabb',
+                    download_path: str = None,
+                    **kwargs):
         subject_id_list = dataset.subject_list
-        dataset.download(subject_list=subject_id_list,
-                         path=os.path.join(io_path, 'raw'),
-                         verbose=False)
+        if download_path is None:
+            download_path = os.path.join(io_path, 'raw')
+        if not os.path.exists(download_path):
+            dataset.download(subject_list=subject_id_list,
+                             path=download_path,
+                             verbose=False)
 
         subject_id = subject_id_list[0]
         session_id_list = list(
@@ -146,7 +216,8 @@ def __getitem__(self, index: int) -> Tuple[any, any, int, int, int]:
         info = self.read_info(index)
 
         eeg_index = str(info['clip_id'])
-        eeg = self.read_eeg(eeg_index)
+        eeg_record = str(info['_record_id'])
+        eeg = self.read_eeg(eeg_record, eeg_index)
 
         signal = eeg
         label = info
diff --git a/torcheeg/datasets/module/__init__.py b/torcheeg/datasets/module/__init__.py
index ef5d8a3..f77f3e0 100644
--- a/torcheeg/datasets/module/__init__.py
+++ b/torcheeg/datasets/module/__init__.py
@@ -1,7 +1,6 @@
 from .base_dataset import BaseDataset
 from .emotion_recognition import *
 from .folder_dataset import FolderDataset
-from .mne_dataset import MNEDataset
-from .numpy_dataset import NumpyDataset
+from .csv_folder_dataset import CSVFolderDataset
 from .personal_identification import *
-from .sspev import *
+from .sspev import *
\ No newline at end of file
diff --git a/torcheeg/datasets/module/base_dataset.py b/torcheeg/datasets/module/base_dataset.py
index 5d9878c..a428057 100644
--- a/torcheeg/datasets/module/base_dataset.py
+++ b/torcheeg/datasets/module/base_dataset.py
@@ -1,8 +1,10 @@
+import logging
 import os
 import shutil
-from multiprocessing import Manager
-from typing import Any, Dict
-
+from typing import Any, Callable, Dict
+import torch
+import numpy as np
+import pandas as pd
 from joblib import Parallel, delayed
 from torch.utils.data import Dataset
 from tqdm import tqdm
@@ -11,6 +13,8 @@
 
 MAX_QUEUE_SIZE = 1024
 
+log = logging.getLogger(__name__)
+
 
 class MockLock():
     def __enter__(self):
@@ -21,7 +25,6 @@ def __exit__(self, exc_type, exc_value, traceback):
 
 
 class BaseDataset(Dataset):
-
     def __init__(self,
                  io_path: str = None,
                  io_size: int = 10485760,
@@ -29,14 +32,19 @@ def __init__(self,
                  in_memory: bool = False,
                  num_worker: int = 0,
                  verbose: bool = True,
+                 after_trial: Callable = None,
+                 after_session: Callable = None,
+                 after_subject: Callable = None,
                  **kwargs):
-
         self.io_path = io_path
         self.io_size = io_size
         self.io_mode = io_mode
         self.in_memory = in_memory
         self.num_worker = num_worker
         self.verbose = verbose
+        self.after_trial = after_trial
+        self.after_session = after_session
+        self.after_subject = after_subject
 
         # new IO
         if not self.exist(self.io_path):
@@ -46,73 +54,68 @@ def __init__(self,
             # make the root dictionary
             os.makedirs(self.io_path, exist_ok=True)
 
-            # init sub-folders
-            meta_info_io_path = os.path.join(self.io_path, 'info.csv')
-            eeg_signal_io_path = os.path.join(self.io_path, 'eeg')
-
-            MetaInfoIO(meta_info_io_path)
-            EEGSignalIO(eeg_signal_io_path,
-                        io_size=self.io_size,
-                        io_mode=self.io_mode)
-
+            records = self.set_records(**kwargs)
             if self.num_worker == 0:
-                lock = MockLock()  # do nothing, just for compatibility
-                # if catch error, then delete the database
                 try:
-                    for file in tqdm(self._set_files(**kwargs),
-                                    disable=not self.verbose,
-                                    desc="[PROCESS]"):
-                        self._process_file(io_path=self.io_path,
-                                        io_size=self.io_size,
-                                        io_mode=self.io_mode,
-                                        file=file,
-                                        lock=lock,
-                                        _load_data=self._load_data,
-                                        **kwargs)
+                    for file_id, file in tqdm(enumerate(records),
+                                              disable=not self.verbose,
+                                              desc="[PROCESS]",
+                                              total=len(records)):
+                        self.save_record(io_path=self.io_path,
+                                         io_size=self.io_size,
+                                         io_mode=self.io_mode,
+                                         file=file,
+                                         file_id=file_id,
+                                         process_record=self.process_record,
+                                         **kwargs)
                 except Exception as e:
                     # shutil to delete the database
                     shutil.rmtree(self.io_path)
                     raise e
             else:
-                # lock for lmdb writter, LMDB only allows single-process writes
-                manager = Manager()
-                lock = manager.Lock()
-
-                # if catch error, then delete the database
+                # catch the exception
                 try:
                     Parallel(n_jobs=self.num_worker)(
-                        delayed(self._process_file)(io_path=io_path,
-                                                    io_size=io_size,
-                                                    io_mode=io_mode,
-                                                    file=file,
-                                                    lock=lock,
-                                                    _load_data=self._load_data,
-                                                    **
-                                                    kwargs)
-                        for file in tqdm(self._set_files(**kwargs),
-                                        disable=not self.verbose,
-                                        desc="[PROCESS]"))
+                        delayed(self.save_record)(
+                            io_path=io_path,
+                            io_size=io_size,
+                            io_mode=io_mode,
+                            file_id=file_id,
+                            file=file,
+                            process_record=self.process_record,
+                            **kwargs)
+                        for file_id, file in tqdm(enumerate(records),
+                                                  disable=not self.verbose,
+                                                  desc="[PROCESS]",
+                                                  total=len(records)))
                 except Exception as e:
                     # shutil to delete the database
                     shutil.rmtree(self.io_path)
                     raise e
 
-        print(
-            f'dataset already exists at path {self.io_path}, reading from path...'
-        )
-
-        meta_info_io_path = os.path.join(self.io_path, 'info.csv')
-        eeg_signal_io_path = os.path.join(self.io_path, 'eeg')
-
-        info_io = MetaInfoIO(meta_info_io_path)
-        self.eeg_io = EEGSignalIO(eeg_signal_io_path,
-                                  io_size=self.io_size,
-                                  io_mode=self.io_mode)
-
-        self.info = info_io.read_all()
+            # get the global io
+            self.eeg_io_router, self.info = self.get_pointer(io_path=io_path,
+                                                             io_size=io_size,
+                                                             io_mode=io_mode)
+            # catch the exception
+            try:
+                self.post_process_record(after_trial=after_trial,
+                                         after_session=after_session,
+                                         after_subject=after_subject)
+            except Exception as e:
+                # shutil to delete the database
+                shutil.rmtree(self.io_path)
+                raise e
+        else:
+            print(
+                f'dataset already exists at path {self.io_path}, reading from path...'
+            )
+            # get the global io
+            self.eeg_io_router, self.info = self.get_pointer(io_path=io_path,
+                                                             io_size=io_size,
+                                                             io_mode=io_mode)
 
-    @staticmethod
-    def _set_files(**kwargs):
+    def set_records(self, **kwargs):
         '''
         The block method for generating the database. It is used to describe which data blocks need to be processed to generate the database. It is called in parallel by :obj:`joblib.Parallel` in :obj:`__init__` of the class.
 
@@ -122,60 +125,95 @@ def _set_files(**kwargs):
 
         .. code-block:: python
 
-            def _set_files(root_path: str = None, **kwargs):
-                # e.g., return file name list for _load_data to process
+        def set_records(self, root_path: str = None, **kwargs):
+                # e.g., return file name list for process_record to process
                 return os.listdir(root_path)
 
         '''
         raise NotImplementedError(
-            "Method _set_files is not implemented in class BaseDataset")
+            "Method set_records is not implemented in class BaseDataset")
 
-    @staticmethod
-    def _process_file(io_path: str = None,
-                      io_size: int = 10485760,
-                      io_mode: str = 'lmdb',
-                      file: Any = None,
-                      lock: Any = None,
-                      _load_data=None,
-                      **kwargs):
+    def get_pointer(self, io_path: str, io_size: int, io_mode: str):
+        '''
+        The method for initializing the database. It is used to describe how to initialize the database. It is called in :obj:`__init__` of the class.
+        '''
+        # get all records
+        records = os.listdir(io_path)
+        # filter the records with the prefix '_record_'
+        records = list(filter(lambda x: '_record_' in x, records))
+
+        # for every record, get the io_path, and init the info_io and eeg_io
+        eeg_io_router = {}
+        info_merged = None
+
+        for record in records:
+            meta_info_io_path = os.path.join(io_path, record, 'info.csv')
+            eeg_signal_io_path = os.path.join(io_path, record, 'eeg')
+            info_io = MetaInfoIO(meta_info_io_path)
+            eeg_io = EEGSignalIO(eeg_signal_io_path,
+                                 io_size=io_size,
+                                 io_mode=io_mode)
+            eeg_io_router[record] = eeg_io
+            info_df = info_io.read_all()
+
+            assert '_record_id' not in info_df.columns, \
+                "column '_record_id' is a forbidden reserved word and is used to index the corresponding IO. Please replace your '_record_id' with another name."
+            info_df['_record_id'] = record
+
+            if info_merged is None:
+                info_merged = info_df
+            else:
+                info_merged = pd.concat([info_merged, info_df],
+                                        ignore_index=True)
+        return eeg_io_router, info_merged
 
-        meta_info_io_path = os.path.join(io_path, 'info.csv')
-        eeg_signal_io_path = os.path.join(io_path, 'eeg')
+    @staticmethod
+    def save_record(io_path: str = None,
+                    io_size: int = 10485760,
+                    io_mode: str = 'lmdb',
+                    file: Any = None,
+                    file_id: int = None,
+                    process_record=None,
+                    **kwargs):
+        _record_id = str(file_id)
+        meta_info_io_path = os.path.join(io_path, f'_record_{_record_id}',
+                                         'info.csv')
+        eeg_signal_io_path = os.path.join(io_path, f'_record_{_record_id}',
+                                          'eeg')
 
         info_io = MetaInfoIO(meta_info_io_path)
         eeg_io = EEGSignalIO(eeg_signal_io_path,
                              io_size=io_size,
                              io_mode=io_mode)
 
-        gen = _load_data(file=file, **kwargs)
-        # loop for data yield by _load_data, until to the end of the data
+        gen = process_record(file=file, **kwargs)
+        # loop for data yield by process_record, until to the end of the data
         while True:
             try:
-                # call _load_data of the class
+                # call process_record of the class
                 # get the current class name
                 obj = next(gen)
 
             except StopIteration:
                 break
 
-            with lock:
-                if 'eeg' in obj and 'key' in obj:
-                    eeg_io.write_eeg(obj['eeg'], obj['key'])
-                if 'info' in obj:
-                    info_io.write_info(obj['info'])
+            if 'eeg' in obj and 'key' in obj:
+                eeg_io.write_eeg(obj['eeg'], obj['key'])
+            if 'info' in obj:
+                info_io.write_info(obj['info'])
 
     @staticmethod
-    def _load_data(file: Any = None, **kwargs):
+    def process_record(file: Any = None, **kwargs):
         '''
         The IO method for generating the database. It is used to describe how files are processed to generate the database. It is called in parallel by :obj:`joblib.Parallel` in :obj:`__init__` of the class.
 
         Args:
-            file (Any): The file to be processed. It is an element in the list returned by _set_files. (default: :obj:`Any`)
+            file (Any): The file to be processed. It is an element in the list returned by set_records. (default: :obj:`Any`)
             **kwargs: The arguments derived from :obj:`__init__` of the class.
 
         .. code-block:: python
 
-            def _load_data(file: Any = None, chunk_size: int = 128, **kwargs):
+            def process_record(file: Any = None, chunk_size: int = 128, **kwargs):
                 # process file
                 eeg = np.ndarray((chunk_size, 64, 128), dtype=np.float32)
                 key = '1'
@@ -194,21 +232,161 @@ def _load_data(file: Any = None, chunk_size: int = 128, **kwargs):
         '''
 
         raise NotImplementedError(
-            "Method _load_data is not implemented in class BaseDataset")
+            "Method process_record is not implemented in class BaseDataset")
 
-    def read_eeg(self, key: str) -> Any:
+    def post_process_record(self,
+                            after_trial: Callable = None,
+                            after_session: Callable = None,
+                            after_subject: Callable = None):
+        '''
+        The hook method for post-processing the data. It is used to describe how to post-process the data.
+        '''
+        pbar = tqdm(total=len(self),
+                    disable=not self.verbose,
+                    desc="[POST-PROCESS]")
+
+        # if all the hooks are None, then return
+        if after_trial is None and after_session is None and after_subject is None:
+            return
+
+        if 'subject_id' in self.info.columns:
+            subject_df = self.info.groupby('subject_id')
+        else:
+            subject_df = [(None, self.info)]
+            if not after_subject is None:
+                print(
+                    "No subject_id column found in info, after_subject hook is ignored."
+                )
+        if after_subject is None:
+            after_subject = lambda x: x
+
+        for _, subject_info in subject_df:
+
+            subject_record_list = []
+            subject_index_list = []
+            subject_samples = []
+
+            # check if have a session_id column
+            if 'session_id' in subject_info.columns:
+                session_df = subject_info.groupby('session_id')
+            else:
+                session_df = [(None, subject_info)]
+                if not after_session is None:
+                    print(
+                        "No session_id column found in info, after_session hook is ignored."
+                    )
+            if after_session is None:
+                after_session = lambda x: x
+
+            for _, session_info in session_df:
+
+                # check if have a trial_id column
+                if 'trial_id' in session_info.columns:
+                    trial_df = session_info.groupby('trial_id')
+                else:
+                    trial_df = [(None, session_info)]
+                    if not after_trial is None:
+                        print(
+                            "No trial_id column found in info, after_trial hook is ignored."
+                        )
+                if after_trial is None:
+                    after_trial = lambda x: x
+
+                session_samples = []
+                for _, trial_info in trial_df:
+                    trial_samples = []
+                    for i in range(len(trial_info)):
+                        eeg_index = str(trial_info.iloc[i]['clip_id'])
+                        eeg_record = str(trial_info.iloc[i]['_record_id'])
+
+                        # if i == 0 and 'baseline_id' in trial_info.columns:
+                        #     baseline_index = str(
+                        #         trial_info.iloc[i]['baseline_id'])
+                        #     subject_record_list.append(eeg_record)
+                        #     subject_index_list.append(baseline_index)
+
+                        #     eeg = self.read_eeg(eeg_record, baseline_index)
+                        #     trial_samples += [eeg]
+
+                        subject_record_list.append(eeg_record)
+                        subject_index_list.append(eeg_index)
+
+                        eeg = self.read_eeg(eeg_record, eeg_index)
+                        trial_samples += [eeg]
+
+                        pbar.update(1)
+
+                    trial_samples = self.hook_data_interface(
+                        after_trial, trial_samples)
+                    session_samples += trial_samples
+
+                session_samples = self.hook_data_interface(
+                    after_session, session_samples)
+                subject_samples += session_samples
+
+            subject_samples = self.hook_data_interface(after_subject,
+                                                       subject_samples)
+
+            # save the data
+            for i, eeg in enumerate(subject_samples):
+                eeg_index = str(subject_index_list[i])
+                eeg_record = str(subject_record_list[i])
+
+                self.write_eeg(eeg_record, eeg_index, eeg)
+
+        pbar.close()
+
+    @staticmethod
+    def hook_data_interface(hook: Callable, data: Any):
+        # like [np.random.randn(32, 128), np.random.randn(32, 128)]
+        if isinstance(data[0], np.ndarray):
+            data = np.stack(data, axis=0)
+        elif isinstance(data[0], torch.Tensor):
+            data = torch.stack(data, axis=0)
+        # else list
+
+        # shape like (2, 32, 128)
+        data = hook(data)
+
+        # back to list like [np.random.randn(32, 128), np.random.randn(32, 128)]
+        if isinstance(data, np.ndarray):
+            data = np.split(data, data.shape[0], axis=0)
+            data = [np.squeeze(d, axis=0) for d in data]
+        elif isinstance(data, torch.Tensor):
+            data = torch.split(data, data.shape[0], dim=0)
+            data = [torch.squeeze(d, axis=0) for d in data]
+        # else list
+        return data
+
+    def read_eeg(self, record: str, key: str) -> Any:
         r'''
         Query the corresponding EEG signal in the EEGSignalIO according to the the given :obj:`key`. If :obj:`self.in_memory` is set to :obj:`True`, then EEGSignalIO will be read into memory and directly index the specified EEG signal in memory with the given :obj:`key` on subsequent reads.
 
         Args:
+            record (str): The record id of the EEG signal to be queried.
             key (str): The index of the EEG signal to be queried.
             
         Returns:
             any: The EEG signal sample.
         '''
+        eeg_io = self.eeg_io_router[record]
+        if self.in_memory:
+            return eeg_io.read_eeg_in_memory(key)
+        return eeg_io.read_eeg(key)
+
+    def write_eeg(self, record: str, key: str, eeg: Any):
+        r'''
+        Update the corresponding EEG signal in the EEGSignalIO according to the the given :obj:`key`.
+
+        Args:
+            record (str): The record id of the EEG signal to be queried.
+            key (str): The index of the EEG signal to be queried.
+            eeg (any): The EEG signal sample to be updated.
+        '''
+        eeg_io = self.eeg_io_router[record]
         if self.in_memory:
-            return self.eeg_io.read_eeg_in_memory(key)
-        return self.eeg_io.read_eeg(key)
+            eeg_io.write_eeg_in_memory(eeg=eeg, key=key)
+        eeg_io.write_eeg(eeg=eeg, key=key)
 
     def read_info(self, index: int) -> Dict:
         r'''
@@ -233,16 +411,30 @@ def exist(self, io_path: str) -> bool:
         Returns:
             bool: True if the database IO exists, otherwise False.
         '''
-        meta_info_io_path = os.path.join(io_path, 'info.csv')
-        eeg_signal_io_path = eeg_signal_io_path = os.path.join(io_path, 'eeg')
 
-        return os.path.exists(meta_info_io_path) and os.path.exists(
-            eeg_signal_io_path)
+        return os.path.exists(io_path)
 
     def __getitem__(self, index: int) -> any:
-        raise NotImplementedError(
-            "Method __getitem__ is not implemented in class " +
-            self.__class__.__name__)
+        info = self.read_info(index)
+
+        eeg_index = str(info['clip_id'])
+        eeg_record = str(info['_record_id'])
+        eeg = self.read_eeg(eeg_record, eeg_index)
+
+        return eeg, info
+
+    def get_labels(self) -> list:
+        '''
+        Get the labels of the dataset.
+
+        Returns:
+            list: The list of labels.
+        '''
+        labels = []
+        for i in range(len(self)):
+            _, label = self.__getitem__(i)
+            labels.append(label)
+        return labels
 
     def __len__(self):
         return len(self.info)
@@ -251,11 +443,11 @@ def __copy__(self) -> 'BaseDataset':
         cls = self.__class__
         result = cls.__new__(cls)
         result.__dict__.update(self.__dict__)
-        eeg_signal_io_path = os.path.join(self.io_path, 'eeg')
 
-        result.eeg_io = EEGSignalIO(eeg_signal_io_path,
-                                    io_size=self.io_size,
-                                    io_mode=self.io_mode)
+        result.eeg_io_router, result.info = self.get_pointer(
+            io_path=result.io_path,
+            io_size=result.io_size,
+            io_mode=result.io_mode)
         return result
 
     @property
diff --git a/torcheeg/datasets/module/csv_folder_dataset.py b/torcheeg/datasets/module/csv_folder_dataset.py
new file mode 100644
index 0000000..18af0fc
--- /dev/null
+++ b/torcheeg/datasets/module/csv_folder_dataset.py
@@ -0,0 +1,177 @@
+from typing import Any, Callable, Dict, Tuple, Union
+import pandas as pd
+import mne
+import numpy as np
+
+from .base_dataset import BaseDataset
+
+
+def default_read_fn(file_path, **kwargs):
+    # Load EEG file
+    raw = mne.io.read_raw(file_path)
+    # Convert raw to epochs
+    epochs = mne.make_fixed_length_epochs(raw, duration=1)
+    # Return EEG data
+    return epochs
+
+
+class CSVFolderDataset(BaseDataset):
+    '''
+    Read meta information from CSV file and read EEG data from folder according to the meta information. The CSV file should contain the following columns:
+
+    - ``subject_id`` (Optional): The subject id of the EEG data. Commonly used in training and testing dataset split.
+    - ``label`` (Optional): The label of the EEG data. Commonly used in training and testing dataset split.
+    - ``file_path`` (Required): The path to the EEG data file.
+
+    .. code-block:: python
+
+        # data.csv
+        # | subject_id | trial_id | label | file_path                 |
+        # | ---------- | -------  | ----- | ------------------------- |
+        # | sub1       | 0        | 0     | './data/label1/sub1.fif' |
+        # | sub1       | 1        | 1     | './data/label2/sub1.fif' |
+        # | sub1       | 2        | 2     | './data/label3/sub1.fif' |
+        # | sub2       | 0        | 0     | './data/label1/sub2.fif' |
+        # | sub2       | 1        | 1     | './data/label2/sub2.fif' |
+        # | sub2       | 2        | 2     | './data/label3/sub2.fif' |
+
+        def default_read_fn(file_path, **kwargs):
+            # Load EEG file
+            raw = mne.io.read_raw(file_path)
+            # Convert raw to epochs
+            epochs = mne.make_fixed_length_epochs(raw, duration=1)
+            # Return EEG data
+            return epochs
+
+        dataset = CSVFolderDataset(csv_path='./data.csv',
+                                   read_fn=default_read_fn,
+                                   online_transform=transforms.ToTensor(),
+                                   label_transform=transforms.Select('label'),
+                                   num_worker=4)
+
+    Args:
+        csv_path (str): The path to the CSV file.
+        num_channel (int): Number of channels in the EEG signal. (default: :obj:`14`)
+        chunk_size (int): Number of data points included in each EEG chunk as training or test samples. If set to -1, the EEG signal of a trial is used as a sample of a chunk. (default: :obj:`128`)
+        overlap (int): The number of overlapping data points between different chunks when dividing EEG chunks. (default: :obj:`0`)
+        read_fn (Callable): Method for reading files in a folder. By default, this class provides methods for reading files using :obj:`mne.io.read_raw`. At the same time, we allow users to pass in custom file reading methods. The first input parameter of whose is file_path, and other parameters are additional parameters passed in when the class is initialized. For example, you can pass :obj:`chunk_size=32` to :obj:`FolderDataset`, then :obj:`chunk_size` will be received here.
+        online_transform (Callable, optional): The transformation of the EEG signals and baseline EEG signals. The input is a :obj:`np.ndarray`, and the ouput is used as the first and second value of each element in the dataset. (default: :obj:`None`)
+        offline_transform (Callable, optional): The usage is the same as :obj:`online_transform`, but executed before generating IO intermediate results. (default: :obj:`None`)
+        label_transform (Callable, optional): The transformation of the label. The input is an information dictionary, and the ouput is used as the third value of each element in the dataset. (default: :obj:`None`)
+        io_path (str): The path to generated unified data IO, cached as an intermediate result. (default: :obj:`./io/deap`)
+        io_size (int): Maximum size database may grow to; used to size the memory mapping. If database grows larger than ``map_size``, an exception will be raised and the user must close and reopen. (default: :obj:`10485760`)
+        io_mode (str): Storage mode of EEG signal. When io_mode is set to :obj:`lmdb`, TorchEEG provides an efficient database (LMDB) for storing EEG signals. LMDB may not perform well on limited operating systems, where a file system based EEG signal storage is also provided. When io_mode is set to :obj:`pickle`, pickle-based persistence files are used. (default: :obj:`lmdb`)
+        num_worker (int): Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process. (default: :obj:`0`)
+        verbose (bool): Whether to display logs during processing, such as progress bars, etc. (default: :obj:`True`)
+        in_memory (bool): Whether to load the entire dataset into memory. If :obj:`in_memory` is set to True, then the first time an EEG sample is read, the entire dataset is loaded into memory for subsequent retrieval. Otherwise, the dataset is stored on disk to avoid the out-of-memory problem. (default: :obj:`False`)
+    '''
+    def __init__(self,
+                 csv_path: str = './data.csv',
+                 read_fn: Union[None, Callable] = default_read_fn,
+                 online_transform: Union[None, Callable] = None,
+                 offline_transform: Union[None, Callable] = None,
+                 label_transform: Union[None, Callable] = None,
+                 io_path: str = './io/folder',
+                 io_size: int = 10485760,
+                 io_mode: str = 'lmdb',
+                 num_worker: int = 0,
+                 verbose: bool = True,
+                 in_memory: bool = False,
+                 **kwargs):
+        # pass all arguments to super class
+        params = {
+            'csv_path': csv_path,
+            'read_fn': read_fn,
+            'online_transform': online_transform,
+            'offline_transform': offline_transform,
+            'label_transform': label_transform,
+            'io_path': io_path,
+            'io_size': io_size,
+            'io_mode': io_mode,
+            'num_worker': num_worker,
+            'verbose': verbose,
+            'in_memory': in_memory
+        }
+        params.update(kwargs)
+        super().__init__(**params)
+        # save all arguments to __dict__
+        self.__dict__.update(params)
+
+    @staticmethod
+    def process_record(file: Any = None,
+                       offline_transform: Union[None, Callable] = None,
+                       read_fn: Union[None, Callable] = None,
+                       **kwargs):
+
+        trial_info = file
+        file_path = trial_info['file_path']
+
+        trial_samples = read_fn(file_path, **kwargs)
+        events = [i[0] for i in trial_samples.events]
+        events.append(
+            events[-1] +
+            np.diff(events)[0])  # time interval between all events are same
+
+        write_pointer = 0
+        for i, trial_signal in enumerate(trial_samples.get_data()):
+            t_eeg = trial_signal
+            if not offline_transform is None:
+                t = offline_transform(eeg=trial_signal)
+                t_eeg = t['eeg']
+
+            clip_id = f'{file_path}_{write_pointer}'
+            write_pointer += 1
+
+            record_info = {
+                **trial_info,
+                'start_at': events[i],
+                'end_at': events[i + 1],
+                'clip_id': clip_id
+            }
+
+            yield {'eeg': t_eeg, 'key': clip_id, 'info': record_info}
+
+    def set_records(self, csv_path: str = './data.csv', **kwargs):
+        # read csv
+        df = pd.read_csv(csv_path)
+
+        assert 'file_path' in df.columns, 'file_path is required in csv file.'
+
+        # df to a list of dict, each dict is a row
+        df_list = df.to_dict('records')
+
+        return df_list
+
+    def __getitem__(self, index: int) -> Tuple[any, any, int, int, int]:
+        info = self.read_info(index)
+
+        eeg_index = str(info['clip_id'])
+        eeg_record = str(info['_record_id'])
+        eeg = self.read_eeg(eeg_record, eeg_index)
+
+        signal = eeg
+        label = info
+
+        if self.online_transform:
+            signal = self.online_transform(eeg=eeg)['eeg']
+
+        if self.label_transform:
+            label = self.label_transform(y=info)['y']
+
+        return signal, label
+
+    @property
+    def repr_body(self) -> Dict:
+        return dict(
+            super().repr_body, **{
+                'csv_path': self.csv_path,
+                'read_fn': self.read_fn,
+                'online_transform': self.online_transform,
+                'offline_transform': self.offline_transform,
+                'label_transform': self.label_transform,
+                'before_trial': self.before_trial,
+                'after_trial': self.after_trial,
+                'num_worker': self.num_worker,
+                'verbose': self.verbose,
+                'io_size': self.io_size
+            })
diff --git a/torcheeg/datasets/module/emotion_recognition/amigos.py b/torcheeg/datasets/module/emotion_recognition/amigos.py
index f8f1d62..76ff558 100644
--- a/torcheeg/datasets/module/emotion_recognition/amigos.py
+++ b/torcheeg/datasets/module/emotion_recognition/amigos.py
@@ -1,4 +1,5 @@
 import os
+import logging
 import re
 from typing import Any, Callable, Dict, List, Tuple, Union
 
@@ -8,6 +9,7 @@
     AMIGOS_ADJACENCY_MATRIX, AMIGOS_CHANNEL_LOCATION_DICT)
 from ..base_dataset import BaseDataset
 
+log = logging.getLogger(__name__)
 
 class AMIGOSDataset(BaseDataset):
     r'''
@@ -144,6 +146,8 @@ def __init__(self,
                  label_transform: Union[None, Callable] = None,
                  before_trial: Union[None, Callable] = None,
                  after_trial: Union[Callable, None] = None,
+                 after_session: Union[Callable, None] = None,
+                 after_subject: Union[Callable, None] = None,
                  io_path: str = './io/amigos',
                  io_size: int = 10485760,
                  io_mode: str = 'lmdb',
@@ -165,6 +169,8 @@ def __init__(self,
             'label_transform': label_transform,
             'before_trial': before_trial,
             'after_trial': after_trial,
+            'after_session': after_session,
+            'after_subject': after_subject,
             'io_path': io_path,
             'io_size': io_size,
             'io_mode': io_mode,
@@ -177,7 +183,7 @@ def __init__(self,
         self.__dict__.update(params)
 
     @staticmethod
-    def _load_data(file: Any = None,
+    def process_record(file: Any = None,
                    root_path: str = './data_preprocessed',
                    chunk_size: int = 128,
                    overlap: int = 0,
@@ -228,7 +234,7 @@ def _load_data(file: Any = None,
                 # 3 of the participants (08,24,28<->32) of the previous experiment did not watch a set of 4 long affective
                 if sum(trial_samples.shape) != sum(trial_rating.shape):
                     print(
-                        f'[WARNING] Find EEG signals without labels, or labels without EEG signals. Please check the {trial_id + 1}-th experiment of the {subject}-th subject in the file {file_name}. TorchEEG currently skipped the mismatched data.'
+                        f'Find EEG signals without labels, or labels without EEG signals. Please check the {trial_id + 1}-th experiment of the {subject}-th subject in the file {file_name}. TorchEEG currently skipped the mismatched data.'
                     )
                 continue
 
@@ -291,36 +297,23 @@ def _load_data(file: Any = None,
                     'clip_id': clip_id
                 }
                 record_info.update(trial_meta_info)
-                if after_trial:
-                    trial_queue.append({
-                        'eeg': t_eeg,
-                        'key': clip_id,
-                        'info': record_info
-                    })
-                else:
-                    yield {'eeg': t_eeg, 'key': clip_id, 'info': record_info}
+                yield {'eeg': t_eeg, 'key': clip_id, 'info': record_info}
 
                 start_at = start_at + step
                 end_at = start_at + chunk_size
 
-            if len(trial_queue) and after_trial:
-                trial_queue = after_trial(trial_queue)
-                for obj in trial_queue:
-                    assert 'eeg' in obj and 'key' in obj and 'info' in obj, 'after_trial must return a list of dictionaries, where each dictionary corresponds to an EEG sample, containing `eeg`, `key` and `info` as keys.'
-                    yield obj
-
-    @staticmethod
-    def _set_files(root_path: str = './data_preprocessed', **kwargs):
+    def set_records(self, root_path: str = './data_preprocessed', **kwargs):
         return os.listdir(root_path)
 
     def __getitem__(self, index: int) -> Tuple[any, any, int, int, int]:
         info = self.read_info(index)
 
         eeg_index = str(info['clip_id'])
-        eeg = self.read_eeg(eeg_index)
+        eeg_record = str(info['_record_id'])
+        eeg = self.read_eeg(eeg_record, eeg_index)
 
         baseline_index = str(info['baseline_id'])
-        baseline = self.read_eeg(baseline_index)
+        baseline = self.read_eeg(eeg_record, baseline_index)
 
         signal = eeg
         label = info
diff --git a/torcheeg/datasets/module/emotion_recognition/bci2022.py b/torcheeg/datasets/module/emotion_recognition/bci2022.py
index 65fce99..404e55e 100644
--- a/torcheeg/datasets/module/emotion_recognition/bci2022.py
+++ b/torcheeg/datasets/module/emotion_recognition/bci2022.py
@@ -240,7 +240,7 @@ def __init__(self,
         self.__dict__.update(params)
 
     @staticmethod
-    def _load_data(file: Any = None,
+    def process_record(file: Any = None,
                    chunk_size: int = 250,
                    overlap: int = 0,
                    channel_num: int = 30,
@@ -298,8 +298,6 @@ def _load_data(file: Any = None,
                 cur_end_at = cur_start_at + chunk_size
                 step = chunk_size - overlap
 
-                trial_queue = []
-
                 if before_trial:
                     samples[:channel_num, cur_start_at:end_at] = before_trial(
                         samples[:channel_num, cur_start_at:end_at])
@@ -318,14 +316,7 @@ def _load_data(file: Any = None,
                         'clip_id': clip_id
                     }
                     record_info.update(trial_meta_info)
-                    if after_trial:
-                        trial_queue.append({
-                            'eeg': t_eeg,
-                            'key': clip_id,
-                            'info': record_info
-                        })
-                    else:
-                        yield {
+                    yield {
                             'eeg': t_eeg,
                             'key': clip_id,
                             'info': record_info
@@ -334,20 +325,13 @@ def _load_data(file: Any = None,
                     cur_start_at = cur_start_at + step
                     cur_end_at = cur_start_at + chunk_size
 
-                if len(trial_queue) and after_trial:
-                    trial_queue = after_trial(trial_queue)
-                    for obj in trial_queue:
-                        assert 'eeg' in obj and 'key' in obj and 'info' in obj, 'after_trial must return a list of dictionaries, where each dictionary corresponds to an EEG sample, containing `eeg`, `key` and `info` as keys.'
-                        yield obj
-
                 # prepare for the next trial
                 trial_id += 1
                 video_id = None
                 start_at = None
                 end_at = None
 
-    @staticmethod
-    def _set_files(root_path: str = './data_preprocessed_python', **kwargs):
+    def set_records(self, root_path: str = './data_preprocessed_python', **kwargs):
         outputs = []
         for train_set_batch in [
                 'TrainSet_first_batch', 'TrainSet_second_batch'
@@ -364,7 +348,8 @@ def __getitem__(self, index: int) -> Tuple[any, any, int, int, int]:
         info = self.read_info(index)
 
         eeg_index = str(info['clip_id'])
-        eeg = self.read_eeg(eeg_index)
+        eeg_record = str(info['_record_id'])
+        eeg = self.read_eeg(eeg_record, eeg_index)
 
         signal = eeg
         label = info
diff --git a/torcheeg/datasets/module/emotion_recognition/deap.py b/torcheeg/datasets/module/emotion_recognition/deap.py
index 19e6007..3dcd2f6 100644
--- a/torcheeg/datasets/module/emotion_recognition/deap.py
+++ b/torcheeg/datasets/module/emotion_recognition/deap.py
@@ -134,6 +134,8 @@ def __init__(self,
                  label_transform: Union[None, Callable] = None,
                  before_trial: Union[None, Callable] = None,
                  after_trial: Union[Callable, None] = None,
+                 after_session: Union[Callable, None] = None,
+                 after_subject: Union[Callable, None] = None,
                  io_path: str = './io/deap',
                  io_size: int = 10485760,
                  io_mode: str = 'lmdb',
@@ -153,6 +155,8 @@ def __init__(self,
             'label_transform': label_transform,
             'before_trial': before_trial,
             'after_trial': after_trial,
+            'after_session': after_session,
+            'after_subject': after_subject,
             'io_path': io_path,
             'io_size': io_size,
             'io_mode': io_mode,
@@ -165,7 +169,7 @@ def __init__(self,
         self.__dict__.update(params)
 
     @staticmethod
-    def _load_data(file: Any = None,
+    def process_record(file: Any = None,
                    root_path: str = './data_preprocessed_python',
                    chunk_size: int = 128,
                    overlap: int = 0,
@@ -219,7 +223,6 @@ def _load_data(file: Any = None,
             # calculate moving step
             step = chunk_size - overlap
 
-            trial_queue = []
             while end_at <= trial_samples.shape[1]:
                 clip_sample = trial_samples[:, start_at:end_at]
 
@@ -249,37 +252,23 @@ def _load_data(file: Any = None,
                     'clip_id': clip_id
                 }
                 record_info.update(trial_meta_info)
-                if after_trial:
-                    trial_queue.append({
-                        'eeg': t_eeg,
-                        'key': clip_id,
-                        'info': record_info
-                    })
-                else:
-                    yield {'eeg': t_eeg, 'key': clip_id, 'info': record_info}
+                yield {'eeg': t_eeg, 'key': clip_id, 'info': record_info}
 
                 start_at = start_at + step
                 end_at = start_at + chunk_size
 
-            if len(trial_queue) and after_trial:
-                trial_queue = after_trial(trial_queue)
-                for obj in trial_queue:
-                    assert 'eeg' in obj and 'key' in obj and 'info' in obj, 'after_trial must return a list of dictionaries, where each dictionary corresponds to an EEG sample, containing `eeg`, `key` and `info` as keys.'
-                    yield obj
-
-    @staticmethod
-    def _set_files(root_path: str = './data_preprocessed_python',
+    def set_records(self, root_path: str = './data_preprocessed_python',
         **kwargs):
         return os.listdir(root_path)
 
     def __getitem__(self, index: int) -> Tuple:
         info = self.read_info(index)
-
         eeg_index = str(info['clip_id'])
-        eeg = self.read_eeg(eeg_index)
+        eeg_record = str(info['_record_id'])
+        eeg = self.read_eeg(eeg_record, eeg_index)
 
         baseline_index = str(info['baseline_id'])
-        baseline = self.read_eeg(baseline_index)
+        baseline = self.read_eeg(eeg_record, baseline_index)
 
         signal = eeg
         label = info
diff --git a/torcheeg/datasets/module/emotion_recognition/dreamer.py b/torcheeg/datasets/module/emotion_recognition/dreamer.py
index 74ff3fd..e3bc7bb 100644
--- a/torcheeg/datasets/module/emotion_recognition/dreamer.py
+++ b/torcheeg/datasets/module/emotion_recognition/dreamer.py
@@ -133,6 +133,8 @@ def __init__(self,
                  label_transform: Union[None, Callable] = None,
                  before_trial: Union[None, Callable] = None,
                  after_trial: Union[Callable, None] = None,
+                 after_session: Union[Callable, None] = None,
+                 after_subject: Union[Callable, None] = None,
                  io_path: str = './io/dreamer',
                  io_size: int = 10485760,
                  io_mode: str = 'lmdb',
@@ -152,6 +154,8 @@ def __init__(self,
             'label_transform': label_transform,
             'before_trial': before_trial,
             'after_trial': after_trial,
+            'after_session': after_session,
+            'after_subject': after_subject,
             'io_path': io_path,
             'io_size': io_size,
             'io_mode': io_mode,
@@ -164,7 +168,7 @@ def __init__(self,
         self.__dict__.update(params)
 
     @staticmethod
-    def _load_data(file: Any = None,
+    def process_record(file: Any = None,
                    mat_path: str = './DREAMER.mat',
                    chunk_size: int = 128,
                    overlap: int = 0,
@@ -259,26 +263,12 @@ def _load_data(file: Any = None,
                     'clip_id': clip_id
                 }
                 record_info.update(trial_meta_info)
-                if after_trial:
-                    trial_queue.append({
-                        'eeg': t_eeg,
-                        'key': clip_id,
-                        'info': record_info
-                    })
-                else:
-                    yield {'eeg': t_eeg, 'key': clip_id, 'info': record_info}
+                yield {'eeg': t_eeg, 'key': clip_id, 'info': record_info}
 
                 start_at = start_at + step
                 end_at = start_at + chunk_size
 
-            if len(trial_queue) and after_trial:
-                trial_queue = after_trial(trial_queue)
-                for obj in trial_queue:
-                    assert 'eeg' in obj and 'key' in obj and 'info' in obj, 'after_trial must return a list of dictionaries, where each dictionary corresponds to an EEG sample, containing `eeg`, `key` and `info` as keys.'
-                    yield obj
-
-    @staticmethod
-    def _set_files(mat_path: str = './DREAMER.mat', **kwargs):
+    def set_records(self, mat_path: str = './DREAMER.mat', **kwargs):
 
         mat_data = scio.loadmat(mat_path,
                                 verify_compressed_data_integrity=False)
@@ -290,10 +280,11 @@ def __getitem__(self, index: int) -> Tuple:
         info = self.read_info(index)
 
         eeg_index = str(info['clip_id'])
-        eeg = self.read_eeg(eeg_index)
+        eeg_record = str(info['_record_id'])
+        eeg = self.read_eeg(eeg_record, eeg_index)
 
         baseline_index = str(info['baseline_id'])
-        baseline = self.read_eeg(baseline_index)
+        baseline = self.read_eeg(eeg_record, baseline_index)
 
         signal = eeg
         label = info
diff --git a/torcheeg/datasets/module/emotion_recognition/mahnob.py b/torcheeg/datasets/module/emotion_recognition/mahnob.py
index aa4ef6d..c923fd3 100644
--- a/torcheeg/datasets/module/emotion_recognition/mahnob.py
+++ b/torcheeg/datasets/module/emotion_recognition/mahnob.py
@@ -153,6 +153,8 @@ def __init__(self,
                  label_transform: Union[None, Callable] = None,
                  before_trial: Union[None, Callable] = None,
                  after_trial: Union[Callable, None] = None,
+                 after_session: Union[Callable, None] = None,
+                 after_subject: Union[Callable, None] = None,
                  io_path: str = './io/mahnob',
                  io_size: int = 10485760,
                  io_mode: str = 'lmdb',
@@ -174,6 +176,8 @@ def __init__(self,
             'label_transform': label_transform,
             'before_trial': before_trial,
             'after_trial': after_trial,
+            'after_session': after_session,
+            'after_subject': after_subject,
             'io_path': io_path,
             'io_size': io_size,
             'io_mode': io_mode,
@@ -186,7 +190,7 @@ def __init__(self,
         self.__dict__.update(params)
 
     @staticmethod
-    def _load_data(
+    def process_record(
                file: Any = None,
                 root_path: str = './Sessions',
                 chunk_size: int = 128,
@@ -198,7 +202,6 @@ def _load_data(
                 num_trial_sample: int = 30,
                 offline_transform: Union[None, Callable] = None,
                 before_trial: Union[None, Callable] = None,
-                after_trial: Union[None, Callable] = None,
                **kwargs):
         file_name = file
 
@@ -278,7 +281,6 @@ def _load_data(
         if not (num_trial_sample <= 0):
             max_len = min(num_trial_sample * chunk_size, trial_samples.shape[1])
 
-        trial_queue = []
         while end_at <= max_len:
             clip_sample = trial_samples[:, start_at:end_at]
 
@@ -309,14 +311,7 @@ def _load_data(
                 'clip_id': clip_id
             }
             record_info.update(trial_meta_info)
-            if after_trial:
-                trial_queue.append({
-                    'eeg': t_eeg,
-                    'key': clip_id,
-                    'info': record_info
-                })
-            else:
-                yield {
+            yield {
                     'eeg': t_eeg,
                     'key': clip_id,
                     'info': record_info
@@ -325,24 +320,18 @@ def _load_data(
             start_at = start_at + step
             end_at = start_at + chunk_size
 
-        if len(trial_queue) and after_trial:
-            trial_queue = after_trial(trial_queue)
-            for obj in trial_queue:
-                assert 'eeg' in obj and 'key' in obj and 'info' in obj, 'after_trial must return a list of dictionaries, where each dictionary corresponds to an EEG sample, containing `eeg`, `key` and `info` as keys.'
-                yield obj
-
-    @staticmethod
-    def _set_files(root_path: str = './Sessions', **kwargs):
+    def set_records(self, root_path: str = './Sessions', **kwargs):
         return os.listdir(root_path)
 
     def __getitem__(self, index: int) -> Tuple:
         info = self.read_info(index)
 
         eeg_index = str(info['clip_id'])
-        eeg = self.read_eeg(eeg_index)
+        eeg_record = str(info['_record_id'])
+        eeg = self.read_eeg(eeg_record, eeg_index)
 
         baseline_index = str(info['baseline_id'])
-        baseline = self.read_eeg(baseline_index)
+        baseline = self.read_eeg(eeg_record, baseline_index)
 
         signal = eeg
         label = info
diff --git a/torcheeg/datasets/module/emotion_recognition/mped_feature.py b/torcheeg/datasets/module/emotion_recognition/mped_feature.py
index 21caa22..424a2e1 100644
--- a/torcheeg/datasets/module/emotion_recognition/mped_feature.py
+++ b/torcheeg/datasets/module/emotion_recognition/mped_feature.py
@@ -140,7 +140,7 @@ def __init__(self,
         self.__dict__.update(params)
 
     @staticmethod
-    def _load_data(file: Any = None,
+    def process_record(file: Any = None,
                    root_path: str = './EEG_feature',
                    feature: list = ['PSD'],
                    num_channel: int = 62,
@@ -195,7 +195,6 @@ def _load_data(file: Any = None,
             }
             num_clips = trial_samples.shape[1]
 
-            trial_queue = []
             for clip_id in range(num_clips):
                 # PSD: (62, 5)
                 # Hjorth: (62, 3 * 6)
@@ -214,23 +213,9 @@ def _load_data(file: Any = None,
                 # record meta info for each signal
                 record_info = {'clip_id': clip_id}
                 record_info.update(trial_meta_info)
-                if after_trial:
-                    trial_queue.append({
-                        'eeg': t_eeg,
-                        'key': clip_id,
-                        'info': record_info
-                    })
-                else:
-                    yield {'eeg': t_eeg, 'key': clip_id, 'info': record_info}
-
-            if len(trial_queue) and after_trial:
-                trial_queue = after_trial(trial_queue)
-                for obj in trial_queue:
-                    assert 'eeg' in obj and 'key' in obj and 'info' in obj, 'after_trial must return a list of dictionaries, where each dictionary corresponds to an EEG sample, containing `eeg`, `key` and `info` as keys.'
-                    yield obj
+                yield {'eeg': t_eeg, 'key': clip_id, 'info': record_info}
 
-    @staticmethod
-    def _set_files(root_path: str = './EEG_feature',
+    def set_records(self, root_path: str = './EEG_feature',
                    feature: list = ['PSD'],
                    **kwargs):
         avaliable_features = os.listdir(root_path)
@@ -246,7 +231,8 @@ def __getitem__(self, index: int) -> Tuple[any, any, int, int, int]:
         info = self.read_info(index)
 
         eeg_index = str(info['clip_id'])
-        eeg = self.read_eeg(eeg_index)
+        eeg_record = str(info['_record_id'])
+        eeg = self.read_eeg(eeg_record, eeg_index)
 
         signal = eeg
         label = info
diff --git a/torcheeg/datasets/module/emotion_recognition/seed.py b/torcheeg/datasets/module/emotion_recognition/seed.py
index a1798cb..e92f837 100644
--- a/torcheeg/datasets/module/emotion_recognition/seed.py
+++ b/torcheeg/datasets/module/emotion_recognition/seed.py
@@ -135,6 +135,8 @@ def __init__(self,
                  label_transform: Union[None, Callable] = None,
                  before_trial: Union[None, Callable] = None,
                  after_trial: Union[Callable, None] = None,
+                 after_session: Union[Callable, None] = None,
+                 after_subject: Union[Callable, None] = None,
                  io_path: str = './io/seed',
                  io_size: int = 10485760,
                  io_mode: str = 'lmdb',
@@ -152,6 +154,8 @@ def __init__(self,
             'label_transform': label_transform,
             'before_trial': before_trial,
             'after_trial': after_trial,
+            'after_session': after_session,
+            'after_subject': after_subject,
             'io_path': io_path,
             'io_size': io_size,
             'io_mode': io_mode,
@@ -164,7 +168,7 @@ def __init__(self,
         self.__dict__.update(params)
 
     @staticmethod
-    def _load_data(file: Any = None,
+    def process_record(file: Any = None,
                    root_path: str = './Preprocessed_EEG',
                    chunk_size: int = 200,
                    overlap: int = 0,
@@ -216,7 +220,6 @@ def _load_data(file: Any = None,
             # calculate moving step
             step = chunk_size - overlap
 
-            trial_queue = []
             while end_at <= trial_samples.shape[1]:
                 clip_sample = trial_samples[:num_channel, start_at:end_at]
 
@@ -234,26 +237,12 @@ def _load_data(file: Any = None,
                     'clip_id': clip_id
                 }
                 record_info.update(trial_meta_info)
-                if after_trial:
-                    trial_queue.append({
-                        'eeg': t_eeg,
-                        'key': clip_id,
-                        'info': record_info
-                    })
-                else:
-                    yield {'eeg': t_eeg, 'key': clip_id, 'info': record_info}
+                yield {'eeg': t_eeg, 'key': clip_id, 'info': record_info}
 
                 start_at = start_at + step
                 end_at = start_at + chunk_size
 
-            if len(trial_queue) and after_trial:
-                trial_queue = after_trial(trial_queue)
-                for obj in trial_queue:
-                    assert 'eeg' in obj and 'key' in obj and 'info' in obj, 'after_trial must return a list of dictionaries, where each dictionary corresponds to an EEG sample, containing `eeg`, `key` and `info` as keys.'
-                    yield obj
-
-    @staticmethod
-    def _set_files(root_path: str = './Preprocessed_EEG', **kwargs):
+    def set_records(self, root_path: str = './Preprocessed_EEG', **kwargs):
         file_list = os.listdir(root_path)
         skip_set = ['label.mat', 'readme.txt']
         file_list = [f for f in file_list if f not in skip_set]
@@ -263,7 +252,8 @@ def __getitem__(self, index: int) -> Tuple[any, any, int, int, int]:
         info = self.read_info(index)
 
         eeg_index = str(info['clip_id'])
-        eeg = self.read_eeg(eeg_index)
+        eeg_record = str(info['_record_id'])
+        eeg = self.read_eeg(eeg_record, eeg_index)
 
         signal = eeg
         label = info
diff --git a/torcheeg/datasets/module/emotion_recognition/seed_feature.py b/torcheeg/datasets/module/emotion_recognition/seed_feature.py
index b40f9f6..78fd5ac 100644
--- a/torcheeg/datasets/module/emotion_recognition/seed_feature.py
+++ b/torcheeg/datasets/module/emotion_recognition/seed_feature.py
@@ -116,6 +116,8 @@ def __init__(self,
                  label_transform: Union[None, Callable] = None,
                  before_trial: Union[None, Callable] = None,
                  after_trial: Union[Callable, None] = None,
+                 after_session: Union[Callable, None] = None,
+                 after_subject: Union[Callable, None] = None,
                  io_path: str = './io/seed_feature',
                  io_size: int = 10485760,
                  io_mode: str = 'lmdb',
@@ -132,6 +134,8 @@ def __init__(self,
             'label_transform': label_transform,
             'before_trial': before_trial,
             'after_trial': after_trial,
+            'after_session': after_session,
+            'after_subject': after_subject,
             'io_path': io_path,
             'io_size': io_size,
             'io_mode': io_mode,
@@ -144,15 +148,15 @@ def __init__(self,
         self.__dict__.update(params)
 
     @staticmethod
-    def _load_data(file: Any = None,
-                   root_path: str = './ExtractedFeatures',
-                   feature: list = ['de_movingAve'],
-                   num_channel: int = 62,
-                   before_trial: Union[None, Callable] = None,
-                   offline_transform: Union[None, Callable] = None,
-                   after_trial: Union[None, Callable] = None,
-                   **kwargs):
-        file_name = file
+    def process_record(file: Any = None,
+                       root_path: str = './ExtractedFeatures',
+                       feature: list = ['de_movingAve'],
+                       num_channel: int = 62,
+                       before_trial: Union[None, Callable] = None,
+                       offline_transform: Union[None, Callable] = None,
+                       after_trial: Union[None, Callable] = None,
+                       **kwargs):
+        file_name, session_id = file
 
         subject = int(os.path.basename(file_name).split('.')[0].split('_')
                       [0])  # subject (15)
@@ -171,21 +175,19 @@ def _load_data(file: Any = None,
             int(re.findall(r"de_movingAve(\d+)", key)[0])
             for key in samples.keys() if 'de_movingAve' in key
         ]
-
         write_pointer = 0
         # loop for each trial
         for trial_id in trial_ids:
             # extract baseline signals
             trial_samples = []
             for cur_feature in feature:
-                trial_samples.append(samples[
-                    cur_feature +
-                    str(trial_id)])  # channel(62), timestep(n), bands(5)
+                trial_samples.append(
+                    samples[cur_feature + str(trial_id)].transpose(
+                        (1, 0, 2)))  # timestep(n), channel(62), bands(5)
             trial_samples = np.concatenate(
                 trial_samples,
-                axis=-1)[:
-                         num_channel]  # channel(62), timestep(n), features(5*k)
-            trial_samples = trial_samples.transpose((1, 0, 2))
+                axis=-1)[:, :
+                         num_channel]  # timestep(n*k), channel(62), features(5)
 
             if before_trial:
                 trial_samples = before_trial(trial_samples)
@@ -195,10 +197,10 @@ def _load_data(file: Any = None,
                 'subject_id': subject,
                 'trial_id': trial_id,
                 'emotion': int(labels[trial_id - 1]),
-                'date': date
+                'date': date,
+                'session_id': session_id
             }
 
-            trial_queue = []
             for i, clip_sample in enumerate(trial_samples):
                 t_eeg = clip_sample
                 if not offline_transform is None:
@@ -214,34 +216,41 @@ def _load_data(file: Any = None,
                     'clip_id': clip_id
                 }
                 record_info.update(trial_meta_info)
-                if after_trial:
-                    trial_queue.append({
-                        'eeg': t_eeg,
-                        'key': clip_id,
-                        'info': record_info
-                    })
-                else:
-                    yield {'eeg': t_eeg, 'key': clip_id, 'info': record_info}
-
-            if len(trial_queue) and after_trial:
-                trial_queue = after_trial(trial_queue)
-                for obj in trial_queue:
-                    assert 'eeg' in obj and 'key' in obj and 'info' in obj, 'after_trial must return a list of dictionaries, where each dictionary corresponds to an EEG sample, containing `eeg`, `key` and `info` as keys.'
-                    yield obj
+                yield {'eeg': t_eeg, 'key': clip_id, 'info': record_info}
 
-    @staticmethod
-    def _set_files(**kwargs):
-        root_path = kwargs.pop('root_path', './ExtractedFeatures')  # str
+    def set_records(self, root_path: str = './ExtractedFeatures', **kwargs):
         file_list = os.listdir(root_path)
         skip_set = ['label.mat', 'readme.txt']
         file_list = [f for f in file_list if f not in skip_set]
-        return file_list
+
+        # get subject_id
+        subject_id_list = set()
+        for file_name in file_list:
+            subject_id = int(file_name.split('_')[0])
+            subject_id_list.add(subject_id)
+        subject_id_list = list(subject_id_list)
+
+        # get session_id
+        file_name_session_id_list = []
+        for subject_id in subject_id_list:
+            subject_file_list = [
+                file_name for file_name in file_list
+                if int(file_name.split('_')[0]) == subject_id
+            ]
+            # sort the subject_file_list, based on the date
+            subject_file_list.sort(key=lambda x: int(x.split('_')[1][:-4]))
+            # get file and corresponding session_id
+            for i, file_name in enumerate(subject_file_list):
+                file_name_session_id_list.append((file_name, i))
+
+        return file_name_session_id_list
 
     def __getitem__(self, index: int) -> Tuple[any, any, int, int, int]:
         info = self.read_info(index)
 
         eeg_index = str(info['clip_id'])
-        eeg = self.read_eeg(eeg_index)
+        eeg_record = str(info['_record_id'])
+        eeg = self.read_eeg(eeg_record, eeg_index)
 
         signal = eeg
         label = info
diff --git a/torcheeg/datasets/module/emotion_recognition/seed_iv.py b/torcheeg/datasets/module/emotion_recognition/seed_iv.py
index 94deab2..fb388c7 100644
--- a/torcheeg/datasets/module/emotion_recognition/seed_iv.py
+++ b/torcheeg/datasets/module/emotion_recognition/seed_iv.py
@@ -129,6 +129,8 @@ def __init__(self,
                  label_transform: Union[None, Callable] = None,
                  before_trial: Union[None, Callable] = None,
                  after_trial: Union[Callable, None] = None,
+                 after_session: Union[Callable, None] = None,
+                 after_subject: Union[Callable, None] = None,
                  io_path: str = './io/seed',
                  io_size: int = 10485760,
                  io_mode: str = 'lmdb',
@@ -146,6 +148,8 @@ def __init__(self,
             'label_transform': label_transform,
             'before_trial': before_trial,
             'after_trial': after_trial,
+            'after_session': after_session,
+            'after_subject': after_subject,
             'io_path': io_path,
             'io_size': io_size,
             'io_mode': io_mode,
@@ -158,7 +162,7 @@ def __init__(self,
         self.__dict__.update(params)
 
     @staticmethod
-    def _load_data(file: Any = None,
+    def process_record(file: Any = None,
                    chunk_size: int = 800,
                    overlap: int = 0,
                    num_channel: int = 62,
@@ -237,7 +241,6 @@ def _load_data(file: Any = None,
             # calculate moving step
             step = chunk_size - overlap
 
-            trial_queue = []
             while end_at <= trial_samples.shape[1]:
                 clip_sample = trial_samples[:num_channel, start_at:end_at]
 
@@ -255,26 +258,12 @@ def _load_data(file: Any = None,
                     'clip_id': clip_id
                 }
                 record_info.update(trial_meta_info)
-                if after_trial:
-                    trial_queue.append({
-                        'eeg': t_eeg,
-                        'key': clip_id,
-                        'info': record_info
-                    })
-                else:
-                    yield {'eeg': t_eeg, 'key': clip_id, 'info': record_info}
+                yield {'eeg': t_eeg, 'key': clip_id, 'info': record_info}
 
                 start_at = start_at + step
                 end_at = start_at + chunk_size
 
-            if len(trial_queue) and after_trial:
-                trial_queue = after_trial(trial_queue)
-                for obj in trial_queue:
-                    assert 'eeg' in obj and 'key' in obj and 'info' in obj, 'after_trial must return a list of dictionaries, where each dictionary corresponds to an EEG sample, containing `eeg`, `key` and `info` as keys.'
-                    yield obj
-
-    @staticmethod
-    def _set_files(root_path: str = './eeg_raw_data', **kwargs):
+    def set_records(self, root_path: str = './eeg_raw_data', **kwargs):
         session_list = ['1', '2', '3']
         file_path_list = []
         for session in session_list:
@@ -288,7 +277,8 @@ def __getitem__(self, index: int) -> Tuple[any, any, int, int, int]:
         info = self.read_info(index)
 
         eeg_index = str(info['clip_id'])
-        eeg = self.read_eeg(eeg_index)
+        eeg_record = str(info['_record_id'])
+        eeg = self.read_eeg(eeg_record, eeg_index)
 
         signal = eeg
         label = info
diff --git a/torcheeg/datasets/module/emotion_recognition/seed_iv_feature.py b/torcheeg/datasets/module/emotion_recognition/seed_iv_feature.py
index 07d8a16..e1d6a2e 100644
--- a/torcheeg/datasets/module/emotion_recognition/seed_iv_feature.py
+++ b/torcheeg/datasets/module/emotion_recognition/seed_iv_feature.py
@@ -105,6 +105,8 @@ def __init__(self,
                  label_transform: Union[None, Callable] = None,
                  before_trial: Union[None, Callable] = None,
                  after_trial: Union[Callable, None] = None,
+                 after_session: Union[Callable, None] = None,
+                 after_subject: Union[Callable, None] = None,
                  io_path: str = './io/seed_iv_feature',
                  io_size: int = 10485760,
                  io_mode: str = 'lmdb',
@@ -121,6 +123,8 @@ def __init__(self,
             'label_transform': label_transform,
             'before_trial': before_trial,
             'after_trial': after_trial,
+            'after_session': after_session,
+            'after_subject': after_subject,
             'io_path': io_path,
             'io_size': io_size,
             'io_mode': io_mode,
@@ -133,13 +137,13 @@ def __init__(self,
         self.__dict__.update(params)
 
     @staticmethod
-    def _load_data(feature: list = ['de_movingAve'],
-                   num_channel: int = 62,
-                   offline_transform: Union[None, Callable] = None,
-                   before_trial: Union[None, Callable] = None,
-                   after_trial: Union[Callable, None] = None,
-                   file: Any = None,
-                   **kwargs):
+    def process_record(feature: list = ['de_movingAve'],
+                       num_channel: int = 62,
+                       offline_transform: Union[None, Callable] = None,
+                       before_trial: Union[None, Callable] = None,
+                       after_trial: Union[Callable, None] = None,
+                       file: Any = None,
+                       **kwargs):
         file_path = file  # an element from file name list
 
         session_id = os.path.basename(os.path.dirname(file_path))
@@ -179,7 +183,6 @@ def _load_data(feature: list = ['de_movingAve'],
             ]
         ]  # The labels with 0, 1, 2, and 3 denote the ground truth, neutral, sad, fear, and happy emotions, respectively.
         session_labels = labels[int(session_id) - 1]
-
         trial_ids = [
             int(re.findall(r"de_movingAve(\d+)", key)[0])
             for key in samples.keys() if 'de_movingAve' in key
@@ -230,23 +233,9 @@ def _load_data(feature: list = ['de_movingAve'],
                     'clip_id': clip_id
                 }
                 record_info.update(trial_meta_info)
-                if after_trial:
-                    trial_queue.append({
-                        'eeg': t_eeg,
-                        'key': clip_id,
-                        'info': record_info
-                    })
-                else:
-                    yield {'eeg': t_eeg, 'key': clip_id, 'info': record_info}
-
-            if len(trial_queue) and after_trial:
-                trial_queue = after_trial(trial_queue)
-                for obj in trial_queue:
-                    assert 'eeg' in obj and 'key' in obj and 'info' in obj, 'after_trial must return a list of dictionaries, where each dictionary corresponds to an EEG sample, containing `eeg`, `key` and `info` as keys.'
-                    yield obj
-
-    @staticmethod
-    def _set_files(root_path: str = './eeg_feature_smooth', **kwargs):
+                yield {'eeg': t_eeg, 'key': clip_id, 'info': record_info}
+                
+    def set_records(self, root_path: str = './eeg_feature_smooth', **kwargs):
 
         session_list = ['1', '2', '3']
         file_path_list = []
@@ -262,7 +251,8 @@ def __getitem__(self, index: int) -> Tuple[any, any, int, int, int]:
         info = self.read_info(index)
 
         eeg_index = str(info['clip_id'])
-        eeg = self.read_eeg(eeg_index)
+        eeg_record = str(info['_record_id'])
+        eeg = self.read_eeg(eeg_record, eeg_index)
 
         signal = eeg
         label = info
diff --git a/torcheeg/datasets/module/folder_dataset.py b/torcheeg/datasets/module/folder_dataset.py
index 8595e42..d5698ad 100644
--- a/torcheeg/datasets/module/folder_dataset.py
+++ b/torcheeg/datasets/module/folder_dataset.py
@@ -1,12 +1,9 @@
-import os
 from pathlib import Path
 from typing import Any, Callable, Dict, Tuple, Union
 
 import mne
 import numpy as np
 
-from torcheeg.io import EEGSignalIO, MetaInfoIO
-
 from .base_dataset import BaseDataset
 
 
@@ -47,6 +44,38 @@ class FolderDataset(BaseDataset):
         |- label01.edf
         |- label02.edf
 
+    An example dataset for GNN-based methods:
+
+    .. code-block:: python
+    
+        sfreq = 128  # Sampling rate
+        n_channels = 14  # Number of channels
+        duration = 5  # Data collected for 5 seconds
+        for i in range(num_files):
+            n_samples = sfreq * duration
+            data = np.random.randn(n_channels, n_samples)
+
+            ch_names = [f'ch_{i+1:03}' for i in range(n_channels)]
+            ch_types = ['eeg'] * n_channels
+            info = mne.create_info(ch_names, sfreq, ch_types)
+            raw = mne.io.RawArray(data, info)
+
+            file_name = f'sub{i+1}.fif'
+            file_path = os.path.join('./root_folder/', file_name)
+            raw.save(file_path)
+
+        label_map = {'folder1': 0, 'folder2': 1}
+        dataset = FolderDataset(io_path='./folder',
+                                root_path='./root_folder',
+                                structure='subject_in_label',
+                                num_channel=14,
+                                online_transform=transforms.ToTensor(),
+                                label_transform=transforms.Compose([
+                                    transforms.Select('label'),
+                                    transforms.Lambda(lambda x: label_map[x])
+                                ]),
+                                num_worker=4)
+
     Args:
         root_path (str): The path to the root folder. (default: :obj:`'./folder'`)
         structure (str): Folder structure, which affects how labels and subjects are mapped to EEG signal samples. Please refer to the above description of the structure of the two folders to select the correct parameters. (default: :obj:`'subject_in_label'`)
@@ -96,7 +125,7 @@ def __init__(self,
         self.__dict__.update(params)
 
     @staticmethod
-    def _load_data(file: Any = None,
+    def process_record(file: Any = None,
                    offline_transform: Union[None, Callable] = None,
                    read_fn: Union[None, Callable] = None,
                    **kwargs):
@@ -121,7 +150,6 @@ def _load_data(file: Any = None,
 
             record_info = {
                 'subject_id': subject_id,
-                'trial_id': i,
                 'start_at': events[i],
                 'end_at': events[i + 1],
                 'clip_id': clip_id,
@@ -130,8 +158,7 @@ def _load_data(file: Any = None,
 
             yield {'eeg': t_eeg, 'key': clip_id, 'info': record_info}
 
-    @staticmethod
-    def _set_files(root_path: str = './folder',
+    def set_records(self, root_path: str = './folder',
                    structure: str = 'subject_in_label',
                    **kwargs):
         # get all the subfolders
@@ -160,7 +187,8 @@ def __getitem__(self, index: int) -> Tuple[any, any, int, int, int]:
         info = self.read_info(index)
 
         eeg_index = str(info['clip_id'])
-        eeg = self.read_eeg(eeg_index)
+        eeg_record = str(info['_record_id'])
+        eeg = self.read_eeg(eeg_record, eeg_index)
 
         signal = eeg
         label = info
diff --git a/torcheeg/datasets/module/mne_dataset.py b/torcheeg/datasets/module/mne_dataset.py
deleted file mode 100644
index a34ad92..0000000
--- a/torcheeg/datasets/module/mne_dataset.py
+++ /dev/null
@@ -1,400 +0,0 @@
-import os
-from multiprocessing import Manager
-from typing import Any, Callable, Dict, List, Tuple, Union
-
-import mne
-import numpy as np
-from joblib import Parallel, delayed
-from tqdm import tqdm
-
-from torcheeg.io import EEGSignalIO, MetaInfoIO
-
-from .base_dataset import BaseDataset
-
-MAX_QUEUE_SIZE = 1024
-
-
-class MockLock():
-    def __enter__(self):
-        pass
-
-    def __exit__(self, exc_type, exc_value, traceback):
-        pass
-
-
-class MNEDataset(BaseDataset):
-    r'''
-    A generic EEG analysis dataset that allows creating datasets from :obj:`mne.Epochs`, and caches the generated results in a unified input and output format (IO). It is generally used to support custom datasets or datasets not yet provided by TorchEEG.
-
-    :obj:`MNEDataset` allows a list of :obj:`mne.Epochs` and the corresponding description dictionary as input, and divides the signals in :obj:`mne.Epochs` into several segments according to the configuration information provided by the user. These segments will be annotated by the description dictionary elements and the event type annotated in :obj:`mne.Epochs`. Here is an example case shows the use of :obj:`MNEDataset`:
-
-    .. code-block:: python
-
-        # subject index and run index of mne.Epochs
-        metadata_list = [{
-            'subject': 1,
-            'run': 3
-        }, {
-            'subject': 1,
-            'run': 7
-        }, {
-            'subject': 1,
-            'run': 11
-        }]
-
-        epochs_list = []
-        for metadata in metadata_list:
-            physionet_path = mne.datasets.eegbci.load_data(metadata['subject'],
-                                                        metadata['run'],
-                                                        update_path=False)[0]
-
-            raw = mne.io.read_raw_edf(physionet_path, preload=True, stim_channel='auto')
-            events, _ = mne.events_from_annotations(raw)
-            picks = mne.pick_types(raw.info,
-                                meg=False,
-                                eeg=True,
-                                stim=False,
-                                eog=False,
-                                exclude='bads')
-            # init Epochs with raw EEG signals and corresponding event annotations
-            epochs_list.append(mne.Epochs(raw, events, picks=picks))
-
-        # split into chunks of 160 data points (1s)
-        dataset = MNEDataset(epochs_list=epochs_list,
-                            metadata_list=metadata_list,
-                            chunk_size=160,
-                            overlap=0,
-                            num_channel=60,
-                            io_path=io_path,
-                            offline_transform=transforms.Compose(
-                                [transforms.BandDifferentialEntropy()]),
-                            online_transform=transforms.ToTensor(),
-                            label_transform=transforms.Compose([
-                                transforms.Select('event')
-                            ]),
-                            num_worker=2)
-        print(dataset[0])
-        # EEG signal (torch.Tensor[60, 4]),
-        # coresponding baseline signal (torch.Tensor[60, 4]),
-        # label (int)
-
-    In particular, TorchEEG utilizes the producer-consumer model to allow multi-process data preprocessing. If your data preprocessing is time consuming, consider increasing :obj:`num_worker` for higher speedup. If running under Windows, please use the proper idiom in the main module:
-
-    .. code-block:: python
-    
-        if __name__ == '__main__':
-            # subject index and run index of mne.Epochs
-            metadata_list = [{
-                'subject': 1,
-                'run': 3
-            }, {
-                'subject': 1,
-                'run': 7
-            }, {
-                'subject': 1,
-                'run': 11
-            }]
-
-            epochs_list = []
-            for metadata in metadata_list:
-                physionet_path = mne.datasets.eegbci.load_data(metadata['subject'],
-                                                            metadata['run'],
-                                                            update_path=False)[0]
-
-                raw = mne.io.read_raw_edf(physionet_path, preload=True, stim_channel='auto')
-                events, _ = mne.events_from_annotations(raw)
-                picks = mne.pick_types(raw.info,
-                                    meg=False,
-                                    eeg=True,
-                                    stim=False,
-                                    eog=False,
-                                    exclude='bads')
-                # init Epochs with raw EEG signals and corresponding event annotations
-                epochs_list.append(mne.Epochs(raw, events, picks=picks))
-
-            # split into chunks of 160 data points (1s)
-            dataset = MNEDataset(epochs_list=epochs_list,
-                                metadata_list=metadata_list,
-                                chunk_size=160,
-                                overlap=0,
-                                num_channel=60,
-                                io_path=io_path,
-                                offline_transform=transforms.Compose(
-                                    [transforms.BandDifferentialEntropy()]),
-                                online_transform=transforms.ToTensor(),
-                                label_transform=transforms.Compose([
-                                    transforms.Select('event')
-                                ]),
-                                num_worker=2)
-            print(dataset[0])
-            # EEG signal (torch.Tensor[60, 4]),
-            # coresponding baseline signal (torch.Tensor[60, 4]),
-            # label (int)
-
-    Args:
-        epochs_list (list): A list of :obj:`mne.Epochs`. :obj:`MNEDataset` will divide the signals in :obj:`mne.Epochs` into several segments according to the :obj:`chunk_size` and :obj:`overlap` information provided by the user. The divided segments will be transformed and cached in a unified input and output format (IO) for accessing. You can also pass through the path list of the mne.Epochs file (can be obtained by .save()).
-        metadata_list (list): A list of dictionaries of the same length as :obj:`epochs_list`. Each of these dictionaries is annotated with meta-information about :obj:`mne.Epochs`, such as subject index, experimental dates, etc. These annotated meta-information will be added to the element corresponding to :obj:`mne.Epochs` for use as labels for the sample.
-        chunk_size (int): Number of data points included in each EEG chunk as training or test samples. If set to -1, the EEG signal of a trial is used as a sample of a chunk. If set to -1, the EEG signal is not segmented, and the length of the chunk is the length of the event. (default: :obj:`-1`)
-        overlap (int): The number of overlapping data points between different chunks when dividing EEG chunks. (default: :obj:`0`)
-        num_channel (int): Number of channels used. If set to -1, all electrodes are used (default: :obj:`-1`)
-        online_transform (Callable, optional): The transformation of the EEG signals and baseline EEG signals. The input is a :obj:`np.ndarray`, and the ouput is used as the first and second value of each element in the dataset. (default: :obj:`None`)
-        offline_transform (Callable, optional): The usage is the same as :obj:`online_transform`, but executed before generating IO intermediate results. (default: :obj:`None`)
-        label_transform (Callable, optional): The transformation of the label. The input is an information dictionary, and the ouput is used as the third value of each element in the dataset. (default: :obj:`None`)
-        before_trial (Callable, optional): The hook performed on the trial to which the sample belongs. It is performed before the offline transformation and thus typically used to implement context-dependent sample transformations, such as moving averages, etc. The input and output of this hook function should be a :obj:`mne.Epoch`.
-        after_trial (Callable, optional): The hook performed on the trial to which the sample belongs. It is performed after the offline transformation and thus typically used to implement context-dependent sample transformations, such as moving averages, etc. The input and output of this hook function should be a sequence of dictionaries representing a sequence of EEG samples. Each dictionary contains two key-value pairs, indexed by :obj:`eeg` (the EEG signal matrix) and :obj:`key` (the index in the database) respectively
-        io_path (str): The path to generated unified data IO, cached as an intermediate result. (default: :obj:`./io/mne`)
-        io_size (int): Maximum size database may grow to; used to size the memory mapping. If database grows larger than ``map_size``, an exception will be raised and the user must close and reopen. (default: :obj:`10485760`)
-        io_mode (str): Storage mode of EEG signal. When io_mode is set to :obj:`lmdb`, TorchEEG provides an efficient database (LMDB) for storing EEG signals. LMDB may not perform well on limited operating systems, where a file system based EEG signal storage is also provided. When io_mode is set to :obj:`pickle`, pickle-based persistence files are used. (default: :obj:`lmdb`)
-        num_worker (int): Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process. (default: :obj:`0`)
-        verbose (bool): Whether to display logs during processing, such as progress bars, etc. (default: :obj:`True`)
-        in_memory (bool): Whether to load the entire dataset into memory. If :obj:`in_memory` is set to True, then the first time an EEG sample is read, the entire dataset is loaded into memory for subsequent retrieval. Otherwise, the dataset is stored on disk to avoid the out-of-memory problem. (default: :obj:`False`)    
-    '''
-    def __init__(self,
-                 epochs_list: List[mne.Epochs],
-                 metadata_list: List[Dict],
-                 chunk_size: int = -1,
-                 overlap: int = 0,
-                 num_channel: int = -1,
-                 online_transform: Union[None, Callable] = None,
-                 offline_transform: Union[None, Callable] = None,
-                 label_transform: Union[None, Callable] = None,
-                 before_trial: Union[None, Callable] = None,
-                 after_trial: Union[Callable, None] = None,
-                 io_path: str = './io/mne',
-                 io_size: int = 10485760,
-                 io_mode: str = 'lmdb',
-                 num_worker: int = 0,
-                 verbose: bool = True,
-                 in_memory: bool = False):
-        # pass all arguments to super class
-        params = {
-            'chunk_size': chunk_size,
-            'overlap': overlap,
-            'num_channel': num_channel,
-            'online_transform': online_transform,
-            'offline_transform': offline_transform,
-            'label_transform': label_transform,
-            'before_trial': before_trial,
-            'after_trial': after_trial
-        }
-        self.__dict__.update(params)
-
-        self.io_path = io_path
-        self.io_size = io_size
-        self.io_mode = io_mode
-        self.in_memory = in_memory
-        self.num_worker = num_worker
-        self.verbose = verbose
-
-        # new IO
-        if not self.exist(self.io_path):
-            print(
-                f'dataset does not exist at path {self.io_path}, generating files to path...'
-            )
-            # make the root dictionary
-            os.makedirs(self.io_path, exist_ok=True)
-
-            # init sub-folders
-            meta_info_io_path = os.path.join(self.io_path, 'info.csv')
-            eeg_signal_io_path = os.path.join(self.io_path, 'eeg')
-
-            MetaInfoIO(meta_info_io_path)
-            EEGSignalIO(eeg_signal_io_path,
-                        io_size=self.io_size,
-                        io_mode=self.io_mode)
-
-            if self.num_worker == 0:
-                lock = MockLock()  # do nothing, just for compatibility
-                for file in tqdm(self._set_files(epochs_list=epochs_list,
-                                                 metadata_list=metadata_list,
-                                                 io_path=io_path,
-                                                 **params),
-                                 disable=not self.verbose,
-                                 desc="[PROCESS]"):
-                    self._process_file(io_path=self.io_path,
-                                       io_size=self.io_size,
-                                       io_mode=self.io_mode,
-                                       file=file,
-                                       lock=lock,
-                                       _load_data=self._load_data,
-                                       **params)
-            else:
-                # lock for lmdb writter, LMDB only allows single-process writes
-                manager = Manager()
-                lock = manager.Lock()
-
-                Parallel(n_jobs=self.num_worker)(
-                    delayed(self._process_file)(io_path=io_path,
-                                                io_size=io_size,
-                                                io_mode=io_mode,
-                                                file=file,
-                                                lock=lock,
-                                                _load_data=self._load_data,
-                                                **params)
-                    for file in tqdm(self._set_files(
-                        epochs_list=epochs_list,
-                        metadata_list=metadata_list,
-                        io_path=io_path,
-                        **params),
-                                     disable=not self.verbose,
-                                     desc="[PROCESS]"))
-
-        print(
-            f'dataset already exists at path {self.io_path}, reading from path...'
-        )
-
-        meta_info_io_path = os.path.join(self.io_path, 'info.csv')
-        eeg_signal_io_path = os.path.join(self.io_path, 'eeg')
-
-        info_io = MetaInfoIO(meta_info_io_path)
-        self.eeg_io = EEGSignalIO(eeg_signal_io_path,
-                                  io_size=self.io_size,
-                                  io_mode=self.io_mode)
-
-        self.info = info_io.read_all()
-
-    @staticmethod
-    def _load_data(file: Any = None,
-                   chunk_size: int = -1,
-                   overlap: int = 0,
-                   num_channel: int = -1,
-                   before_trial: Union[None, Callable] = None,
-                   offline_transform: Union[None, Callable] = None,
-                   after_trial: Union[None, Callable] = None,
-                   **kwargs):
-        mne.set_log_level('CRITICAL')
-
-        tmp_path, metadata, block_id = file
-        epochs = mne.read_epochs(tmp_path, preload=True)
-
-        assert (epochs.tmax - epochs.tmin) * epochs.info[
-            "sfreq"] >= chunk_size, f'chunk_size cannot be larger than (tmax - tmin) * sfreq. Here, tmax is set to {epochs.tmax}, tmin is set to {epochs.tmin}, and sfreq is {epochs.info["sfreq"]}. In the current configuration, chunk_size {chunk_size} is greater than {(epochs.tmax - epochs.tmin) * epochs.info["sfreq"]}!'
-
-        if chunk_size <= 0:
-            chunk_size = len(epochs.times)
-
-        if num_channel == -1:
-            num_channel = len(epochs.info['chs'])
-
-        trial_event_index = epochs.events[:, 2]
-        trial_start_at_list = epochs.events[:, 0]
-
-        trial_end_at = len(epochs.times) - chunk_size
-
-        clip_sample_start_at_list = np.arange(0, trial_end_at + 1,
-                                              chunk_size - overlap)
-
-        sample_events = [[clip_sample_start_at, chunk_size, -1]
-                         for clip_sample_start_at in clip_sample_start_at_list]
-
-        epoch_meta_info = metadata
-
-        write_pointer = 0
-
-        # for loop of trials
-        for trial_id, trial in enumerate(epochs):
-            # split sample from epochs
-            start_at_list = clip_sample_start_at_list + trial_start_at_list[
-                trial_id]
-            end_at_list = clip_sample_start_at_list + trial_start_at_list[
-                trial_id] + chunk_size
-            event_index_list = len(sample_events) * [
-                trial_event_index[trial_id]
-            ]
-
-            trial_samples = mne.Epochs(mne.io.RawArray(trial, epochs.info),
-                                       sample_events,
-                                       baseline=None,
-                                       tmin=0,
-                                       tmax=(chunk_size - 1) /
-                                       epochs.info["sfreq"])
-            trial_samples.drop_bad(reject=None, flat=None)
-            if before_trial:
-                trial_samples = before_trial(trial_samples)
-
-            # for loop of samples
-            trial_queue = []
-            for i, trial_signal in enumerate(trial_samples.get_data()):
-                t_eeg = trial_signal[:num_channel, :]
-                if not offline_transform is None:
-                    t = offline_transform(eeg=trial_signal[:num_channel, :])
-                    t_eeg = t['eeg']
-
-                clip_id = f'{block_id}_{write_pointer}'
-                write_pointer += 1
-
-                record_info = {
-                    'trial_id': trial_id,
-                    'start_at': start_at_list[i],
-                    'end_at': end_at_list[i],
-                    'event': event_index_list[i],
-                    'clip_id': clip_id
-                }
-                record_info.update(epoch_meta_info)
-                if after_trial:
-                    trial_queue.append({
-                        'eeg': t_eeg,
-                        'key': clip_id,
-                        'info': record_info
-                    })
-                else:
-                    yield {'eeg': t_eeg, 'key': clip_id, 'info': record_info}
-
-            if len(trial_queue) and after_trial:
-                trial_queue = after_trial(trial_queue)
-                for obj in trial_queue:
-                    assert 'eeg' in obj and 'key' in obj and 'info' in obj, 'after_trial must return a list of dictionaries, where each dictionary corresponds to an EEG sample, containing `eeg`, `key` and `info` as keys.'
-                    yield obj
-
-    @staticmethod
-    def _set_files(epochs_list: Union[List[str], List[mne.Epochs]],
-                   metadata_list: List[Dict[str, Any]], io_path: str, **kwargs):
-        epochs_metadata_block_id_list = []
-        for block_id, (epochs,
-                       metadata) in enumerate(zip(epochs_list, metadata_list)):
-            if isinstance(epochs, str):
-                epochs_path = epochs
-            else:
-                if not os.path.exists(os.path.join(io_path, 'tmp')):
-                    os.makedirs(os.path.join(io_path, 'tmp'))
-                epochs_path = os.path.join(io_path, 'tmp', f'{block_id}.epochs')
-                epochs.save(epochs_path)
-
-            epochs_metadata_block_id_list.append(
-                (epochs_path, metadata, block_id))
-
-        return epochs_metadata_block_id_list
-
-    def __getitem__(self, index: int) -> Tuple:
-        info = self.read_info(index)
-
-        eeg_index = str(info['clip_id'])
-        eeg = self.read_eeg(eeg_index)
-
-        signal = eeg
-        label = info
-
-        if self.online_transform:
-            signal = self.online_transform(eeg=eeg)['eeg']
-
-        if self.label_transform:
-            label = self.label_transform(y=info)['y']
-
-        return signal, label
-
-    @property
-    def repr_body(self) -> Dict:
-        return dict(
-            super().repr_body, **{
-                'chunk_size': self.chunk_size,
-                'overlap': self.overlap,
-                'num_channel': self.num_channel,
-                'online_transform': self.online_transform,
-                'offline_transform': self.offline_transform,
-                'label_transform': self.label_transform,
-                'before_trial': self.before_trial,
-                'after_trial': self.after_trial,
-                'num_worker': self.num_worker,
-                'verbose': self.verbose,
-                'io_size': self.io_size
-            })
diff --git a/torcheeg/datasets/module/numpy_dataset.py b/torcheeg/datasets/module/numpy_dataset.py
deleted file mode 100644
index 7828abf..0000000
--- a/torcheeg/datasets/module/numpy_dataset.py
+++ /dev/null
@@ -1,317 +0,0 @@
-import os
-from typing import Any, List, Callable, Dict, Tuple, Union
-from multiprocessing import Manager
-from joblib import Parallel, delayed
-from tqdm import tqdm
-import numpy as np
-
-from torcheeg.io import EEGSignalIO, MetaInfoIO
-
-from .base_dataset import BaseDataset
-
-
-class MockLock():
-    def __enter__(self):
-        pass
-
-    def __exit__(self, exc_type, exc_value, traceback):
-        pass
-
-
-class NumpyDataset(BaseDataset):
-    r'''
-    A general dataset, this class converts EEG signals and annotations in Numpy format into dataset types, and caches the generated results in a unified input and output format (IO).
-
-    A tiny case shows the use of :obj:`NumpyDataset`:
-
-    .. code-block:: python
-
-        # Mock 100 EEG samples. Each EEG signal contains a signal of length 1 s at a sampling rate of 128 sampled by 32 electrodes.
-        X = np.random.randn(100, 32, 128)
-
-        # Mock 100 labels, denoting valence and arousal of subjects during EEG recording.
-        y = {
-            'valence': np.random.randint(10, size=100),
-            'arousal': np.random.randint(10, size=100)
-        }
-        dataset = NumpyDataset(X=X,
-                               y=y,
-                               io_path=io_path,
-                               offline_transform=transforms.Compose(
-                                   [transforms.BandDifferentialEntropy()]),
-                               online_transform=transforms.ToTensor(),
-                               label_transform=transforms.Compose([
-                                   transforms.Select('valence'),
-                                   transforms.Binary(5.0),
-                               ]),
-                               num_worker=2,
-                               num_samples_per_worker=50)
-        print(dataset[0])
-        # EEG signal (torch.Tensor[32, 4]),
-        # coresponding baseline signal (torch.Tensor[32, 4]),
-        # label (int)
-
-    In particular, TorchEEG utilizes the producer-consumer model to allow multi-process data preprocessing. If your data preprocessing is time consuming, consider increasing :obj:`num_worker` for higher speedup. If running under Windows, please use the proper idiom in the main module:
-
-    .. code-block:: python
-    
-        if __name__ == '__main__':
-            # Mock 100 EEG samples. Each EEG signal contains a signal of length 1 s at a sampling rate of 128 sampled by 32 electrodes.
-            X = np.random.randn(100, 32, 128)
-
-            # Mock 100 labels, denoting valence and arousal of subjects during EEG recording.
-            y = np.random.randint(10, size=100, 2)
-            dataset = NumpyDataset(X=X,
-                                y=y,
-                                io_path=io_path,
-                                offline_transform=transforms.Compose(
-                                    [transforms.BandDifferentialEntropy()]),
-                                online_transform=transforms.ToTensor(),
-                                label_transform=transforms.Compose([
-                                    transforms.Select('0'),
-                                    transforms.Binary(5.0),
-                                ]),
-                                num_worker=2,
-                                num_samples_per_worker=50)
-            print(dataset[0])
-            # EEG signal (torch.Tensor[32, 4]),
-            # coresponding baseline signal (torch.Tensor[32, 4]),
-            # label (int)
-
-    Args:
-        X (np.ndarray): An array in :obj:`numpy.ndarray` format representing the EEG signal samples in the dataset. The shape of the array is :obj:`[num_sample, ...]` where :obj:`num_sample` is the number of samples. :obj:`X` and :obj:`y` also allow lists of strings to be used together, representing paths to corresponding files of :obj:`X` and :obj:`y` (generated using :obj:`np.save`).
-        y (np.ndarray):An array in :obj:`numpy.ndarray` format representing the labels of EEG signal samples, and the values are lists of labels whose length is consistent with the EEG signal samples. The name of the label is automatically generated from the string corresponding to its index.
-        online_transform (Callable, optional): The transformation of the EEG signals and baseline EEG signals. The input is a :obj:`np.ndarray`, and the ouput is used as the first and second value of each element in the dataset. (default: :obj:`None`)
-        offline_transform (Callable, optional): The usage is the same as :obj:`online_transform`, but executed before generating IO intermediate results. (default: :obj:`None`)
-        label_transform (Callable, optional): The transformation of the label. The input is an information dictionary, and the ouput is used as the third value of each element in the dataset. (default: :obj:`None`)
-        before_trial (Callable, optional): The hook performed on the trial to which the sample belongs. It is performed before the offline transformation and thus typically used to implement context-dependent sample transformations, such as moving averages, etc. The input and output of this hook function should be a :obj:`np.ndarray`, whose shape is (number of EEG samples per trial, ...).
-        after_trial (Callable, optional): The hook performed on the trial to which the sample belongs. It is performed after the offline transformation and thus typically used to implement context-dependent sample transformations, such as moving averages, etc. The input and output of this hook function should be a sequence of dictionaries representing a sequence of EEG samples. Each dictionary contains two key-value pairs, indexed by :obj:`eeg` (the EEG signal matrix) and :obj:`key` (the index in the database) respectively
-        io_path (str): The path to generated unified data IO, cached as an intermediate result. (default: :obj:`./io/deap`)
-        io_size (int): Maximum size database may grow to; used to size the memory mapping. If database grows larger than ``map_size``, an exception will be raised and the user must close and reopen. (default: :obj:`10485760`)
-        io_mode (str): Storage mode of EEG signal. When io_mode is set to :obj:`lmdb`, TorchEEG provides an efficient database (LMDB) for storing EEG signals. LMDB may not perform well on limited operating systems, where a file system based EEG signal storage is also provided. When io_mode is set to :obj:`pickle`, pickle-based persistence files are used. (default: :obj:`lmdb`)
-        num_worker (int): How many subprocesses to use for data processing. (default: :obj:`0`)
-        num_samples_per_worker (int): The number of samples processed by each process. Once the specified number of samples are processed, the process will be destroyed and new processes will be created to perform new tasks. (default: :obj:`100`)
-        verbose (bool): Whether to display logs during processing, such as progress bars, etc. (default: :obj:`True`)
-        in_memory (bool): Whether to load the entire dataset into memory. If :obj:`in_memory` is set to True, then the first time an EEG sample is read, the entire dataset is loaded into memory for subsequent retrieval. Otherwise, the dataset is stored on disk to avoid the out-of-memory problem. (default: :obj:`False`)    
-    '''
-    def __init__(self,
-                 X: np.ndarray,
-                 y: Dict,
-                 online_transform: Union[None, Callable] = None,
-                 offline_transform: Union[None, Callable] = None,
-                 label_transform: Union[None, Callable] = None,
-                 before_trial: Union[None, Callable] = None,
-                 after_trial: Union[Callable, None] = None,
-                 io_path: str = './io/numpy',
-                 io_size: int = 10485760,
-                 io_mode: str = 'lmdb',
-                 num_worker: int = 0,
-                 num_samples_per_worker: int = 100,
-                 verbose: bool = True,
-                 in_memory: bool = False):
-        # pass all arguments to super class
-        params = {
-            'online_transform': online_transform,
-            'offline_transform': offline_transform,
-            'label_transform': label_transform,
-            'before_trial': before_trial,
-            'after_trial': after_trial,
-            'num_samples_per_worker': num_samples_per_worker
-        }
-        self.__dict__.update(params)
-
-        self.io_path = io_path
-        self.io_size = io_size
-        self.io_mode = io_mode
-        self.in_memory = in_memory
-        self.num_worker = num_worker
-        self.verbose = verbose
-
-        # new IO
-        if not self.exist(self.io_path):
-            print(
-                f'dataset does not exist at path {self.io_path}, generating files to path...'
-            )
-            # make the root dictionary
-            os.makedirs(self.io_path, exist_ok=True)
-
-            # init sub-folders
-            meta_info_io_path = os.path.join(self.io_path, 'info.csv')
-            eeg_signal_io_path = os.path.join(self.io_path, 'eeg')
-
-            MetaInfoIO(meta_info_io_path)
-            EEGSignalIO(eeg_signal_io_path,
-                        io_size=self.io_size,
-                        io_mode=self.io_mode)
-
-            if self.num_worker == 0:
-                lock = MockLock()  # do nothing, just for compatibility
-                for file in tqdm(self._set_files(X=X,
-                                                 y=y,
-                                                 io_path=io_path,
-                                                 **params),
-                                 disable=not self.verbose,
-                                 desc="[PROCESS]"):
-                    self._process_file(io_path=self.io_path,
-                                       io_size=self.io_size,
-                                       io_mode=self.io_mode,
-                                       file=file,
-                                       lock=lock,
-                                       _load_data=self._load_data,
-                                       **params)
-            else:
-                # lock for lmdb writter, LMDB only allows single-process writes
-                manager = Manager()
-                lock = manager.Lock()
-
-                Parallel(n_jobs=self.num_worker)(
-                    delayed(self._process_file)(io_path=io_path,
-                                                io_size=io_size,
-                                                io_mode=io_mode,
-                                                file=file,
-                                                lock=lock,
-                                                _load_data=self._load_data,
-                                                **params)
-                    for file in tqdm(self._set_files(
-                        X=X, y=y, io_path=io_path, **params),
-                                     disable=not self.verbose,
-                                     desc="[PROCESS]"))
-
-        print(
-            f'dataset already exists at path {self.io_path}, reading from path...'
-        )
-
-        meta_info_io_path = os.path.join(self.io_path, 'info.csv')
-        eeg_signal_io_path = os.path.join(self.io_path, 'eeg')
-
-        info_io = MetaInfoIO(meta_info_io_path)
-        self.eeg_io = EEGSignalIO(eeg_signal_io_path,
-                                  io_size=self.io_size,
-                                  io_mode=self.io_mode)
-
-        self.info = info_io.read_all()
-
-    @staticmethod
-    def _load_data(file: Any = None,
-                   before_trial: Union[None, Callable] = None,
-                   offline_transform: Union[None, Callable] = None,
-                   after_trial: Union[None, Callable] = None,
-                   **kwargs):
-        X_file_path, y_file_path, file_id = file
-        X = np.load(X_file_path)
-        y = np.load(y_file_path)
-
-        if before_trial:
-            X = before_trial(X)
-
-        trial_queue = []
-        for write_pointer, clip_sample in enumerate(X):
-
-            t_eeg = clip_sample
-
-            if not offline_transform is None:
-                t = offline_transform(eeg=clip_sample)
-                t_eeg = t['eeg']
-
-            clip_id = f'{file_id}_{write_pointer}'
-
-            # record meta info for each signal
-            record_info = {'clip_id': clip_id}
-            # y is np.ndarray (n_samples, n_labels)
-            record_info.update(
-                {f'{i}': y[write_pointer, i]
-                 for i in range(y.shape[1])})
-
-            if after_trial:
-                trial_queue.append({
-                    'eeg': t_eeg,
-                    'key': clip_id,
-                    'info': record_info
-                })
-            else:
-                yield {'eeg': t_eeg, 'key': clip_id, 'info': record_info}
-
-        if len(trial_queue) and after_trial:
-            trial_queue = after_trial(trial_queue)
-            for obj in trial_queue:
-                assert 'eeg' in obj and 'key' in obj and 'info' in obj, 'after_trial must return a list of dictionaries, where each dictionary corresponds to an EEG sample, containing `eeg`, `key` and `info` as keys.'
-                yield obj
-
-    @staticmethod
-    def _set_files(X: Union[np.ndarray, List[str]] = None,
-                   y: Union[np.ndarray, List[str]] = None,
-                   io_path: str = None,
-                   num_samples_per_worker: int = 100,
-                   **kwargs):
-        # check if X is a list of str
-        X_str = isinstance(X, list) and all([isinstance(x, str) for x in X])
-        y_str = isinstance(y, list) and all([isinstance(y, str) for y in y])
-
-        if X_str and y_str:
-            X_y_file_id_list = []
-            for block_id, (X_file_path, y_file_path) in enumerate(zip(X, y)):
-                X_y_file_id_list.append((X_file_path, y_file_path, block_id))
-            return X_y_file_id_list
-
-        X_ndarray = isinstance(X, np.ndarray)
-        y_ndarray = isinstance(y, np.ndarray)
-
-        if X_ndarray and y_ndarray:
-            indices = np.arange(len(X))
-            block_samples_list = np.array_split(
-                indices,
-                len(X) // num_samples_per_worker)
-
-            X_y_file_id_list = []
-            for block_id, sample_indices in enumerate(block_samples_list):
-                X_file = X[sample_indices]
-                y_file = y[sample_indices]
-                # save block to disk
-                if not os.path.exists(os.path.join(io_path, 'tmp')):
-                    os.makedirs(os.path.join(io_path, 'tmp'))
-
-                X_file_path = os.path.join(io_path, 'tmp', f'{block_id}_x.npy')
-                np.save(X_file_path, X_file)
-
-                y_file_path = os.path.join(io_path, 'tmp', f'{block_id}_y.npy')
-                np.save(y_file_path, y_file)
-
-                X_y_file_id_list.append((X_file_path, y_file_path, block_id))
-
-            return X_y_file_id_list
-
-        raise ValueError(
-            'X and y must be either a list of paths to np.ndarray, or a np.ndarray.'
-        )
-
-    def __getitem__(self, index: int) -> Tuple:
-        info = self.read_info(index)
-
-        eeg_index = str(info['clip_id'])
-        eeg = self.read_eeg(eeg_index)
-
-        signal = eeg
-        label = info
-
-        if self.online_transform:
-            signal = self.online_transform(eeg=eeg)['eeg']
-
-        if self.label_transform:
-            label = self.label_transform(y=info)['y']
-
-        return signal, label
-
-    @property
-    def repr_body(self) -> Dict:
-        return dict(
-            super().repr_body, **{
-                'online_transform': self.online_transform,
-                'offline_transform': self.offline_transform,
-                'label_transform': self.label_transform,
-                'before_trial': self.before_trial,
-                'after_trial': self.after_trial,
-                'num_worker': self.num_worker,
-                'num_samples_per_worker': self.num_samples_per_worker,
-                'verbose': self.verbose,
-                'io_size': self.io_size
-            })
diff --git a/torcheeg/datasets/module/personal_identification/m3cv.py b/torcheeg/datasets/module/personal_identification/m3cv.py
index 85cc312..b18e309 100644
--- a/torcheeg/datasets/module/personal_identification/m3cv.py
+++ b/torcheeg/datasets/module/personal_identification/m3cv.py
@@ -163,7 +163,7 @@ def __init__(self,
         self.__dict__.update(params)
 
     @staticmethod
-    def _load_data(file: Any = None,
+    def process_record(file: Any = None,
                    root_path: str = './aistudio',
                    subset: str = 'Enrollment',
                    chunk_size: int = 1000,
@@ -247,8 +247,7 @@ def _load_data(file: Any = None,
                     assert 'eeg' in obj and 'key' in obj and 'info' in obj, 'after_trial must return a list of dictionaries, where each dictionary corresponds to an EEG sample, containing `eeg`, `key` and `info` as keys.'
                     yield obj
 
-    @staticmethod
-    def _set_files(root_path: str = './aistudio',
+    def set_records(self, root_path: str = './aistudio',
                    subset: str = 'Enrollment',
                    **kwargs):
         assert subset in [
@@ -271,7 +270,8 @@ def __getitem__(self, index: int) -> Tuple:
         info = self.read_info(index)
 
         eeg_index = str(info['clip_id'])
-        eeg = self.read_eeg(eeg_index)
+        eeg_record = str(info['_record_id'])
+        eeg = self.read_eeg(eeg_record, eeg_index)
 
         signal = eeg
         label = info
diff --git a/torcheeg/datasets/module/sspev/tsu_benchmark.py b/torcheeg/datasets/module/sspev/tsu_benchmark.py
index 4d9f032..c0c3676 100644
--- a/torcheeg/datasets/module/sspev/tsu_benchmark.py
+++ b/torcheeg/datasets/module/sspev/tsu_benchmark.py
@@ -156,7 +156,7 @@ def __init__(self,
         self.__dict__.update(params)
 
     @staticmethod
-    def _load_data(file: Any = None,
+    def process_record(file: Any = None,
                    root_path: str = './TSUBenchmark',
                    chunk_size: int = 250,
                    overlap: int = 0,
@@ -248,8 +248,7 @@ def _load_data(file: Any = None,
                         assert 'eeg' in obj and 'key' in obj and 'info' in obj, 'after_trial must return a list of dictionaries, where each dictionary corresponds to an EEG sample, containing `eeg`, `key` and `info` as keys.'
                         yield obj
 
-    @staticmethod
-    def _set_files(**kwargs):
+    def set_records(self, **kwargs):
         root_path = kwargs.pop('root_path', './TSUBenchmark')  # str
 
         file_list = os.listdir(root_path)
@@ -264,7 +263,8 @@ def __getitem__(self, index: int) -> Tuple[any, any, int, int, int]:
         info = self.read_info(index)
 
         eeg_index = str(info['clip_id'])
-        eeg = self.read_eeg(eeg_index)
+        eeg_record = str(info['_record_id'])
+        eeg = self.read_eeg(eeg_record, eeg_index)
 
         signal = eeg
         label = info
diff --git a/torcheeg/io/eeg_signal.py b/torcheeg/io/eeg_signal.py
index 8154531..5cc9d56 100644
--- a/torcheeg/io/eeg_signal.py
+++ b/torcheeg/io/eeg_signal.py
@@ -185,3 +185,19 @@ def read_eeg_in_memory(self, key: str) -> any:
             self._in_memory = self.to_dict()
 
         return self._in_memory[key]
+
+    def write_eeg_in_memory(self, eeg: any, key: str) -> None:
+        r'''
+        Write EEG signal to memory.
+
+        .. warning::
+            This method will write all the data in EEGSignalIO into memory, which may cause memory overflow. Thus, it is only recommended for fast reading of small-scale datasets.
+
+        Args:
+            eeg (any): EEG signal samples to be written into the database.
+            key (str): The key of the EEG signal to be inserted, if not specified, it will be an auto-incrementing
+        '''
+        if self._in_memory is None:
+            self._in_memory = self.to_dict()
+        
+        self._in_memory[key] = eeg
\ No newline at end of file
diff --git a/torcheeg/io/meta_info.py b/torcheeg/io/meta_info.py
index fed3e2c..68eed17 100644
--- a/torcheeg/io/meta_info.py
+++ b/torcheeg/io/meta_info.py
@@ -33,6 +33,7 @@ class MetaInfoIO:
     def __init__(self, io_path: str) -> None:
         self.io_path = io_path
         if not os.path.exists(self.io_path):
+            os.makedirs(os.path.dirname(io_path), exist_ok=True)
             open(self.io_path, 'x').close()
             self.write_pointer = 0
         else:
@@ -82,4 +83,6 @@ def read_all(self) -> pd.DataFrame:
         Returns:
             pd.DataFrame: The EEG descriptions.
         '''
+        if os.path.getsize(self.io_path) == 0:
+            return pd.DataFrame()
         return pd.read_csv(self.io_path)
\ No newline at end of file
diff --git a/torcheeg/models/cnn/ccnn.py b/torcheeg/models/cnn/ccnn.py
index e9a03c7..e7faeaf 100644
--- a/torcheeg/models/cnn/ccnn.py
+++ b/torcheeg/models/cnn/ccnn.py
@@ -55,18 +55,8 @@ def __init__(self, in_channels: int = 4, grid_size: Tuple[int, int] = (9, 9), nu
         )
         self.lin2 = nn.Linear(1024, self.num_classes)
 
-    @property
     def feature_dim(self):
-        with torch.no_grad():
-            mock_eeg = torch.zeros(1, self.in_channels, *self.grid_size)
-
-            mock_eeg = self.conv1(mock_eeg)
-            mock_eeg = self.conv2(mock_eeg)
-            mock_eeg = self.conv3(mock_eeg)
-            mock_eeg = self.conv4(mock_eeg)
-            mock_eeg = mock_eeg.flatten(start_dim=1)
-
-            return mock_eeg.shape[1]
+        return self.grid_size[0] * self.grid_size[1] * 64
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         r'''
@@ -82,6 +72,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         x = self.conv4(x)
 
         x = x.flatten(start_dim=1)
+
         x = self.lin1(x)
         x = self.lin2(x)
         return x
\ No newline at end of file
diff --git a/torcheeg/models/cnn/eegnet.py b/torcheeg/models/cnn/eegnet.py
index 1c8d418..c42b60a 100644
--- a/torcheeg/models/cnn/eegnet.py
+++ b/torcheeg/models/cnn/eegnet.py
@@ -99,9 +99,8 @@ def __init__(self,
             nn.BatchNorm2d(self.F2, momentum=0.01, affine=True, eps=1e-3), nn.ELU(), nn.AvgPool2d((1, 8), stride=8),
             nn.Dropout(p=dropout))
 
-        self.lin = nn.Linear(self.F2 * self.feature_dim, num_classes, bias=False)
+        self.lin = nn.Linear(self.feature_dim(), num_classes, bias=False)
 
-    @property
     def feature_dim(self):
         with torch.no_grad():
             mock_eeg = torch.zeros(1, 1, self.num_electrodes, self.chunk_size)
@@ -109,7 +108,7 @@ def feature_dim(self):
             mock_eeg = self.block1(mock_eeg)
             mock_eeg = self.block2(mock_eeg)
 
-        return mock_eeg.shape[3]
+        return self.F2 * mock_eeg.shape[3]
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         r'''
@@ -124,4 +123,4 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         x = x.flatten(start_dim=1)
         x = self.lin(x)
 
-        return x
+        return x
\ No newline at end of file
diff --git a/torcheeg/models/cnn/fbcnet.py b/torcheeg/models/cnn/fbcnet.py
index 40acc3f..7c35673 100644
--- a/torcheeg/models/cnn/fbcnet.py
+++ b/torcheeg/models/cnn/fbcnet.py
@@ -205,3 +205,6 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         x = torch.flatten(x, start_dim=1)
         x = self.last_layer(x)
         return x
+
+    def feature_dim(self):
+        return self.num_S * self.in_channels * self.stride_factor
\ No newline at end of file
diff --git a/torcheeg/models/cnn/tsception.py b/torcheeg/models/cnn/tsception.py
index 1e14d2b..950c375 100644
--- a/torcheeg/models/cnn/tsception.py
+++ b/torcheeg/models/cnn/tsception.py
@@ -4,7 +4,7 @@
 
 class TSCeption(nn.Module):
     r'''
-    Continuous Convolutional Neural Network (CCNN). For more details, please refer to the following information.
+    TSCeption. For more details, please refer to the following information.
 
     - Paper: Ding Y, Robinson N, Zhang S, et al. Tsception: Capturing temporal dynamics and spatial asymmetry from EEG for emotion recognition[J]. arXiv preprint arXiv:2104.02935, 2021.
     - URL: https://arxiv.org/abs/2104.02935
@@ -122,3 +122,6 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         out = torch.squeeze(torch.mean(out, dim=-1), dim=-1)
         out = self.fc(out)
         return out
+
+    def feature_dim(self):
+        return self.num_S
diff --git a/torcheeg/models/transformer/__init__.py b/torcheeg/models/transformer/__init__.py
index ac59e0a..1ccf37e 100644
--- a/torcheeg/models/transformer/__init__.py
+++ b/torcheeg/models/transformer/__init__.py
@@ -2,3 +2,4 @@
 from .simple_vit import SimpleViT
 from .vanilla_transformer import VanillaTransformer
 from .vit import ViT
+from .conformer import Conformer
\ No newline at end of file
diff --git a/torcheeg/models/transformer/conformer.py b/torcheeg/models/transformer/conformer.py
new file mode 100644
index 0000000..f575ba1
--- /dev/null
+++ b/torcheeg/models/transformer/conformer.py
@@ -0,0 +1,260 @@
+import torch.nn as nn
+import math
+import torch
+
+import torch.nn.functional as F
+from einops import rearrange
+from einops.layers.torch import Rearrange, Reduce
+
+
+class PatchEmbedding(nn.Module):
+    def __init__(self,
+                 num_electrodes: int,
+                 hid_channels: int = 40,
+                 dropout: float = 0.5):
+        # self.patch_size = patch_size
+        super().__init__()
+
+        self.shallownet = nn.Sequential(
+            nn.Conv2d(1, hid_channels, (1, 25), (1, 1)),
+            nn.Conv2d(hid_channels, hid_channels, (num_electrodes, 1), (1, 1)),
+            nn.BatchNorm2d(hid_channels),
+            nn.ELU(),
+            nn.AvgPool2d(
+                (1, 75), (1, 15)
+            ),  # pooling acts as slicing to obtain 'patch' along the time dimension as in ViT
+            nn.Dropout(dropout),
+        )
+
+        self.projection = nn.Sequential(
+            nn.Conv2d(hid_channels, hid_channels, (1, 1), stride=(
+                1, 1)),  # transpose, conv could enhance fiting ability slightly
+            Rearrange('b e (h) (w) -> b (h w) e'),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        b, _, _, _ = x.shape
+        x = self.shallownet(x)
+        x = self.projection(x)
+        return x
+
+
+class MultiHeadAttention(nn.Module):
+    def __init__(self, hid_channels: int, heads: int, dropout: float):
+        super().__init__()
+        self.hid_channels = hid_channels
+        self.heads = heads
+        self.keys = nn.Linear(hid_channels, hid_channels)
+        self.queries = nn.Linear(hid_channels, hid_channels)
+        self.values = nn.Linear(hid_channels, hid_channels)
+        self.att_drop = nn.Dropout(dropout)
+        self.projection = nn.Linear(hid_channels, hid_channels)
+
+    def forward(self,
+                x: torch.Tensor,
+                mask: torch.Tensor = None) -> torch.Tensor:
+        queries = rearrange(self.queries(x),
+                            "b n (h d) -> b h n d",
+                            h=self.heads)
+        keys = rearrange(self.keys(x), "b n (h d) -> b h n d", h=self.heads)
+        values = rearrange(self.values(x), "b n (h d) -> b h n d", h=self.heads)
+        energy = torch.einsum('bhqd, bhkd -> bhqk', queries, keys)
+        if mask is not None:
+            fill_value = torch.finfo(torch.float32).min
+            energy.mask_fill(~mask, fill_value)
+
+        scaling = self.hid_channels**(1 / 2)
+        att = F.softmax(energy / scaling, dim=-1)
+        att = self.att_drop(att)
+        out = torch.einsum('bhal, bhlv -> bhav ', att, values)
+        out = rearrange(out, "b h n d -> b n (h d)")
+        out = self.projection(out)
+        return out
+
+
+class ResidualAdd(nn.Module):
+    def __init__(self, fn):
+        super().__init__()
+        self.fn = fn
+
+    def forward(self, x: torch.Tensor, **kwargs) -> torch.Tensor:
+        res = x
+        x = self.fn(x, **kwargs)
+        x += res
+        return x
+
+
+class FeedForwardBlock(nn.Sequential):
+    def __init__(self,
+                 hid_channels: int,
+                 expansion: int = 4,
+                 dropout: float = 0.):
+        super().__init__(
+            nn.Linear(hid_channels, expansion * hid_channels),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(expansion * hid_channels, hid_channels),
+        )
+
+
+class GELU(nn.Module):
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        return input * 0.5 * (1.0 + torch.erf(input / math.sqrt(2.0)))
+
+
+class TransformerEncoderBlock(nn.Sequential):
+    def __init__(self, hid_channels: int, heads: int, dropout: float,
+                 forward_expansion: int, forward_dropout: float):
+        super().__init__(
+            ResidualAdd(
+                nn.Sequential(nn.LayerNorm(hid_channels),
+                              MultiHeadAttention(hid_channels, heads, dropout),
+                              nn.Dropout(dropout))),
+            ResidualAdd(
+                nn.Sequential(
+                    nn.LayerNorm(hid_channels),
+                    FeedForwardBlock(hid_channels,
+                                     expansion=forward_expansion,
+                                     dropout=forward_dropout),
+                    nn.Dropout(dropout))))
+
+
+class TransformerEncoder(nn.Sequential):
+    def __init__(self,
+                 depth: int,
+                 hid_channels: int,
+                 heads: int = 10,
+                 dropout: float = 0.5,
+                 forward_expansion: int = 4,
+                 forward_dropout: float = 0.5):
+        super().__init__(*[
+            TransformerEncoderBlock(hid_channels=hid_channels,
+                                    heads=heads,
+                                    dropout=dropout,
+                                    forward_expansion=forward_expansion,
+                                    forward_dropout=forward_dropout)
+            for _ in range(depth)
+        ])
+
+
+class ClassificationHead(nn.Sequential):
+    def __init__(self,
+                 in_channels: int,
+                 num_classes: int,
+                 hid_channels: int = 32,
+                 dropout: float = 0.5):
+        super().__init__()
+        self.fc = nn.Sequential(nn.Linear(in_channels, hid_channels * 8),
+                                nn.ELU(), nn.Dropout(dropout),
+                                nn.Linear(hid_channels * 8, hid_channels),
+                                nn.ELU(), nn.Dropout(dropout),
+                                nn.Linear(hid_channels, num_classes))
+
+    def forward(self, x):
+        x = x.contiguous().view(x.size(0), -1)
+        x = self.fc(x)
+        return x
+
+
+class Conformer(nn.Module):
+    r'''
+    The EEG Conformer model is based on the paper "EEG Conformer: Convolutional Transformer for EEG Decoding and Visualization". For more details, please refer to the following information. 
+
+    It is worth noting that this model is not designed for EEG analysis, but shows good performance and can serve as a good research start.
+
+    - Paper: Song Y, Zheng Q, Liu B, et al. EEG Conformer: Convolutional Transformer for EEG Decoding and Visualization[J]. IEEE Transactions on Neural Systems and Rehabilitation Engineering, 2022.
+    - URL: https://ieeexplore.ieee.org/document/9991178
+    - Related Project: https://github.com/eeyhsong/EEG-Conformer
+
+    Below is a recommended suite for use in emotion recognition tasks:
+
+    .. code-block:: python
+
+        dataset = SEEDDataset(io_path=f'./seed',
+                              root_path='./Preprocessed_EEG',
+                              offline_transform=transforms.Compose([
+                                  transforms.MinMaxNormalize(axis=-1),
+                                  transforms.To2d()
+                              ]),
+                              online_transform=transforms.ToTensor(),
+                              label_transform=transforms.Compose([
+                                  transforms.Select('emotion'),
+                                  transforms.Lambda(lambda x: x + 1)
+                              ]))
+        model = Conformer(num_electrodes=62,
+                          sampling_rate=200,
+                          hid_channels=40,
+                          depth=6,
+                          heads=10,
+                          dropout=0.5,
+                          forward_expansion=4,
+                          forward_dropout=0.5,
+                          num_classes=2)
+
+    Args:
+        num_electrodes (int): The number of electrodes. (default: :obj:`62`)
+        sampling_rate (int): The sampling rate of EEG signals. (default: :obj:`200`)
+        hid_channels (int): The feature dimension of embeded patch. (default: :obj:`40`)
+        depth (int): The number of attention layers for each transformer block. (default: :obj:`6`)
+        heads (int): The number of attention heads for each attention layer. (default: :obj:`10`)
+        dropout (float): The dropout rate of the attention layer. (default: :obj:`0.5`)
+        forward_expansion (int): The expansion factor of the feedforward layer. (default: :obj:`4`)
+        forward_dropout (float): The dropout rate of the feedforward layer. (default: :obj:`0.5`)
+        num_classes (int): The number of classes. (default: :obj:`2`)
+    '''
+    def __init__(self,
+                 num_electrodes: int = 62,
+                 sampling_rate: int = 200,
+                 embed_dropout: float = 0.5,
+                 hid_channels: int = 40,
+                 depth: int = 6,
+                 heads: int = 10,
+                 dropout: float = 0.5,
+                 forward_expansion: int = 4,
+                 forward_dropout: float = 0.5,
+                 cls_channels: int = 32,
+                 cls_dropout: float = 0.5,
+                 num_classes: int = 2):
+        super().__init__()
+        self.num_electrodes = num_electrodes
+        self.sampling_rate = sampling_rate
+        self.embed_dropout = embed_dropout
+        self.hid_channels = hid_channels
+        self.depth = depth
+        self.heads = heads
+        self.dropout = dropout
+        self.forward_expansion = forward_expansion
+        self.forward_dropout = forward_dropout
+        self.cls_channels = cls_channels
+        self.cls_dropout = cls_dropout
+        self.num_classes = num_classes
+
+        self.embd = PatchEmbedding(num_electrodes, hid_channels, embed_dropout)
+        self.encoder = TransformerEncoder(depth,
+                                          hid_channels,
+                                          heads=heads,
+                                          dropout=dropout,
+                                          forward_expansion=forward_expansion,
+                                          forward_dropout=forward_dropout)
+        self.cls = ClassificationHead(in_channels=self.feature_dim(),
+                                      num_classes=num_classes,
+                                      hid_channels=cls_channels,
+                                      dropout=cls_dropout)
+
+    def feature_dim(self):
+        with torch.no_grad():
+            mock_eeg = torch.zeros(1, 1, self.num_electrodes,
+                                   self.sampling_rate)
+
+            mock_eeg = self.embd(mock_eeg)
+            mock_eeg = self.encoder(mock_eeg)
+            mock_eeg = mock_eeg.flatten(start_dim=1)
+
+            return mock_eeg.shape[1]
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.embd(x)
+        x = self.encoder(x)
+        x = x.flatten(start_dim=1)
+        x = self.cls(x)
+        return x
\ No newline at end of file
diff --git a/torcheeg/trainers/__init__.py b/torcheeg/trainers/__init__.py
index a3d4ced..e42e3ed 100644
--- a/torcheeg/trainers/__init__.py
+++ b/torcheeg/trainers/__init__.py
@@ -1,4 +1,5 @@
 from .classifier import ClassifierTrainer
 from .domain_adaption import *
 from .generative import *
-from .self_supervised import *
\ No newline at end of file
+from .self_supervised import *
+from .finetune import *
\ No newline at end of file
diff --git a/torcheeg/trainers/classifier.py b/torcheeg/trainers/classifier.py
index 7697e27..371460d 100644
--- a/torcheeg/trainers/classifier.py
+++ b/torcheeg/trainers/classifier.py
@@ -1,3 +1,4 @@
+import logging
 from typing import Any, Dict, List, Tuple
 
 import pytorch_lightning as pl
@@ -9,6 +10,8 @@
 
 _EVALUATE_OUTPUT = List[Dict[str, float]]  # 1 dict per DataLoader
 
+log = logging.getLogger(__name__)
+
 
 def classification_metrics(metric_list: List[str], num_classes: int):
     allowed_metrics = ['precision', 'recall', 'f1score', 'accuracy']
@@ -77,12 +80,12 @@ def __init__(self,
         self.num_classes = num_classes
         self.lr = lr
         self.weight_decay = weight_decay
-        
+
         self.devices = devices
         self.accelerator = accelerator
         self.metrics = metrics
 
-        self._ce_fn = nn.CrossEntropyLoss()
+        self.ce_fn = nn.CrossEntropyLoss()
 
         self.init_metrics(metrics, num_classes)
 
@@ -95,7 +98,11 @@ def init_metrics(self, metrics: List[str], num_classes: int) -> None:
         self.val_metrics = classification_metrics(metrics, num_classes)
         self.test_metrics = classification_metrics(metrics, num_classes)
 
-    def fit(self, train_loader: DataLoader, val_loader: DataLoader, max_epochs: int = 300, *args,
+    def fit(self,
+            train_loader: DataLoader,
+            val_loader: DataLoader,
+            max_epochs: int = 300,
+            *args,
             **kwargs) -> Any:
         r'''
         Args:
@@ -129,7 +136,7 @@ def training_step(self, batch: Tuple[torch.Tensor],
                       batch_idx: int) -> torch.Tensor:
         x, y = batch
         y_hat = self(x)
-        loss = self._ce_fn(y_hat, y)
+        loss = self.ce_fn(y_hat, y)
 
         # log to prog_bar
         self.log("train_loss",
@@ -138,7 +145,7 @@ def training_step(self, batch: Tuple[torch.Tensor],
                  on_epoch=False,
                  logger=False,
                  on_step=True)
-        
+
         for i, metric_value in enumerate(self.train_metrics.values()):
             self.log(f"train_{self.metrics[i]}",
                      metric_value(y_hat, y),
@@ -179,7 +186,7 @@ def validation_step(self, batch: Tuple[torch.Tensor],
                         batch_idx: int) -> torch.Tensor:
         x, y = batch
         y_hat = self(x)
-        loss = self._ce_fn(y_hat, y)
+        loss = self.ce_fn(y_hat, y)
 
         self.val_loss.update(loss)
         self.val_metrics.update(y_hat, y)
@@ -214,7 +221,7 @@ def test_step(self, batch: Tuple[torch.Tensor],
                   batch_idx: int) -> torch.Tensor:
         x, y = batch
         y_hat = self(x)
-        loss = self._ce_fn(y_hat, y)
+        loss = self.ce_fn(y_hat, y)
 
         self.test_loss.update(loss)
         self.test_metrics.update(y_hat, y)
@@ -246,7 +253,9 @@ def on_test_epoch_end(self) -> None:
         self.test_metrics.reset()
 
     def configure_optimizers(self):
-        optimizer = torch.optim.Adam(self.model.parameters(),
+        parameters = list(self.model.parameters())
+        trainable_parameters = list(filter(lambda p: p.requires_grad, parameters))
+        optimizer = torch.optim.Adam(trainable_parameters,
                                      lr=self.lr,
                                      weight_decay=self.weight_decay)
         return optimizer
diff --git a/torcheeg/trainers/domain_adaption/ada.py b/torcheeg/trainers/domain_adaption/ada.py
index 4b3f997..828cf30 100644
--- a/torcheeg/trainers/domain_adaption/ada.py
+++ b/torcheeg/trainers/domain_adaption/ada.py
@@ -120,25 +120,28 @@ def __init__(self,
                  weight_visit: float = 1.0,
                  weight_domain: float = 1.0,
                  weight_decay: float = 0.0,
-                 weight_scheduler: bool = False,
-                 lr_scheduler: bool = False,
+                 weight_scheduler: bool = True,
+                 lr_scheduler_gamma: float = 0.0,
+                 lr_scheduler_decay: float = 0.75,
                  warmup_epochs: int = 0,
                  devices: int = 1,
                  accelerator: str = "cpu",
                  metrics: List[str] = ["accuracy"]):
 
-        super(ADATrainer, self).__init__(extractor=extractor,
-                                         classifier=classifier,
-                                         num_classes=num_classes,
-                                         lr=lr,
-                                         weight_decay=weight_decay,
-                                         weight_domain=weight_domain,
-                                         weight_scheduler=weight_scheduler,
-                                         lr_scheduler=lr_scheduler,
-                                         warmup_epochs=warmup_epochs,
-                                         devices=devices,
-                                         accelerator=accelerator,
-                                         metrics=metrics)
+        super(ADATrainer,
+              self).__init__(extractor=extractor,
+                             classifier=classifier,
+                             num_classes=num_classes,
+                             lr=lr,
+                             weight_decay=weight_decay,
+                             weight_domain=weight_domain,
+                             weight_scheduler=weight_scheduler,
+                             lr_scheduler_gamma=lr_scheduler_gamma,
+                             lr_scheduler_decay=lr_scheduler_decay,
+                             warmup_epochs=warmup_epochs,
+                             devices=devices,
+                             accelerator=accelerator,
+                             metrics=metrics)
         self.weight_walker = weight_walker
         self.weight_visit = weight_visit
 
@@ -155,7 +158,7 @@ def training_step(self, batch: Tuple[torch.Tensor],
         domain_loss = self._domain_loss_fn(x_source_feat, x_target_feat,
                                            y_source)
 
-        task_loss = self._ce_fn(y_source_pred, y_source)
+        task_loss = self.ce_fn(y_source_pred, y_source)
         if self.current_epoch >= self.warmup_epochs:
             loss = task_loss + self.scheduled_weight_domain * domain_loss
         else:
diff --git a/torcheeg/trainers/domain_adaption/coral.py b/torcheeg/trainers/domain_adaption/coral.py
index 6246b30..975ba58 100644
--- a/torcheeg/trainers/domain_adaption/coral.py
+++ b/torcheeg/trainers/domain_adaption/coral.py
@@ -6,6 +6,24 @@
 from .mmd_like import _MMDLikeTrainer
 
 
+def compute_covariance(data):
+    n = data.size(0)
+
+    device = data.device
+
+    ones_vector = torch.ones(n).resize(1, n).to(device=device)
+
+    one_onto_D = torch.mm(ones_vector, data)
+    mult_right_terms = torch.mm(one_onto_D.t(), one_onto_D)
+    mult_right_terms = torch.div(mult_right_terms, n)
+
+    mult_left_terms = torch.mm(data.t(), data)
+    covariance_matrix = 1 / (n - 1) * torch.add(mult_left_terms, -1 *
+                                                (mult_right_terms))
+
+    return covariance_matrix
+
+
 class CORALTrainer(_MMDLikeTrainer):
     r'''
     This class supports the implementation of CORrelation ALignment (CORAL) for deep domain adaptation.
@@ -51,45 +69,44 @@ def __init__(self,
                  lr: float = 1e-4,
                  weight_decay: float = 0.0,
                  weight_domain: float = 1.0,
-                 weight_scheduler: bool = False,
-                 lr_scheduler: bool = False,
+                 weight_scheduler: bool = True,
+                 lr_scheduler_gamma: float = 0.0,
+                 lr_scheduler_decay: float = 0.75,
                  warmup_epochs: int = 0,
                  devices: int = 1,
                  accelerator: str = "cpu",
                  metrics: List[str] = ["accuracy"]):
-        super(CORALTrainer, self).__init__(extractor=extractor,
-                                              classifier=classifier,
-                                              num_classes=num_classes,
-                                              lr=lr,
-                                              weight_decay=weight_decay,
-                                              weight_domain=weight_domain,
-                                              weight_scheduler=weight_scheduler,
-                                              lr_scheduler=lr_scheduler,
-                                              warmup_epochs=warmup_epochs,
-                                              devices=devices,
-                                              accelerator=accelerator,
-                                              metrics=metrics)
+        super(CORALTrainer,
+              self).__init__(extractor=extractor,
+                             classifier=classifier,
+                             num_classes=num_classes,
+                             lr=lr,
+                             weight_decay=weight_decay,
+                             weight_domain=weight_domain,
+                             weight_scheduler=weight_scheduler,
+                             lr_scheduler_gamma=lr_scheduler_gamma,
+                             lr_scheduler_decay=lr_scheduler_decay,
+                             warmup_epochs=warmup_epochs,
+                             devices=devices,
+                             accelerator=accelerator,
+                             metrics=metrics)
 
     def _domain_loss_fn(self, x_source_feat: torch.Tensor,
                         x_target_feat: torch.Tensor) -> torch.Tensor:
 
-        # Compute the loss value
-        batch_size = x_source_feat.shape[1]
-
-        # source covariance
-        x_source_devi = torch.mean(x_source_feat, 1,
-                                   keepdim=True) - x_source_feat
-        x_source_cova = torch.matmul(torch.transpose(x_source_devi, 0, 1),
-                                     x_source_devi)
-
-        # target covariance
-        x_target_devi = torch.mean(x_target_feat, 1,
-                                   keepdim=True) - x_target_feat
-        x_target_cova = torch.matmul(torch.transpose(x_target_devi, 0, 1),
-                                     x_target_devi)
-        # frobenius norm between source and target
-        loss = torch.mean(
-            torch.mul((x_source_cova - x_target_cova),
-                      (x_source_cova - x_target_cova))) / (4 * batch_size * 4)
+        d = x_source_feat.size(
+            1)  # d-dimensional vectors (same for source, target)
+
+        source_covariance = compute_covariance(x_source_feat)
+        target_covariance = compute_covariance(x_target_feat)
+
+        # take Frobenius norm (https://pytorch.org/docs/stable/torch.html)
+        loss = torch.norm(torch.mul((source_covariance - target_covariance),
+                                    (source_covariance - target_covariance)),
+                          p="fro")
+
+        # loss = torch.norm(torch.mm((source_covariance-target_covariance),
+        # 							(source_covariance-target_covariance)), p="fro")
 
+        loss = loss / (4 * d * d)
         return loss
\ No newline at end of file
diff --git a/torcheeg/trainers/domain_adaption/dan.py b/torcheeg/trainers/domain_adaption/dan.py
index 5bd8f45..ba499a5 100644
--- a/torcheeg/trainers/domain_adaption/dan.py
+++ b/torcheeg/trainers/domain_adaption/dan.py
@@ -95,8 +95,9 @@ def __init__(self,
                  lr: float = 1e-4,
                  weight_decay: float = 0.0,
                  weight_domain: float = 1.0,
-                 weight_scheduler: bool = False,
-                 lr_scheduler: bool = False,
+                 weight_scheduler: bool = True,
+                 lr_scheduler_gamma: float = 0.0,
+                 lr_scheduler_decay: float = 0.75,
                  warmup_epochs: int = 0,
                  devices: int = 1,
                  accelerator: str = "cpu",
@@ -107,8 +108,9 @@ def __init__(self,
                                          lr=lr,
                                          weight_decay=weight_decay,
                                          weight_domain=weight_domain,
-                                         weight_scheduler=weight_scheduler,
-                                         lr_scheduler=lr_scheduler,
+                                        weight_scheduler=weight_scheduler,
+                                         lr_scheduler_gamma=lr_scheduler_gamma,
+                                         lr_scheduler_decay=lr_scheduler_decay,
                                          warmup_epochs=warmup_epochs,
                                          devices=devices,
                                          accelerator=accelerator,
diff --git a/torcheeg/trainers/domain_adaption/dann.py b/torcheeg/trainers/domain_adaption/dann.py
index 3d67426..c9208b0 100644
--- a/torcheeg/trainers/domain_adaption/dann.py
+++ b/torcheeg/trainers/domain_adaption/dann.py
@@ -53,7 +53,8 @@ def __init__(self,
                  weight_decay: float = 0.0,
                  weight_domain: float = 1.0,
                  alpha_scheduler: bool = True,
-                 lr_scheduler: bool = False,
+                 lr_scheduler_gamma: float = 0.0,
+                 lr_scheduler_decay: float = 0.75,
                  warmup_epochs: int = 0,
                  devices: int = 1,
                  accelerator: str = "cpu",
@@ -67,7 +68,8 @@ def __init__(self,
                                           weight_decay=weight_decay,
                                           weight_domain=weight_domain,
                                           alpha_scheduler=alpha_scheduler,
-                                          lr_scheduler=lr_scheduler,
+                                          lr_scheduler_gamma=lr_scheduler_gamma,
+                                          lr_scheduler_decay=lr_scheduler_decay,
                                           warmup_epochs=warmup_epochs,
                                           devices=devices,
                                           accelerator=accelerator,
diff --git a/torcheeg/trainers/domain_adaption/dann_like.py b/torcheeg/trainers/domain_adaption/dann_like.py
index 58bef17..910dffc 100644
--- a/torcheeg/trainers/domain_adaption/dann_like.py
+++ b/torcheeg/trainers/domain_adaption/dann_like.py
@@ -1,3 +1,4 @@
+import logging
 from itertools import chain
 from typing import List, Tuple
 
@@ -11,6 +12,8 @@
 
 from ..classifier import ClassifierTrainer, classification_metrics
 
+log = logging.getLogger(__name__)
+
 
 class DualDataLoader:
     def __init__(self, ref_dataloader: DataLoader,
@@ -57,7 +60,8 @@ def __init__(self,
                  weight_decay: float = 0.0,
                  weight_domain: float = 1.0,
                  alpha_scheduler: bool = True,
-                 lr_scheduler: bool = False,
+                 lr_scheduler_gamma: float = 0.0,
+                 lr_scheduler_decay: float = 0.75,
                  warmup_epochs: int = 0,
                  devices: int = 1,
                  accelerator: str = "cpu",
@@ -72,7 +76,9 @@ def __init__(self,
         self.weight_decay = weight_decay
         self.weight_domain = weight_domain
         self.alpha_scheduler = alpha_scheduler
-        self.lr_scheduler = lr_scheduler
+        self.lr_scheduler = not lr_scheduler_gamma == 0.0
+        self.lr_scheduler_gamma = lr_scheduler_gamma
+        self.lr_scheduler_decay = lr_scheduler_decay
         self.warmup_epochs = warmup_epochs
 
         self.devices = devices
@@ -83,7 +89,7 @@ def __init__(self,
         self.metrics = metrics
         self.init_metrics(metrics, num_classes)
 
-        self._ce_fn = nn.CrossEntropyLoss()
+        self.ce_fn = nn.CrossEntropyLoss()
 
         self.num_batches = None  # init in 'fit' method
         self.non_warmup_epochs = None  # init in 'fit' method
@@ -138,7 +144,8 @@ def on_train_batch_start(self, batch, batch_idx):
                 self.scheduled_alpha = 2.0 / (1.0 + np.exp(-10 * p)) - 1
 
             if self.lr_scheduler:
-                self.lr_factor = 1.0 / ((1.0 + 10 * p)**0.75)
+                self.lr_factor = 1.0 / ((1.0 + self.lr_scheduler_gamma * p)**
+                                        self.lr_scheduler_decay)
 
     def _domain_loss_fn(self, x_source_feat: torch.Tensor,
                         x_target_feat: torch.Tensor) -> torch.Tensor:
@@ -160,16 +167,16 @@ def training_step(self, batch: Tuple[torch.Tensor],
                                                 self.scheduled_alpha)
         y_target_disc = self.domain_classifier(x_target_feat_r)
 
-        task_loss = self._ce_fn(y_source_pred, y_source)
-        domain_loss = self._ce_fn(
+        task_loss = self.ce_fn(y_source_pred, y_source)
+        domain_loss = self.ce_fn(
             y_source_disc,
             torch.zeros(len(y_source_disc),
                         dtype=torch.long,
-                        device=x_source.device)) + self._ce_fn(
+                        device=x_source.device)) + self.ce_fn(
                             y_target_disc,
                             torch.ones(len(y_target_disc),
-                                        dtype=torch.long,
-                                        device=x_target.device))
+                                       dtype=torch.long,
+                                       device=x_target.device))
 
         if self.current_epoch >= self.warmup_epochs:
             loss = task_loss + self.weight_domain * domain_loss
diff --git a/torcheeg/trainers/domain_adaption/ddc.py b/torcheeg/trainers/domain_adaption/ddc.py
index ad32a0b..80cf59a 100644
--- a/torcheeg/trainers/domain_adaption/ddc.py
+++ b/torcheeg/trainers/domain_adaption/ddc.py
@@ -6,8 +6,7 @@
 from .mmd_like import _MMDLikeTrainer
 
 
-def maximum_mean_discrepancy_linear(x_source,
-                             x_target):
+def maximum_mean_discrepancy_linear(x_source, x_target):
     delta = x_source - x_target
     loss = torch.mean(torch.mm(delta, torch.transpose(delta, 0, 1)))
     return loss
@@ -58,24 +57,27 @@ def __init__(self,
                  lr: float = 1e-4,
                  weight_decay: float = 0.0,
                  weight_domain: float = 1.0,
-                 weight_scheduler: bool = False,
-                 lr_scheduler: bool = False,
+                 weight_scheduler: bool = True,
+                 lr_scheduler_gamma: float = 0.0,
+                 lr_scheduler_decay: float = 0.75,
                  warmup_epochs: int = 0,
                  devices: int = 1,
                  accelerator: str = "cpu",
                  metrics: List[str] = ["accuracy"]):
-        super(DDCTrainer, self).__init__(extractor=extractor,
-                                         classifier=classifier,
-                                         num_classes=num_classes,
-                                         lr=lr,
-                                         weight_decay=weight_decay,
-                                         weight_domain=weight_domain,
-                                         weight_scheduler=weight_scheduler,
-                                         lr_scheduler=lr_scheduler,
-                                         warmup_epochs=warmup_epochs,
-                                         devices=devices,
-                                         accelerator=accelerator,
-                                         metrics=metrics)
+        super(DDCTrainer,
+              self).__init__(extractor=extractor,
+                             classifier=classifier,
+                             num_classes=num_classes,
+                             lr=lr,
+                             weight_decay=weight_decay,
+                             weight_domain=weight_domain,
+                             weight_scheduler=weight_scheduler,
+                             lr_scheduler_gamma=lr_scheduler_gamma,
+                             lr_scheduler_decay=lr_scheduler_decay,
+                             warmup_epochs=warmup_epochs,
+                             devices=devices,
+                             accelerator=accelerator,
+                             metrics=metrics)
 
     def _domain_loss_fn(self, x_source_feat: torch.Tensor,
                         x_target_feat: torch.Tensor) -> torch.Tensor:
diff --git a/torcheeg/trainers/domain_adaption/jan.py b/torcheeg/trainers/domain_adaption/jan.py
index 605390f..4cba429 100644
--- a/torcheeg/trainers/domain_adaption/jan.py
+++ b/torcheeg/trainers/domain_adaption/jan.py
@@ -76,7 +76,7 @@ class JANTrainer(_MMDLikeTrainer):
         mul_kernel (tuple of float): The kernel multiplier for the Gaussian kernel. (default: :obj:`(2.0, 2.0)`)
         num_kernel (tuple of int): The number of kernels for the Gaussian kernel. (default: :obj:`(5, 1)`)
         weight_domain (float): The weight of the associative loss (default: :obj:`1.0`)
-        weight_scheduler (bool): Whether to use a scheduler for the weight of the associative loss, growing from 0 to 1 following the schedule from the DANN paper. (default: :obj:`False`)
+        lr_scheduler (bool): Whether to use a scheduler for the weight of the associative loss, growing from 0 to 1 following the schedule from the DANN paper. (default: :obj:`False`)
         lr_scheduler (bool): Whether to use a scheduler for the learning rate, as defined in the DANN paper. (default: :obj:`False`)
         warmup_epochs (int): The number of epochs for the warmup phase, during which the weight of the associative loss is 0. (default: :obj:`0`)
         devices (int): The number of devices to use. (default: :obj:`1`)
@@ -95,25 +95,28 @@ def __init__(self,
                  mul_kernel=(2.0, 2.0),
                  num_kernel=(5, 1),
                  weight_domain: float = 1.0,
-                 weight_scheduler: bool = False,
-                 lr_scheduler: bool = False,
+                 weight_scheduler: bool = True,
+                 lr_scheduler_gamma: float = 0.0,
+                 lr_scheduler_decay: float = 0.75,
                  warmup_epochs: int = 0,
                  devices: int = 1,
                  accelerator: str = "cpu",
                  metrics: List[str] = ["accuracy"]):
 
-        super(JANTrainer, self).__init__(extractor=extractor,
-                                         classifier=classifier,
-                                         num_classes=num_classes,
-                                         lr=lr,
-                                         weight_decay=weight_decay,
-                                         weight_domain=weight_domain,
-                                         weight_scheduler=weight_scheduler,
-                                         lr_scheduler=lr_scheduler,
-                                         warmup_epochs=warmup_epochs,
-                                         devices=devices,
-                                         accelerator=accelerator,
-                                         metrics=metrics)
+        super(JANTrainer,
+              self).__init__(extractor=extractor,
+                             classifier=classifier,
+                             num_classes=num_classes,
+                             lr=lr,
+                             weight_decay=weight_decay,
+                             weight_domain=weight_domain,
+                             weight_scheduler=weight_scheduler,
+                             lr_scheduler_gamma=lr_scheduler_gamma,
+                             lr_scheduler_decay=lr_scheduler_decay,
+                             warmup_epochs=warmup_epochs,
+                             devices=devices,
+                             accelerator=accelerator,
+                             metrics=metrics)
         self.mul_kernel = mul_kernel
         self.num_kernel = num_kernel
 
@@ -129,7 +132,7 @@ def training_step(self, batch: Tuple[torch.Tensor],
         domain_loss = self._domain_loss_fn(x_source_feat, x_target_feat,
                                            y_source_pred, y_target_pred)
 
-        task_loss = self._ce_fn(y_source_pred, y_source)
+        task_loss = self.ce_fn(y_source_pred, y_source)
 
         if self.current_epoch >= self.warmup_epochs:
             loss = task_loss + self.scheduled_weight_domain * domain_loss
diff --git a/torcheeg/trainers/domain_adaption/mmd_like.py b/torcheeg/trainers/domain_adaption/mmd_like.py
index 168efe1..f477c89 100644
--- a/torcheeg/trainers/domain_adaption/mmd_like.py
+++ b/torcheeg/trainers/domain_adaption/mmd_like.py
@@ -1,6 +1,6 @@
 from itertools import chain
 from typing import List, Tuple
-
+import logging
 import numpy as np
 import pytorch_lightning as pl
 import torch
@@ -42,8 +42,9 @@ def __init__(self,
                  lr: float = 1e-4,
                  weight_decay: float = 0.0,
                  weight_domain: float = 1.0,
-                 weight_scheduler: bool = False,
-                 lr_scheduler: bool = False,
+                 weight_scheduler: bool = True,
+                 lr_scheduler_gamma: float = 0.0,
+                 lr_scheduler_decay: float = 0.75,
                  warmup_epochs: int = 0,
                  devices: int = 1,
                  accelerator: str = "cpu",
@@ -57,7 +58,10 @@ def __init__(self,
         self.weight_decay = weight_decay
         self.weight_domain = weight_domain
         self.weight_scheduler = weight_scheduler
-        self.lr_scheduler = lr_scheduler
+
+        self.lr_scheduler_gamma = lr_scheduler_gamma
+        self.lr_scheduler_decay = lr_scheduler_decay
+        self.lr_scheduler = not lr_scheduler_gamma == 0.0
         self.warmup_epochs = warmup_epochs
 
         self.devices = devices
@@ -68,7 +72,7 @@ def __init__(self,
         self.metrics = metrics
         self.init_metrics(metrics, num_classes)
 
-        self._ce_fn = nn.CrossEntropyLoss()
+        self.ce_fn = nn.CrossEntropyLoss()
 
         self.num_batches = None  # init in 'fit' method
         self.non_warmup_epochs = None  # init in 'fit' method
@@ -122,7 +126,8 @@ def on_train_batch_start(self, batch, batch_idx):
             self.weight_factor = 2.0 / (1.0 + np.exp(-10 * p)) - 1
 
             if self.lr_scheduler:
-                self.lr_factor = 1.0 / ((1.0 + 10 * p)**0.75)
+                self.lr_factor = 1.0 / ((1.0 + self.lr_scheduler_gamma * p)**
+                                        self.lr_scheduler_decay)
 
         if self.weight_scheduler:
             self.scheduled_weight_domain = self.weight_domain * self.weight_factor
@@ -141,21 +146,21 @@ def training_step(self, batch: Tuple[torch.Tensor],
 
         domain_loss = self._domain_loss_fn(x_source_feat, x_target_feat)
 
-        task_loss = self._ce_fn(y_source_pred, y_source)
-        
+        task_loss = self.ce_fn(y_source_pred, y_source)
+
         if self.current_epoch >= self.warmup_epochs:
             loss = task_loss + self.scheduled_weight_domain * domain_loss
         else:
             loss = task_loss
 
         self.log("train_domain_loss",
-                 self.train_domain_loss(loss),
+                 self.train_domain_loss(domain_loss),
                  prog_bar=True,
                  on_epoch=False,
                  logger=False,
                  on_step=True)
         self.log("train_task_loss",
-                 self.train_task_loss(loss),
+                 self.train_task_loss(task_loss),
                  prog_bar=True,
                  on_epoch=False,
                  logger=False,
@@ -209,8 +214,12 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return x
 
     def configure_optimizers(self):
-        optimizer = torch.optim.Adam(chain(self.extractor.parameters(),
-                                           self.classifier.parameters()),
+        optimizer = torch.optim.Adam([{
+            "params": self.extractor.parameters()
+        }, {
+            "params": self.classifier.parameters(),
+            "lr": 10 * self.lr
+        }],
                                      lr=self.lr,
                                      weight_decay=self.weight_decay)
         if self.lr_scheduler:
diff --git a/torcheeg/trainers/finetune.py b/torcheeg/trainers/finetune.py
new file mode 100644
index 0000000..9f78b15
--- /dev/null
+++ b/torcheeg/trainers/finetune.py
@@ -0,0 +1,408 @@
+import logging
+from typing import Any, Dict, List, Tuple
+
+import numpy as np
+import pytorch_lightning as pl
+import torch
+import torch.nn as nn
+import torchmetrics
+from torch.optim.lr_scheduler import MultiStepLR
+from torch.utils.data import DataLoader
+from torchmetrics import MetricCollection
+from pytorch_lightning.callbacks.finetuning import BaseFinetuning
+from torch.optim import Optimizer
+
+_EVALUATE_OUTPUT = List[Dict[str, float]]  # 1 dict per DataLoader
+
+log = logging.getLogger(__name__)
+
+
+class FinetuningCallback(BaseFinetuning):
+    def __init__(self, freeze_epochs: int = 10, freeze_layers: List[str] = []):
+        super().__init__()
+        self.freeze_epochs = freeze_epochs
+        self.freeze_layers = freeze_layers
+
+    def freeze_before_training(self, pl_module: pl.LightningModule):
+        for layer_name in self.freeze_layers:
+            module = pl_module
+            for attr in layer_name.split('.'):
+                module = getattr(module, attr)
+            self.freeze(module)
+
+    def finetune_function(self, pl_module: pl.LightningModule, epoch: int,
+                          optimizer: Optimizer):
+        if epoch == self.freeze_epochs:
+            # unfreeze remaing layers
+            for layer_name in self.freeze_layers:
+                module = pl_module
+                for attr in layer_name.split('.'):
+                    module = getattr(module, attr)
+                self.unfreeze_and_add_param_group(module, optimizer=optimizer)
+
+
+def mixup_data(x, y, alpha=1.0):
+    '''Compute the mixup data. Return mixed inputs, pairs of targets, and lambda'''
+    if alpha > 0.0:
+        lam = np.random.beta(alpha, alpha)
+    else:
+        lam = 1.0
+
+    batch_size = x.size()[0]
+    index = torch.randperm(batch_size)
+
+    mixed_x = lam * x + (1 - lam) * x[index, :]
+    y_a, y_b = y, y[index]
+    return mixed_x, y_a, y_b, lam
+
+
+def mixup_criterion(criterion, pred, y_a, y_b, lam):
+    return lam * criterion(pred, y_a) + \
+           (1 - lam) * criterion(pred, y_b)
+
+
+from torch.optim.lr_scheduler import _LRScheduler
+
+
+class GradualWarmupScheduler(_LRScheduler):
+    def __init__(self, optimizer, multiplier, total_iter, after_scheduler=None):
+        self.multiplier = multiplier
+        if self.multiplier <= 1.:
+            raise ValueError('multiplier should be greater than 1.')
+        self.total_iter = total_iter
+        self.after_scheduler = after_scheduler
+        self.finished = False
+        super().__init__(optimizer)
+
+    def get_lr(self):
+        if self.last_epoch > self.total_iter:
+            if self.after_scheduler:
+                if not self.finished:
+                    self.after_scheduler.base_lrs = [
+                        base_lr * self.multiplier for base_lr in self.base_lrs
+                    ]
+                    self.finished = True
+                return self.after_scheduler.get_lr()
+            return [base_lr * self.multiplier for base_lr in self.base_lrs]
+
+        return [
+            base_lr *
+            ((self.multiplier - 1.) * self.last_epoch / self.total_iter + 1.)
+            for base_lr in self.base_lrs
+        ]
+
+    def step(self, epoch=None):
+        if self.finished and self.after_scheduler:
+            return self.after_scheduler.step(epoch)
+        else:
+            return super(GradualWarmupScheduler, self).step(epoch)
+
+
+def classification_metrics(metric_list: List[str], num_classes: int):
+    allowed_metrics = ['precision', 'recall', 'f1score', 'accuracy']
+
+    for metric in metric_list:
+        if metric not in allowed_metrics:
+            raise ValueError(
+                f"{metric} is not allowed. Please choose 'precision', 'recall', 'f1_score', 'accuracy'"
+            )
+    metric_dict = {
+        'accuracy':
+        torchmetrics.Accuracy(task='multiclass',
+                              num_classes=num_classes,
+                              top_k=1),
+        'precision':
+        torchmetrics.Precision(task='multiclass',
+                               average='macro',
+                               num_classes=num_classes),
+        'recall':
+        torchmetrics.Recall(task='multiclass',
+                            average='macro',
+                            num_classes=num_classes),
+        'f1score':
+        torchmetrics.F1Score(task='multiclass',
+                             average='macro',
+                             num_classes=num_classes)
+    }
+    metrics = [metric_dict[name] for name in metric_list]
+    return MetricCollection(metrics)
+
+
+class FinetuneTrainer(pl.LightningModule):
+    r'''
+        A generic trainer class for EEG classification.
+
+        .. code-block:: python
+
+            trainer = FinetuneTrainer(model)
+            trainer.fit(train_loader, val_loader)
+            trainer.test(test_loader)
+
+        Args:
+            model (nn.Module): The classification model, and the dimension of its output should be equal to the number of categories in the dataset. The output layer does not need to have a softmax activation function.
+            num_classes (int, optional): The number of categories in the dataset. If :obj:`None`, the number of categories will be inferred from the attribute :obj:`num_classes` of the model. (defualt: :obj:`None`)
+            lr (float): The learning rate. (default: :obj:`0.001`)
+            weight_decay (float): The weight decay. (default: :obj:`0.0`)
+            mixup_alpha (float): The alpha parameter of the mixup data augmentation. If set to 0.0, no mixup will be used. (default: :obj:`0.2`)
+            cosine_scheduler (bool): Whether to use cosine annealing scheduler. (default: :obj:`False`)
+            milestones (list of int): The milestones of the learning rate scheduler. (default: :obj:`[100]`)
+            freeze_epochs (int): The number of epochs to freeze the layers. (default: :obj:`10`)
+            freeze_layers (list of str): The names of the layers to freeze. (default: :obj:`[]`)
+            warmup_epochs (int): The number of epochs to warmup the learning rate. (default: :obj:`10`)
+            warmup_lr (float): The learning rate to warmup to. (default: :obj:`1e-5`)
+            devices (int): The number of devices to use. (default: :obj:`1`)
+            accelerator (str): The accelerator to use. Available options are: 'cpu', 'gpu'. (default: :obj:`"cpu"`)
+            metrics (list of str): The metrics to use. Available options are: 'precision', 'recall', 'f1score', 'accuracy'. (default: :obj:`["accuracy"]`)
+        
+        .. automethod:: fit
+        .. automethod:: test
+    '''
+    def __init__(self,
+                 model: nn.Module,
+                 num_classes: int,
+                 lr: float = 1e-3,
+                 weight_decay: float = 0.0,
+                 mixup_alpha: float = 0.0,
+                 cosine_scheduler: bool = False,
+                 milestones: List[int] = [100],
+                 freeze_epochs: int = 1,
+                 freeze_layers: List[str] = [],
+                 warmup_epochs: int = 10,
+                 warmup_lr: float = 1e-5,
+                 devices: int = 1,
+                 accelerator: str = "cpu",
+                 metrics: List[str] = ["accuracy"]):
+
+        super().__init__()
+        self.model = model
+
+        self.num_classes = num_classes
+        self.lr = lr
+        self.weight_decay = weight_decay
+
+        self.mixup_alpha = mixup_alpha
+        self.cosine_scheduler = cosine_scheduler
+        self.milestones = milestones
+        self.warmup_epochs = warmup_epochs
+        self.warmup_lr = warmup_lr
+        self.freeze_epochs = freeze_epochs
+        self.freeze_layers = freeze_layers
+
+        self.devices = devices
+        self.accelerator = accelerator
+        self.metrics = metrics
+
+        self.ce_fn = nn.CrossEntropyLoss()
+        self.len_train_loader = 0
+
+        self.init_metrics(metrics, num_classes)
+
+    def init_metrics(self, metrics: List[str], num_classes: int) -> None:
+        self.train_loss = torchmetrics.MeanMetric()
+        self.val_loss = torchmetrics.MeanMetric()
+        self.test_loss = torchmetrics.MeanMetric()
+
+        self.train_metrics = classification_metrics(metrics, num_classes)
+        self.val_metrics = classification_metrics(metrics, num_classes)
+        self.test_metrics = classification_metrics(metrics, num_classes)
+
+    def fit(self,
+            train_loader: DataLoader,
+            val_loader: DataLoader,
+            max_epochs: int = 300,
+            callbacks: List[pl.Callback] = [],
+            *args,
+            **kwargs) -> Any:
+        r'''
+        Args:
+            train_loader (DataLoader): Iterable DataLoader for traversing the training data batch (:obj:`torch.utils.data.dataloader.DataLoader`, :obj:`torch_geometric.loader.DataLoader`, etc).
+            val_loader (DataLoader): Iterable DataLoader for traversing the validation data batch (:obj:`torch.utils.data.dataloader.DataLoader`, :obj:`torch_geometric.loader.DataLoader`, etc).
+            max_epochs (int): Maximum number of epochs to train the model. (default: :obj:`300`)
+        '''
+        callbacks += [
+            FinetuningCallback(self.freeze_epochs, self.freeze_layers)
+        ]
+        self.len_train_loader = len(train_loader)
+        trainer = pl.Trainer(devices=self.devices,
+                             accelerator=self.accelerator,
+                             max_epochs=max_epochs,
+                             callbacks=callbacks,
+                             *args,
+                             **kwargs)
+        return trainer.fit(self, train_loader, val_loader)
+
+    def test(self, test_loader: DataLoader, *args,
+             **kwargs) -> _EVALUATE_OUTPUT:
+        r'''
+        Args:
+            test_loader (DataLoader): Iterable DataLoader for traversing the test data batch (torch.utils.data.dataloader.DataLoader, torch_geometric.loader.DataLoader, etc).
+        '''
+        trainer = pl.Trainer(devices=self.devices,
+                             accelerator=self.accelerator,
+                             *args,
+                             **kwargs)
+        return trainer.test(self, test_loader)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.model(x)
+
+    def training_step(self, batch: Tuple[torch.Tensor],
+                      batch_idx: int) -> torch.Tensor:
+        x, y = batch
+
+        y_hat = self(x)
+        if self.mixup_alpha > 0.0:
+            data, y_a, y_b, lam = mixup_data(x, y, self.mixup_alpha)
+            loss = mixup_criterion(self.ce_fn, self(data), y_a, y_b, lam)
+        else:
+            loss = self.ce_fn(y_hat, y)
+
+        # log to prog_bar
+        self.log("train_loss",
+                 self.train_loss(loss),
+                 prog_bar=True,
+                 on_epoch=False,
+                 logger=False,
+                 on_step=True)
+
+        for i, metric_value in enumerate(self.train_metrics.values()):
+            self.log(f"train_{self.metrics[i]}",
+                     metric_value(y_hat, y),
+                     prog_bar=True,
+                     on_epoch=False,
+                     logger=False,
+                     on_step=True)
+
+        return loss
+
+    def on_train_epoch_end(self) -> None:
+        self.log("train_loss",
+                 self.train_loss.compute(),
+                 prog_bar=False,
+                 on_epoch=True,
+                 on_step=False,
+                 logger=True)
+        for i, metric_value in enumerate(self.train_metrics.values()):
+            self.log(f"train_{self.metrics[i]}",
+                     metric_value.compute(),
+                     prog_bar=False,
+                     on_epoch=True,
+                     on_step=False,
+                     logger=True)
+
+        # print the metrics
+        str = "\n[Train] "
+        for key, value in self.trainer.logged_metrics.items():
+            if key.startswith("train_"):
+                str += f"{key}: {value:.3f} "
+        print(str + '\n')
+
+        # reset the metrics
+        self.train_loss.reset()
+        self.train_metrics.reset()
+
+    def validation_step(self, batch: Tuple[torch.Tensor],
+                        batch_idx: int) -> torch.Tensor:
+        x, y = batch
+        y_hat = self(x)
+        loss = self.ce_fn(y_hat, y)
+
+        self.val_loss.update(loss)
+        self.val_metrics.update(y_hat, y)
+        return loss
+
+    def on_validation_epoch_end(self) -> None:
+        self.log("val_loss",
+                 self.val_loss.compute(),
+                 prog_bar=False,
+                 on_epoch=True,
+                 on_step=False,
+                 logger=True)
+        for i, metric_value in enumerate(self.val_metrics.values()):
+            self.log(f"val_{self.metrics[i]}",
+                     metric_value.compute(),
+                     prog_bar=False,
+                     on_epoch=True,
+                     on_step=False,
+                     logger=True)
+
+        # print the metrics
+        str = "\n[Val] "
+        for key, value in self.trainer.logged_metrics.items():
+            if key.startswith("val_"):
+                str += f"{key}: {value:.3f} "
+        print(str + '\n')
+
+        self.val_loss.reset()
+        self.val_metrics.reset()
+
+    def test_step(self, batch: Tuple[torch.Tensor],
+                  batch_idx: int) -> torch.Tensor:
+        x, y = batch
+        y_hat = self(x)
+        loss = self.ce_fn(y_hat, y)
+
+        self.test_loss.update(loss)
+        self.test_metrics.update(y_hat, y)
+        return loss
+
+    def on_test_epoch_end(self) -> None:
+        self.log("test_loss",
+                 self.test_loss.compute(),
+                 prog_bar=False,
+                 on_epoch=True,
+                 on_step=False,
+                 logger=True)
+        for i, metric_value in enumerate(self.test_metrics.values()):
+            self.log(f"test_{self.metrics[i]}",
+                     metric_value.compute(),
+                     prog_bar=False,
+                     on_epoch=True,
+                     on_step=False,
+                     logger=True)
+
+        # print the metrics
+        str = "\n[Test] "
+        for key, value in self.trainer.logged_metrics.items():
+            if key.startswith("test_"):
+                str += f"{key}: {value:.3f} "
+        print(str + '\n')
+
+        self.test_loss.reset()
+        self.test_metrics.reset()
+
+    def configure_optimizers(self):
+        parameters = list(self.model.parameters())
+        trainable_parameters = list(
+            filter(lambda p: p.requires_grad, parameters))
+        optimizer = torch.optim.SGD(trainable_parameters,
+                                    lr=self.lr,
+                                    weight_decay=self.weight_decay,
+                                    momentum=0.9)
+
+        if self.cosine_scheduler:
+            scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
+                optimizer, T_max=self.trainer.max_epochs)
+        else:
+            milestones = [
+                (e - self.warmup_epochs) * self.len_train_loader
+                for e in self.milestones
+            ]
+            scheduler = MultiStepLR(optimizer, milestones=milestones, gamma=0.1)
+        if self.warmup_epochs > 0:
+            total_iter = self.trainer.max_epochs * self.len_train_loader
+            multiplier = self.lr / self.warmup_lr
+            scheduler = GradualWarmupScheduler(optimizer,
+                                               multiplier=multiplier,
+                                               total_iter=total_iter,
+                                               after_scheduler=scheduler)
+
+        return [optimizer], [scheduler]
+
+    def predict_step(self,
+                     batch: Tuple[torch.Tensor],
+                     batch_idx: int,
+                     dataloader_idx: int = 0):
+        x, y = batch
+        y_hat = self(x)
+        return y_hat
\ No newline at end of file
diff --git a/torcheeg/trainers/generative/beta_vae.py b/torcheeg/trainers/generative/beta_vae.py
index 2f5c4e5..8b48f7f 100644
--- a/torcheeg/trainers/generative/beta_vae.py
+++ b/torcheeg/trainers/generative/beta_vae.py
@@ -1,3 +1,4 @@
+import logging
 import warnings
 from itertools import chain
 from typing import Any, Dict, List, Tuple
@@ -13,6 +14,7 @@
 
 _EVALUATE_OUTPUT = List[Dict[str, float]]  # 1 dict per DataLoader
 
+log = logging.getLogger(__name__)
 
 class BetaVAETrainer(pl.LightningModule):
     r'''
diff --git a/torcheeg/trainers/generative/ddpm.py b/torcheeg/trainers/generative/ddpm.py
index 0f6cf5a..a8de921 100644
--- a/torcheeg/trainers/generative/ddpm.py
+++ b/torcheeg/trainers/generative/ddpm.py
@@ -1,3 +1,4 @@
+import logging
 import warnings
 from contextlib import contextmanager
 from functools import partial
@@ -16,6 +17,8 @@
 
 _EVALUATE_OUTPUT = List[Dict[str, float]]  # 1 dict per DataLoader
 
+log = logging.getLogger(__name__)
+
 
 def extract_into_tensor(a: torch.Tensor, t: int,
                         x_shape: Tuple[int, ...]) -> torch.Tensor:
diff --git a/torcheeg/trainers/generative/glow.py b/torcheeg/trainers/generative/glow.py
index 467e4bf..0303716 100644
--- a/torcheeg/trainers/generative/glow.py
+++ b/torcheeg/trainers/generative/glow.py
@@ -1,3 +1,4 @@
+import logging
 import warnings
 from typing import Any, Dict, List, Tuple
 
@@ -12,6 +13,8 @@
 
 _EVALUATE_OUTPUT = List[Dict[str, float]]  # 1 dict per DataLoader
 
+log = logging.getLogger(__name__)
+
 
 class GlowTrainer(pl.LightningModule):
     r'''
diff --git a/torcheeg/trainers/generative/wgan_gp.py b/torcheeg/trainers/generative/wgan_gp.py
index f2293a7..9008163 100644
--- a/torcheeg/trainers/generative/wgan_gp.py
+++ b/torcheeg/trainers/generative/wgan_gp.py
@@ -1,3 +1,4 @@
+import logging
 import warnings
 from typing import Any, Dict, List, Tuple
 
@@ -13,6 +14,8 @@
 
 _EVALUATE_OUTPUT = List[Dict[str, float]]  # 1 dict per DataLoader
 
+log = logging.getLogger(__name__)
+
 
 def gradient_penalty(model, real, fake, *args, **kwargs):
     device = real.device
diff --git a/torcheeg/trainers/self_supervised/__init__.py b/torcheeg/trainers/self_supervised/__init__.py
index 8d23f02..0d27b50 100644
--- a/torcheeg/trainers/self_supervised/__init__.py
+++ b/torcheeg/trainers/self_supervised/__init__.py
@@ -1,1 +1,2 @@
-from .sim_clr import SimCLRTrainer
\ No newline at end of file
+from .sim_clr import SimCLRTrainer
+from .byol import BYOLTrainer
\ No newline at end of file
diff --git a/torcheeg/trainers/self_supervised/byol.py b/torcheeg/trainers/self_supervised/byol.py
new file mode 100644
index 0000000..13ec8f6
--- /dev/null
+++ b/torcheeg/trainers/self_supervised/byol.py
@@ -0,0 +1,474 @@
+import copy
+import logging
+from itertools import chain
+from typing import Any, List, Tuple
+
+import pytorch_lightning as pl
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchmetrics
+from torch.utils.data import DataLoader
+
+log = logging.getLogger(__name__)
+
+
+class BYOLTrainer(pl.LightningModule):
+    r'''
+    This class supports the implementation of Bootstrap Your Own Latent (BYOL) for self-supervised pre-training.
+
+    - Paper: Grill J B, Strub F, Altché F, et al. Bootstrap your own latent-a new approach to self-supervised learning[J]. Advances in neural information processing systems, 2020, 33: 21271-21284.
+    - URL: https://proceedings.neurips.cc/paper/2020/hash/f3ada80d5c4ee70142b17b8192b2958e-Abstract.html
+    - Related Project: https://github.com/lucidrains/byol-pytorch
+
+    .. code-block:: python
+
+        trainer = BYOLTrainer(extractor,
+                              extract_channels=256,
+                              devices=1,
+                              accelerator='gpu')
+        trainer.fit(train_loader, val_loader)
+
+    NOTE: The first element of each batch in :obj:`train_loader` and :obj:`val_loader` should be a two-tuple, representing two random transformations (views) of data. You can use :obj:`Contrastive` to achieve this functionality.
+
+    .. code-block:: python
+
+        contras_dataset = DEAPDataset(
+            io_path=f'./io/deap',
+            root_path='./data_preprocessed_python',
+            offline_transform=transforms.Compose([
+                transforms.BandDifferentialEntropy(sampling_rate=128,
+                                                apply_to_baseline=True),
+                transforms.BaselineRemoval(),
+                transforms.ToGrid(DEAP_CHANNEL_LOCATION_DICT)
+            ]),
+            online_transform=transforms.Compose([
+                transforms.ToTensor(),
+                transforms.Contrastive(transforms.Compose( # see here
+                    [transforms.RandomMask(p=0.5),
+                    transforms.RandomNoise(p=0.5)]),
+                                    num_views=2)
+            ]),
+            chunk_size=128,
+            baseline_chunk_size=128,
+            num_baseline=3)
+
+        trainer = BYOLTrainer(extractor,
+                              extract_channels=256,
+                              devices=1,
+                              accelerator='gpu')
+        trainer.fit(train_loader, val_loader)
+
+    Args:
+        extractor (nn.Module): The feature extraction model learns the feature representation of the EEG signal by forcing the correlation matrixes of source and target data to be close.
+        extract_channels (int): The feature dimensions of the output of the feature extraction model.
+        proj_channels (int): The feature dimensions of the output of the projection head. (default: :obj:`256`)
+        proj_hid_channels (int): The feature dimensions of the hidden layer of the projection head. (default: :obj:`512`)
+        lr (float): The learning rate. (default: :obj:`0.0001`)
+        weight_decay (float): The weight decay. (default: :obj:`0.0`)
+        devices (int): The number of GPUs to use. (default: :obj:`1`)
+        accelerator (str): The accelerator to use. Available options are: 'cpu', 'gpu'. (default: :obj:`"cpu"`)
+        metrics (List[str]): The metrics to use. Available options are: 'acc_top1', 'acc_top5', 'acc_mean_pos'. (default: :obj:`["acc_top1"]`)
+
+    .. automethod:: fit
+    '''
+    def __init__(self,
+                 extractor: nn.Module,
+                 extract_channels: int,
+                 proj_channels: int = 256,
+                 proj_hid_channels: int = 512,
+                 lr: float = 1e-4,
+                 weight_decay: float = 0.0,
+                 moving_average_decay=0.99,
+                 devices: int = 1,
+                 accelerator: str = "cpu",
+                 metrics: List[str] = ["acc_top1"]):
+        super().__init__()
+
+        self.student_model = extractor
+        self.student_projector = self.MLP(extract_channels, proj_hid_channels,
+                                          proj_channels)
+        self.student_predictor = self.MLP(proj_channels, proj_hid_channels,
+                                          proj_channels)
+
+        self.teacher_model, self.teacher_projector = self.teacher()
+
+        self.extract_channels = extract_channels
+        self.proj_channels = proj_channels
+        self.proj_hid_channels = proj_hid_channels
+
+        self.lr = lr
+        self.weight_decay = weight_decay
+
+        self.moving_average_decay = moving_average_decay
+
+        self.devices = devices
+        self.accelerator = accelerator
+        self.metrics = metrics
+
+        self.init_metrics(metrics)
+
+    def MLP(self, in_channels: int, hid_channels: int, out_channels: int):
+        return nn.Sequential(
+            nn.Linear(in_channels, hid_channels),
+            nn.BatchNorm1d(hid_channels),
+            nn.ReLU(inplace=True),
+            nn.Linear(hid_channels, out_channels),
+        )
+
+    def teacher(self) -> Tuple[nn.Module, nn.Module]:
+        r'''
+        The teacher model is a copy of the student model, but the weights are not updated during training.
+
+        Returns:
+            tuple: The teacher model and the projection head.
+        '''
+        with torch.no_grad():
+            teacher_model = copy.deepcopy(self.student_model)
+            teacher_projector = copy.deepcopy(self.student_projector)
+            return teacher_model, teacher_projector
+
+    def update_moving_average(self) -> None:
+        r'''
+        Update the weights of the teacher model and the projection head.
+        '''
+        with torch.no_grad():
+            for param_q, param_k in zip(self.student_model.parameters(),
+                                        self.teacher_model.parameters()):
+                param_k.data = param_k.data * self.moving_average_decay + param_q.data * (
+                    1. - self.moving_average_decay)
+            for param_q, param_k in zip(self.student_projector.parameters(),
+                                        self.teacher_projector.parameters()):
+                param_k.data = param_k.data * self.moving_average_decay + param_q.data * (
+                    1. - self.moving_average_decay)
+
+    def loss_fn(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+        r'''
+        The loss function of BYOL.
+
+        Args:
+            x (torch.Tensor): The output of the projection head.
+            y (torch.Tensor): The output of the projection head of the teacher model.
+
+        Returns:
+            torch.Tensor: The loss.
+        '''
+        x = F.normalize(x, dim=-1, p=2)
+        y = F.normalize(y, dim=-1, p=2)
+        return 2 - 2 * (x * y).sum(dim=-1)
+
+    def init_metrics(self, metrics) -> None:
+        self.train_loss = torchmetrics.MeanMetric()
+        self.val_loss = torchmetrics.MeanMetric()
+
+        if "acc_top1" in metrics:
+            self.train_acc_top1 = torchmetrics.MeanMetric()
+            self.val_acc_top1 = torchmetrics.MeanMetric()
+        if "acc_top5" in metrics:
+            self.train_acc_top5 = torchmetrics.MeanMetric()
+            self.val_acc_top5 = torchmetrics.MeanMetric()
+        if "acc_mean_pos" in metrics:
+            self.train_acc_mean_pos = torchmetrics.MeanMetric()
+            self.val_acc_mean_pos = torchmetrics.MeanMetric()
+
+    def fit(self,
+            train_loader: DataLoader,
+            val_loader: DataLoader,
+            max_epochs: int = 300,
+            *args,
+            **kwargs) -> Any:
+        r'''
+        NOTE: The first element of each batch in :obj:`train_loader` and :obj:`val_loader` should be a two-tuple, representing two random transformations (views) of data. You can use :obj:`Contrastive` to achieve this functionality.
+
+        Args:
+            train_loader (DataLoader): Iterable DataLoader for traversing the training data batch (:obj:`torch.utils.data.dataloader.DataLoader`, :obj:`torch_geometric.loader.DataLoader`, etc).
+            val_loader (DataLoader): Iterable DataLoader for traversing the validation data batch (:obj:`torch.utils.data.dataloader.DataLoader`, :obj:`torch_geometric.loader.DataLoader`, etc).
+            max_epochs (int): Maximum number of epochs to train the model. (default: :obj:`300`)
+        '''
+        trainer = pl.Trainer(devices=self.devices,
+                             accelerator=self.accelerator,
+                             max_epochs=max_epochs,
+                             *args,
+                             **kwargs)
+        return trainer.fit(self, train_loader, val_loader)
+
+    def training_step(self, batch: Tuple[torch.Tensor],
+                      batch_idx: int) -> torch.Tensor:
+        xs, _ = batch
+
+        assert len(xs) == 2, "The number of views must be two in BYOL."
+
+        for i, x in enumerate(xs):
+            # # batch size must greater than one
+            # assert x.shape[
+            #     0] > 1, "Batch size must greater than one, due to the batch normalization layer in the projection head (BYOL)."
+            # copy the first element of the batch make the batch size greater than one
+            if x.shape[0] == 1:
+                xs[i] = torch.cat([x, x[0].unsqueeze(0)], dim=0)
+
+        # Student model
+        eeg_one, eeg_two = xs
+
+        student_proj_one = self.student_model(eeg_one)
+        student_proj_one = self.student_projector(student_proj_one)
+
+        student_proj_two = self.student_model(eeg_two)
+        student_proj_two = self.student_projector(student_proj_two)
+
+        student_pred_one = self.student_predictor(student_proj_one)
+        student_pred_two = self.student_predictor(student_proj_two)
+
+        with torch.no_grad():
+            # Teacher model
+            teacher_proj_one = self.teacher_model(eeg_one)
+            teacher_proj_one = self.teacher_projector(teacher_proj_one)
+
+            teacher_proj_two = self.teacher_model(eeg_two)
+            teacher_proj_two = self.teacher_projector(teacher_proj_two)
+
+        loss_one = self.loss_fn(student_pred_one, teacher_proj_two)
+        loss_two = self.loss_fn(student_pred_two, teacher_proj_one)
+        loss = (loss_one + loss_two).mean()
+
+        # Get ranking position of positive example
+        xs = torch.cat(xs, dim=0)
+        feats = self.student_model(xs)
+        cos_sim = F.cosine_similarity(feats[:, None, :],
+                                      feats[None, :, :],
+                                      dim=-1)
+        # Mask out cosine similarity to itself
+        self_mask = torch.eye(cos_sim.shape[0],
+                              dtype=torch.bool,
+                              device=cos_sim.device)
+        cos_sim.masked_fill_(self_mask, -9e15)
+        # Find positive example -> batch_size//2 away from the original example
+        pos_mask = self_mask.roll(shifts=cos_sim.shape[0] // 2, dims=0)
+        comb_sim = torch.cat(
+            [cos_sim[pos_mask][:, None],
+             cos_sim.masked_fill(pos_mask, -9e15)
+             ],  # First position positive example
+            dim=-1,
+        )
+        sim_argsort = comb_sim.argsort(dim=-1, descending=True).argmin(dim=-1)
+
+        self.log("train_loss",
+                 self.train_loss(loss),
+                 prog_bar=True,
+                 on_epoch=False,
+                 logger=False,
+                 on_step=True)
+        if "acc_top1" in self.metrics:
+            # Logging ranking metrics
+            self.log("train_acc_top1",
+                     self.train_acc_top1((sim_argsort == 0).float()),
+                     prog_bar=True,
+                     on_epoch=False,
+                     logger=False,
+                     on_step=True)
+        if "acc_top5" in self.metrics:
+            self.log("train_acc_top5",
+                     self.train_acc_top5((sim_argsort < 5).float()),
+                     prog_bar=True,
+                     on_epoch=False,
+                     logger=False,
+                     on_step=True)
+        if "acc_mean_pos" in self.metrics:
+            self.log("train_acc_mean_pos",
+                     self.train_acc_mean_pos(1 + sim_argsort.float()),
+                     prog_bar=True,
+                     on_epoch=False,
+                     logger=False,
+                     on_step=True)
+
+        return loss
+
+    def on_train_epoch_end(self) -> None:
+        self.log("train_loss",
+                 self.train_loss.compute(),
+                 prog_bar=False,
+                 on_epoch=True,
+                 on_step=False,
+                 logger=True)
+        if "acc_top1" in self.metrics:
+            self.log("train_acc_top1",
+                     self.train_acc_top1.compute(),
+                     prog_bar=False,
+                     on_epoch=True,
+                     on_step=False,
+                     logger=True)
+        if "acc_top5" in self.metrics:
+            self.log("train_acc_top5",
+                     self.train_acc_top5.compute(),
+                     prog_bar=False,
+                     on_epoch=True,
+                     on_step=False,
+                     logger=True)
+        if "acc_mean_pos" in self.metrics:
+            self.log("train_acc_mean_pos",
+                     self.train_acc_mean_pos.compute(),
+                     prog_bar=False,
+                     on_epoch=True,
+                     on_step=False,
+                     logger=True)
+
+        # print the metrics
+        str = "\n[Train] "
+        for key, value in self.trainer.logged_metrics.items():
+            if key.startswith("train_"):
+                str += f"{key}: {value:.3f} "
+        print(str + '\n')
+
+        # reset the metrics
+        self.train_loss.reset()
+        if "acc_top1" in self.metrics:
+            self.train_acc_top1.reset()
+        if "acc_top5" in self.metrics:
+            self.train_acc_top5.reset()
+        if "acc_mean_pos" in self.metrics:
+            self.train_acc_mean_pos.reset()
+
+    def on_after_backward(self) -> None:
+        self.update_moving_average()
+        return super().on_after_backward()
+
+    def validation_step(self, batch: Tuple[torch.Tensor],
+                        batch_idx: int) -> torch.Tensor:
+        xs, _ = batch
+
+        assert len(xs) == 2, "The number of views must be two in BYOL."
+        for i, x in enumerate(xs):
+            # # batch size must greater than one
+            # assert x.shape[
+            #     0] > 1, "Batch size must greater than one, due to the batch normalization layer in the projection head (BYOL)."
+            # copy the first element of the batch make the batch size greater than one
+            if x.shape[0] == 1:
+                xs[i] = torch.cat([x, x[0].unsqueeze(0)], dim=0)
+
+        # Student model
+        eeg_one, eeg_two = xs
+
+        student_proj_one = self.student_model(eeg_one)
+        student_proj_one = self.student_projector(student_proj_one)
+
+        student_proj_two = self.student_model(eeg_two)
+        student_proj_two = self.student_projector(student_proj_two)
+
+        student_pred_one = self.student_predictor(student_proj_one)
+        student_pred_two = self.student_predictor(student_proj_two)
+
+        with torch.no_grad():
+            # Teacher model
+            teacher_proj_one = self.teacher_model(eeg_one)
+            teacher_proj_one = self.teacher_projector(teacher_proj_one)
+
+            teacher_proj_two = self.teacher_model(eeg_two)
+            teacher_proj_two = self.teacher_projector(teacher_proj_two)
+
+        loss_one = self.loss_fn(student_pred_one, teacher_proj_two)
+        loss_two = self.loss_fn(student_pred_two, teacher_proj_one)
+        loss = (loss_one + loss_two).mean()
+
+        # Get ranking position of positive example
+        xs = torch.cat(xs, dim=0)
+        feats = self.student_model(xs)
+        cos_sim = F.cosine_similarity(feats[:, None, :],
+                                      feats[None, :, :],
+                                      dim=-1)
+        # Mask out cosine similarity to itself
+        self_mask = torch.eye(cos_sim.shape[0],
+                              dtype=torch.bool,
+                              device=cos_sim.device)
+        cos_sim.masked_fill_(self_mask, -9e15)
+        # Find positive example -> batch_size//2 away from the original example
+        pos_mask = self_mask.roll(shifts=cos_sim.shape[0] // 2, dims=0)
+        comb_sim = torch.cat(
+            [cos_sim[pos_mask][:, None],
+             cos_sim.masked_fill(pos_mask, -9e15)
+             ],  # First position positive example
+            dim=-1,
+        )
+        sim_argsort = comb_sim.argsort(dim=-1, descending=True).argmin(dim=-1)
+
+        self.log("val_loss",
+                 self.val_loss(loss),
+                 prog_bar=True,
+                 on_epoch=False,
+                 logger=False,
+                 on_step=True)
+        if "acc_top1" in self.metrics:
+            # Logging ranking metrics
+            self.log("val_acc_top1",
+                     self.val_acc_top1((sim_argsort == 0).float()),
+                     prog_bar=True,
+                     on_epoch=False,
+                     logger=False,
+                     on_step=True)
+        if "acc_top5" in self.metrics:
+            self.log("val_acc_top5",
+                     self.val_acc_top5((sim_argsort < 5).float()),
+                     prog_bar=True,
+                     on_epoch=False,
+                     logger=False,
+                     on_step=True)
+        if "acc_mean_pos" in self.metrics:
+            self.log("val_acc_mean_pos",
+                     self.val_acc_mean_pos(1 + sim_argsort.float()),
+                     prog_bar=True,
+                     on_epoch=False,
+                     logger=False,
+                     on_step=True)
+
+        return loss
+
+    def on_validation_epoch_end(self) -> None:
+        self.log("val_loss",
+                 self.val_loss.compute(),
+                 prog_bar=False,
+                 on_epoch=True,
+                 on_step=False,
+                 logger=True)
+        if "acc_top1" in self.metrics:
+            self.log("val_acc_top1",
+                     self.val_acc_top1.compute(),
+                     prog_bar=False,
+                     on_epoch=True,
+                     on_step=False,
+                     logger=True)
+        if "acc_top5" in self.metrics:
+            self.log("val_acc_top5",
+                     self.val_acc_top5.compute(),
+                     prog_bar=False,
+                     on_epoch=True,
+                     on_step=False,
+                     logger=True)
+        if "acc_mean_pos" in self.metrics:
+            self.log("val_acc_mean_pos",
+                     self.val_acc_mean_pos.compute(),
+                     prog_bar=False,
+                     on_epoch=True,
+                     on_step=False,
+                     logger=True)
+
+        # print the metrics
+        str = "\n[VAL] "
+        for key, value in self.trainer.logged_metrics.items():
+            if key.startswith("val_"):
+                str += f"{key}: {value:.3f} "
+        print(str + '\n')
+
+        # reset the metrics
+        self.val_loss.reset()
+        if "acc_top1" in self.metrics:
+            self.val_acc_top1.reset()
+        if "acc_top5" in self.metrics:
+            self.val_acc_top5.reset()
+        if "acc_mean_pos" in self.metrics:
+            self.val_acc_mean_pos.reset()
+
+    def configure_optimizers(self):
+        optimizer = torch.optim.Adam(chain(self.student_model.parameters(),
+                                           self.student_projector.parameters(),
+                                           self.student_predictor.parameters()),
+                                     lr=self.lr,
+                                     weight_decay=self.weight_decay)
+        return optimizer
\ No newline at end of file
diff --git a/torcheeg/trainers/self_supervised/sim_clr.py b/torcheeg/trainers/self_supervised/sim_clr.py
index 171d00f..05ba362 100644
--- a/torcheeg/trainers/self_supervised/sim_clr.py
+++ b/torcheeg/trainers/self_supervised/sim_clr.py
@@ -1,12 +1,15 @@
-import torch
+import logging
+from itertools import chain
+from typing import Any, List, Tuple
+
 import pytorch_lightning as pl
+import torch
 import torch.nn as nn
-
-import torchmetrics
 import torch.nn.functional as F
+import torchmetrics
 from torch.utils.data import DataLoader
 
-from typing import Any, Tuple, List
+log = logging.getLogger(__name__)
 
 
 class SimCLRTrainer(pl.LightningModule):
@@ -18,15 +21,46 @@ class SimCLRTrainer(pl.LightningModule):
     - Related Project: https://github.com/sthalles/SimCLR
 
     .. code-block:: python
-    
+
+        trainer = SimCLRTrainer(extractor,
+                                devices=1,
+                                accelerator='gpu')
+        trainer.fit(train_loader, val_loader)
+
+    NOTE: The first element of each batch in :obj:`train_loader` and :obj:`val_loader` should be a two-tuple, representing two random transformations (views) of data. You can use :obj:`Contrastive` to achieve this functionality.
+
+    .. code-block:: python
+
+        contras_dataset = DEAPDataset(
+            io_path=f'./io/deap',
+            root_path='./data_preprocessed_python',
+            offline_transform=transforms.Compose([
+                transforms.BandDifferentialEntropy(sampling_rate=128,
+                                                apply_to_baseline=True),
+                transforms.BaselineRemoval(),
+                transforms.ToGrid(DEAP_CHANNEL_LOCATION_DICT)
+            ]),
+            online_transform=transforms.Compose([
+                transforms.ToTensor(),
+                transforms.Contrastive(transforms.Compose( # see here
+                    [transforms.RandomMask(p=0.5),
+                    transforms.RandomNoise(p=0.5)]),
+                                    num_views=2)
+            ]),
+            chunk_size=128,
+            baseline_chunk_size=128,
+            num_baseline=3)
+
         trainer = SimCLRTrainer(extractor,
                                 devices=1,
                                 accelerator='gpu')
         trainer.fit(train_loader, val_loader)
-        trainer.test(test_loader)
 
     Args:
         extractor (nn.Module): The feature extraction model learns the feature representation of the EEG signal by forcing the correlation matrixes of source and target data to be close.
+        extract_channels (int): The feature dimensions of the output of the feature extraction model.
+        proj_channels (int): The feature dimensions of the output of the projection head. (default: :obj:`256`)
+        proj_hid_channels (int): The feature dimensions of the hidden layer of the projection head. (default: :obj:`512`)
         lr (float): The learning rate. (default: :obj:`0.0001`)
         weight_decay (float): The weight decay. (default: :obj:`0.0`)
         temperature (float): The temperature. (default: :obj:`0.1`)
@@ -38,6 +72,9 @@ class SimCLRTrainer(pl.LightningModule):
     '''
     def __init__(self,
                  extractor: nn.Module,
+                 extract_channels: int,
+                 proj_channels: int = 256,
+                 proj_hid_channels: int = 512,
                  lr: float = 1e-4,
                  weight_decay: float = 0.0,
                  temperature: float = 0.1,
@@ -47,6 +84,8 @@ def __init__(self,
         super().__init__()
 
         self.extractor = extractor
+        self.projector = self.MLP(extract_channels, proj_hid_channels,
+                                  proj_channels)
         self.lr = lr
         self.weight_decay = weight_decay
         self.temperature = temperature
@@ -56,6 +95,14 @@ def __init__(self,
 
         self.init_metrics(metrics)
 
+    def MLP(self, in_channels: int, hid_channels: int, out_channels: int):
+        return nn.Sequential(
+            nn.Linear(in_channels, hid_channels),
+            nn.BatchNorm1d(hid_channels),
+            nn.ReLU(inplace=True),
+            nn.Linear(hid_channels, out_channels),
+        )
+
     def init_metrics(self, metrics) -> None:
         self.train_loss = torchmetrics.MeanMetric()
         self.val_loss = torchmetrics.MeanMetric()
@@ -97,6 +144,7 @@ def training_step(self, batch: Tuple[torch.Tensor],
         xs = torch.cat(xs, dim=0)
 
         feats = self.extractor(xs)
+        feats = self.projector(feats)
         cos_sim = F.cosine_similarity(feats[:, None, :],
                                       feats[None, :, :],
                                       dim=-1)
@@ -120,7 +168,6 @@ def training_step(self, batch: Tuple[torch.Tensor],
             dim=-1,
         )
         sim_argsort = comb_sim.argsort(dim=-1, descending=True).argmin(dim=-1)
-
         self.log("train_loss",
                  self.train_loss(nll),
                  prog_bar=True,
@@ -203,6 +250,7 @@ def validation_step(self, batch: Tuple[torch.Tensor],
         xs = torch.cat(xs, dim=0)
 
         feats = self.extractor(xs)
+        feats = self.projector(feats)
         cos_sim = F.cosine_similarity(feats[:, None, :],
                                       feats[None, :, :],
                                       dim=-1)
@@ -288,7 +336,7 @@ def on_validation_epoch_end(self) -> None:
                      logger=True)
 
         # print the metrics
-        str = "\n[Train] "
+        str = "\n[VAL] "
         for key, value in self.trainer.logged_metrics.items():
             if key.startswith("val_"):
                 str += f"{key}: {value:.3f} "
@@ -304,7 +352,9 @@ def on_validation_epoch_end(self) -> None:
             self.val_acc_mean_pos.reset()
 
     def configure_optimizers(self):
-        optimizer = torch.optim.Adam(self.extractor.parameters(),
-                                     lr=self.lr,
-                                     weight_decay=self.weight_decay)
+        optimizer = torch.optim.Adam(
+            chain(self.extractor.parameters(), self.projector.parameters()),
+            lr=self.lr,
+            weight_decay=self.weight_decay,
+        )
         return optimizer
\ No newline at end of file
diff --git a/torcheeg/transforms/label/__init__.py b/torcheeg/transforms/label/__init__.py
index c8a59e3..624db9b 100644
--- a/torcheeg/transforms/label/__init__.py
+++ b/torcheeg/transforms/label/__init__.py
@@ -1,4 +1,5 @@
 from .binary import *
 from .select import *
 from .string import *
-from .fix import *
\ No newline at end of file
+from .fix import *
+from .mapping import *
\ No newline at end of file
diff --git a/torcheeg/transforms/label/mapping.py b/torcheeg/transforms/label/mapping.py
new file mode 100644
index 0000000..c533c01
--- /dev/null
+++ b/torcheeg/transforms/label/mapping.py
@@ -0,0 +1,62 @@
+from typing import Dict, List, Union
+
+from ..base_transform import LabelTransform
+
+
+class Mapping(LabelTransform):
+    r'''
+    Mapping the label according to a certain dictionary.
+    
+    .. code-block:: python
+
+        transform = Mapping({
+            'left_hand': 0,
+            'right_hand': 1,
+        })
+        transform(y='left_hand')['y']
+        >>> 0
+
+    :obj:`Mapping` allows simultaneous binarization using the same threshold for multiple labels.
+
+    .. code-block:: python
+
+        transform = Mapping({
+            'left_hand': 0,
+            'right_hand': 1,
+            'left_feet': 0,
+            'right_feet': 1
+        })
+        transform(y=['left_hand', 'left_feet'])['y']
+        >>> [0, 0]
+
+    Args:
+        map_dict (float): The mapping dictionary.
+        default (float, optional): The default value when the input label is not in the dictionary. (default: :obj:`-1`)
+
+    .. automethod:: __call__
+    '''
+    def __init__(self, map_dict: float, default: float = -1) -> None:
+        super(Mapping, self).__init__()
+        self.map_dict = map_dict
+        self.default = default
+
+    def __call__(self, *args, y: Union[int, float, List],
+                 **kwargs) -> Union[int, List]:
+        r'''
+        Args:
+            label (int, float, or list): The input label or list of labels.
+            
+        Returns:
+            int, float, or list: The output label or list of labels after binarization.
+        '''
+        return super().__call__(*args, y=y, **kwargs)
+
+    def apply(self, y: Union[int, float, List], **kwargs) -> Union[int, List]:
+        if isinstance(y, list):
+            return [self.map_dict.get(label, self.default) for label in y]
+        else:
+            return self.map_dict.get(y, self.default)
+
+    @property
+    def repr_body(self) -> Dict:
+        return dict(super().repr_body, **{'map_dict': self.map_dict})
\ No newline at end of file
diff --git a/torcheeg/transforms/numpy/__init__.py b/torcheeg/transforms/numpy/__init__.py
index 779b41c..02e5ec1 100644
--- a/torcheeg/transforms/numpy/__init__.py
+++ b/torcheeg/transforms/numpy/__init__.py
@@ -8,3 +8,5 @@
 from .correlation import *
 from .band_pyeeg import *
 from .downsample import *
+from .rearrange import *
+from .flatten import *
\ No newline at end of file
diff --git a/torcheeg/transforms/numpy/flatten.py b/torcheeg/transforms/numpy/flatten.py
new file mode 100644
index 0000000..d9fa34b
--- /dev/null
+++ b/torcheeg/transforms/numpy/flatten.py
@@ -0,0 +1,36 @@
+from typing import Dict, Union
+
+import numpy as np
+
+from ..base_transform import EEGTransform
+
+
+class Flatten(EEGTransform):
+    r'''
+    Flatten the input EEG representation.
+
+    .. code-block:: python
+
+        transform = Flatten()
+        transform(eeg=np.random.randn(62, 5))['eeg'].shape
+        >>> (310,)
+
+    .. automethod:: __call__
+    '''
+    def __call__(self,
+                 *args,
+                 eeg: np.ndarray,
+                 baseline: Union[np.ndarray, None] = None,
+                 **kwargs) -> Dict[str, np.ndarray]:
+        r'''
+        Args:
+            eeg (np.ndarray): The input EEG signals.
+            baseline (np.ndarray, optional) : The corresponding baseline signal, if apply_to_baseline is set to True and baseline is passed, the baseline signal will be transformed with the same way as the experimental signal.
+
+        Returns:
+            np.ndarray: The transformed results.
+        '''
+        return super().__call__(*args, eeg=eeg, baseline=baseline, **kwargs)
+
+    def apply(self, eeg: np.ndarray, **kwargs) -> np.ndarray:
+        return eeg.reshape(-1)
diff --git a/torcheeg/transforms/numpy/rearrange.py b/torcheeg/transforms/numpy/rearrange.py
new file mode 100644
index 0000000..45ca87a
--- /dev/null
+++ b/torcheeg/transforms/numpy/rearrange.py
@@ -0,0 +1,88 @@
+from typing import Dict, List, Union
+
+import numpy as np
+
+from ..base_transform import EEGTransform
+
+
+class RearrangeElectrode(EEGTransform):
+    r'''
+    Select parts of electrode signals based on a given electrode index list.
+
+    .. code-block:: python
+
+        transform = RearrangeElectrode(
+            source=['FP1', 'F3', 'F7'],
+            target=['F3', 'F7', 'FP1', 'AF2'],
+            missing='mean'
+        )
+        transform(eeg=np.random.randn(3, 128))['eeg'].shape
+        >>> (4, 128)
+
+    Args:
+        source (list): The list of electrode names to be rearranged.
+        target (list): The list of electrode names to be rearranged to.
+        missing (str): The method to deal with missing electrodes. (default: :obj:`'random'`)
+
+    .. automethod:: __call__
+    '''
+    def __init__(self,
+                 source: List[str],
+                 target: List[str],
+                 missing: str = 'mean',
+                 apply_to_baseline: bool = False):
+        super(RearrangeElectrode,
+              self).__init__(apply_to_baseline=apply_to_baseline)
+        self.source = source
+        self.target = target
+
+        assert missing in [
+            'random', 'zero', 'mean'
+        ], f'Invalid missing method {missing}, should be one of [random, zero, mean].'
+
+        self.missing = missing
+
+    def __call__(self,
+                 *args,
+                 eeg: np.ndarray,
+                 baseline: Union[np.ndarray, None] = None,
+                 **kwargs) -> Dict[str, np.ndarray]:
+        r'''
+        Args:
+            eeg (np.ndarray): The input EEG signals in shape of [number of electrodes, number of data points].
+            baseline (np.ndarray, optional) : The corresponding baseline signal, if apply_to_baseline is set to True and baseline is passed, the baseline signal will be transformed with the same way as the experimental signal.
+
+        Returns:
+            np.ndarray: The output signals with the shape of [number of picked electrodes, number of data points].
+        '''
+        return super().__call__(*args, eeg=eeg, baseline=baseline, **kwargs)
+
+    def apply(self, eeg: np.ndarray, **kwargs) -> np.ndarray:
+        r'''
+        Args:
+            eeg (np.ndarray): The input EEG signals in shape of [number of electrodes, number of data points].
+
+        Returns:
+            np.ndarray: The output signals with the shape of [number of rearranged electrodes, number of data points].
+        '''
+        output = np.zeros((len(self.target), eeg.shape[1]))
+        for i, target in enumerate(self.target):
+            if target in self.source:
+                output[i] = eeg[self.source.index(target)]
+            else:
+                if self.missing == 'random':
+                    output[i] = np.random.randn(eeg.shape[1])
+                elif self.missing == 'zero':
+                    output[i] = np.zeros(eeg.shape[1])
+                elif self.missing == 'mean':
+                    output[i] = np.mean(eeg, axis=0)
+        return output
+
+    @property
+    def repr_body(self) -> Dict:
+        return dict(
+            super().repr_body, **{
+                'source': self.source,
+                'target': self.target,
+                'missing': self.missing
+            })
diff --git a/torcheeg/transforms/numpy/to.py b/torcheeg/transforms/numpy/to.py
index 6f28d94..7517eca 100644
--- a/torcheeg/transforms/numpy/to.py
+++ b/torcheeg/transforms/numpy/to.py
@@ -68,7 +68,10 @@ def __init__(self,
 
         loc_x_list = []
         loc_y_list = []
-        for _, (loc_y, loc_x) in channel_location_dict.items():
+        for _, locs in channel_location_dict.items():
+            if locs is None:
+                continue
+            (loc_y, loc_x) = locs
             loc_x_list.append(loc_x)
             loc_y_list.append(loc_y)
 
@@ -94,7 +97,10 @@ def apply(self, eeg: np.ndarray, **kwargs) -> np.ndarray:
         # num_electrodes x timestep
         outputs = np.zeros([self.height, self.width, eeg.shape[-1]])
         # 9 x 9 x timestep
-        for i, (loc_y, loc_x) in enumerate(self.channel_location_dict.values()):
+        for i, locs in enumerate(self.channel_location_dict.values()):
+            if locs is None:
+                continue
+            (loc_y, loc_x) = locs
             outputs[loc_y][loc_x] = eeg[i]
 
         outputs = outputs.transpose(2, 0, 1)
diff --git a/torcheeg/transforms/torch/random.py b/torcheeg/transforms/torch/random.py
index 189cbfb..015e510 100644
--- a/torcheeg/transforms/torch/random.py
+++ b/torcheeg/transforms/torch/random.py
@@ -131,6 +131,112 @@ def repr_body(self) -> Dict:
         return dict(super().repr_body, **{'ratio': self.ratio})
 
 
+class RandomMask2D(RandomEEGTransform):
+    '''
+    Overlay the EEG signal using a random mask, and the value of the overlaid data points was set to 0.0. The mask is applied to the three dimensions of the input tensor, where the first dimension represents the number of channels (bands, or something like that), the second dimension denotes the number of electrodes and the second dimension represents the number of time points. The mask is applied to the electrode dimension.
+    
+    .. code-block:: python
+
+        transform = RandomMask()
+        transform(eeg=torch.randn(1, 32, 128))['eeg'].shape
+        >>> (1, 32, 128)
+
+    Args:
+        ratio (float): The proportion of data points covered by the mask out of all data points for each EEG signal sample. (default: :obj:`0.5`)
+        p (float): Probability of applying random mask on EEG signal samples. Should be between 0.0 and 1.0, where 0.0 means no mask is applied to every sample and 1.0 means that masks are applied to every sample. (default: :obj:`0.5`)
+        apply_to_baseline: (bool): Whether to act on the baseline signal at the same time, if the baseline is passed in when calling. (default: :obj:`False`)
+
+    .. automethod:: __call__
+    '''
+    def __init__(self,
+                 ratio: float = 0.5,
+                 p: float = 0.5,
+                 apply_to_baseline: bool = False):
+        super(RandomMask2D, self).__init__(p=p,
+                                           apply_to_baseline=apply_to_baseline)
+        self.ratio = ratio
+
+    def __call__(self,
+                 *args,
+                 eeg: torch.Tensor,
+                 baseline: Union[torch.Tensor, None] = None,
+                 **kwargs) -> Dict[str, torch.Tensor]:
+        r'''
+        Args:
+            eeg (torch.Tensor): The input EEG signal.
+            baseline (torch.Tensor, optional) : The corresponding baseline signal, if apply_to_baseline is set to True and baseline is passed, the baseline signal will be transformed with the same way as the experimental signal.
+
+        Returns:
+            torch.Tensor: The output EEG signal after applying a random mask.
+        '''
+        return super().__call__(*args, eeg=eeg, baseline=baseline, **kwargs)
+
+    def random_apply(self, eeg: torch.Tensor, **kwargs) -> torch.Tensor:
+        # mask at the electrode dimension
+        mask = torch.rand(eeg.shape[1], device=eeg.device)
+        mask = (mask < self.ratio).to(eeg.dtype)
+        # to the same shape of eeg
+        mask = mask.unsqueeze(0).unsqueeze(2).repeat(eeg.shape[0], 1,
+                                                     eeg.shape[2])
+        return eeg * mask
+
+    @property
+    def repr_body(self) -> Dict:
+        return dict(super().repr_body, **{'ratio': self.ratio})
+
+
+class RandomMaskGrid(RandomEEGTransform):
+    '''
+    Overlay the EEG signal using a random mask, and the value of the overlaid data points was set to 0.0. The mask is applied to the last three dimensions of the input tensor, where the second dimension represents the height of the grid and the third dimension represents the width of the grid. The mask is applied to the height and width dimensions.
+    
+    .. code-block:: python
+
+        transform = RandomMaskGrid()
+        transform(eeg=torch.randn(4, 9, 9))['eeg'].shape
+        >>> (4, 9, 9)
+
+    Args:
+        ratio (float): The proportion of data points covered by the mask out of all data points for each EEG signal sample. (default: :obj:`0.5`)
+        p (float): Probability of applying random mask on EEG signal samples. Should be between 0.0 and 1.0, where 0.0 means no mask is applied to every sample and 1.0 means that masks are applied to every sample. (default: :obj:`0.5`)
+        apply_to_baseline: (bool): Whether to act on the baseline signal at the same time, if the baseline is passed in when calling. (default: :obj:`False`)
+
+    .. automethod:: __call__
+    '''
+    def __init__(self,
+                 ratio: float = 0.5,
+                 p: float = 0.5,
+                 apply_to_baseline: bool = False):
+        super(RandomMaskGrid,
+              self).__init__(p=p, apply_to_baseline=apply_to_baseline)
+        self.ratio = ratio
+
+    def __call__(self,
+                 *args,
+                 eeg: torch.Tensor,
+                 baseline: Union[torch.Tensor, None] = None,
+                 **kwargs) -> Dict[str, torch.Tensor]:
+        r'''
+        Args:
+            eeg (torch.Tensor): The input EEG signal.
+            baseline (torch.Tensor, optional) : The corresponding baseline signal, if apply_to_baseline is set to True and baseline is passed, the baseline signal will be transformed with the same way as the experimental signal.
+
+        Returns:
+            torch.Tensor: The output EEG signal after applying a random mask.
+        '''
+        return super().__call__(*args, eeg=eeg, baseline=baseline, **kwargs)
+
+    def random_apply(self, eeg: torch.Tensor, **kwargs) -> torch.Tensor:
+        # mask at the electrode dimension
+        mask = torch.rand(eeg.shape[1], eeg.shape[2], device=eeg.device)
+        mask = (mask < self.ratio).to(eeg.dtype)
+        mask = mask.unsqueeze(0).repeat(eeg.shape[0], 1, 1)
+        return eeg * mask
+
+    @property
+    def repr_body(self) -> Dict:
+        return dict(super().repr_body, **{'ratio': self.ratio})
+
+
 class RandomWindowSlice(RandomEEGTransform):
     '''
     Randomly applies a slice transformation with a given probability, where the original time series is sliced by a window, and the sliced data is scaled to the original size. It is worth noting that the random position where each channel slice starts is the same.


Test patch:
diff --git a/test/datasets/test_folder.py b/test/datasets/test_folder.py
index 8f5ee8d..492401b 100644
--- a/test/datasets/test_folder.py
+++ b/test/datasets/test_folder.py
@@ -1,12 +1,12 @@
 import os
 import shutil
 import unittest
-
+import pandas as pd
 import mne
 import numpy as np
 
 from torcheeg import transforms
-from torcheeg.datasets import FolderDataset
+from torcheeg.datasets import FolderDataset, CSVFolderDataset
 
 mne.set_log_level('CRITICAL')
 
@@ -56,7 +56,6 @@ def test_folder_dataset(self):
         label_map = {'folder1': 0, 'folder2': 1}
         dataset = FolderDataset(io_path=io_path,
                                 root_path=root_path,
-                                num_channel=14,
                                 online_transform=transforms.ToTensor(),
                                 label_transform=transforms.Compose([
                                     transforms.Select('label'),
@@ -67,11 +66,50 @@ def test_folder_dataset(self):
         self.assertEqual(len(dataset), 20)
         first_item = dataset[0]
         self.assertEqual(first_item[0].shape, (14, 128))
-        self.assertEqual(first_item[1], 0)
         last_item = dataset[-1]
         self.assertEqual(last_item[0].shape, (14, 128))
-        self.assertEqual(last_item[1], 1)
 
+    def test_csv_folder_dataset(self):
+        # Generate dummy EEG files and save them to folders
+        folder1 = './tmp_in/data/label1'
+        folder2 = './tmp_in/data/label2'
+
+        if not os.path.exists(folder1):
+            os.makedirs(folder1)
+            self.generate_dummy_eeg_files(2, folder1)
+        if not os.path.exists(folder2):
+            os.makedirs(folder2)
+            self.generate_dummy_eeg_files(2, folder2)
+
+        # Define input and output paths for folder_constructor
+        csv_path = './tmp_out/data.csv'
+        io_path = './tmp_out/eeg_folder'
+
+        df = pd.DataFrame({
+            'subject_id': ['sub1', 'sub1', 'sub2', 'sub2'],
+            'trial_id': [0, 1, 0, 1],
+            'label': [0, 1, 0, 1],
+            'file_path': [
+                os.path.join(folder1, 'sub1.fif'),
+                os.path.join(folder2, 'sub1.fif'),
+                os.path.join(folder1, 'sub2.fif'),
+                os.path.join(folder2, 'sub2.fif'),
+            ]
+        })
+
+        df.to_csv(csv_path, index=False)
+
+        dataset = CSVFolderDataset(csv_path=csv_path,
+                                   io_path=io_path,
+                                   online_transform=transforms.ToTensor(),
+                                   label_transform=transforms.Select('label'),
+                                   num_worker=0)
+
+        self.assertEqual(len(dataset), 20)
+        first_item = dataset[0]
+        self.assertEqual(first_item[0].shape, (14, 128))
+        last_item = dataset[-1]
+        self.assertEqual(last_item[0].shape, (14, 128))
 
 if __name__ == '__main__':
     unittest.main()
diff --git a/test/models/test_transformer.py b/test/models/test_transformer.py
index cab2456..c403ed5 100644
--- a/test/models/test_transformer.py
+++ b/test/models/test_transformer.py
@@ -2,7 +2,7 @@
 
 import torch
 
-from torcheeg.models import SimpleViT, ArjunViT, VanillaTransformer, ViT
+from torcheeg.models import SimpleViT, ArjunViT, VanillaTransformer, ViT, Conformer
 
 
 class TestTransformer(unittest.TestCase):
@@ -59,6 +59,17 @@ def test_vit(self):
         pred = model(eeg).cpu()
         self.assertEqual(tuple(pred.shape), (1, 2))
 
+    def test_conformer(self):
+        eeg = torch.randn(2, 1, 32, 128)
+        model = Conformer(num_electrodes=32, sampling_rate=128, num_classes=2)
+        pred = model(eeg)
+        self.assertEqual(tuple(pred.shape), (2, 2))
+
+        eeg = eeg.cuda()
+        model = model.cuda()
+        pred = model(eeg).cpu()
+        self.assertEqual(tuple(pred.shape), (2, 2))
+
 
 if __name__ == '__main__':
     unittest.main()
diff --git a/test/trainers/test_glow_trainer.py b/test/trainers/test_glow_trainer.py
index 30fafb4..3d931e7 100644
--- a/test/trainers/test_glow_trainer.py
+++ b/test/trainers/test_glow_trainer.py
@@ -22,20 +22,20 @@ def __getitem__(self, i) -> int:
 
 
 class TestGlowTrainer(unittest.TestCase):
-    # def test_glow_trainer(self):
-    #     train_dataset = DummyDataset()
-    #     val_dataset = DummyDataset()
-    #     test_dataset = DummyDataset()
+    def test_glow_trainer(self):
+        train_dataset = DummyDataset()
+        val_dataset = DummyDataset()
+        test_dataset = DummyDataset()
 
-    #     train_loader = DataLoader(train_dataset, batch_size=4)
-    #     val_loader = DataLoader(val_dataset, batch_size=4)
-    #     test_loader = DataLoader(test_dataset, batch_size=4)
+        train_loader = DataLoader(train_dataset, batch_size=4)
+        val_loader = DataLoader(val_dataset, batch_size=4)
+        test_loader = DataLoader(test_dataset, batch_size=4)
 
-    #     model = BGlow(in_channels=4)
+        model = BGlow(in_channels=4)
 
-    #     trainer = GlowTrainer(model, accelerator='gpu')
-    #     trainer.fit(train_loader, val_loader, max_epochs=1)
-    #     trainer.test(test_loader)
+        trainer = GlowTrainer(model, accelerator='gpu')
+        trainer.fit(train_loader, val_loader, max_epochs=1)
+        trainer.test(test_loader)
 
     def test_cglow_trainer(self):
         train_dataset = DummyDataset()
diff --git a/test/trainers/test_contrastive_trainer.py b/test/trainers/test_self_supervised_trainer.py
similarity index 61%
rename from test/trainers/test_contrastive_trainer.py
rename to test/trainers/test_self_supervised_trainer.py
index 050e8d7..aba67f1 100644
--- a/test/trainers/test_contrastive_trainer.py
+++ b/test/trainers/test_self_supervised_trainer.py
@@ -5,7 +5,7 @@
 import torch.nn as nn
 from torch.utils.data import Dataset
 from torch.utils.data.dataloader import DataLoader
-from torcheeg.trainers import ContrastiveTrainer
+from torcheeg.trainers import SimCLRTrainer, BYOLTrainer
 
 
 class DummyDataset(Dataset):
@@ -22,7 +22,7 @@ def __getitem__(self, i) -> int:
 
 class DummyModel(nn.Module):
 
-    def __init__(self, in_channels=120, out_channels=2):
+    def __init__(self, in_channels=120, out_channels=10):
         super().__init__()
         self.in_channels = in_channels
         self.out_channels = out_channels
@@ -32,9 +32,9 @@ def forward(self, x):
         return self.fc(x)
 
 
-class TestClassificationTrainer(unittest.TestCase):
+class TestSelfSupervisedTrainer(unittest.TestCase):
 
-    def test_classification_trainer(self):
+    def test_sim_clr_trainer(self):
         train_dataset = DummyDataset()
         val_dataset = DummyDataset()
 
@@ -43,14 +43,21 @@ def test_classification_trainer(self):
 
         model = DummyModel()
 
-        trainer = ContrastiveTrainer(model)
+        trainer = SimCLRTrainer(model, extract_channels=10)
+        trainer.fit(train_loader, val_loader, max_epochs=1)
+
+    def test_byol_trainer(self):
+        train_dataset = DummyDataset()
+        val_dataset = DummyDataset()
+
+        train_loader = DataLoader(train_dataset, batch_size=2)
+        val_loader = DataLoader(val_dataset, batch_size=2)
+
+        model = DummyModel()
+
+        trainer = BYOLTrainer(model, extract_channels=10)
         trainer.fit(train_loader, val_loader, max_epochs=1)
 
-        trainer = ContrastiveTrainer(model,
-                                     devices=1,
-                                     accelerator='gpu')
-        trainer.fit(train_loader, val_loader,
-                                     max_epochs=1)
 
 
 if __name__ == '__main__':
diff --git a/test/transforms/test_label.py b/test/transforms/test_label.py
index 14f2145..a6c5906 100644
--- a/test/transforms/test_label.py
+++ b/test/transforms/test_label.py
@@ -1,6 +1,6 @@
 import unittest
 
-from torcheeg.transforms import Select, Binary, BinariesToCategory, StringToInt, BinaryOneVSRest, FixCategory
+from torcheeg.transforms import Select, Binary, BinariesToCategory, StringToInt, BinaryOneVSRest, FixCategory, Mapping
 
 
 class TestLabelTransforms(unittest.TestCase):
@@ -41,6 +41,14 @@ def test_fix_category(self):
         self.assertEqual(FixCategory(value=0)(y=3)['y'], 0)
         self.assertEqual(FixCategory(value=[0, 1])(y=[1, 2])['y'], [0, 1])
 
+    def test_mapping(self):
+        transform = Mapping({
+            'left_hand': 0,
+            'right_hand': 1,
+        })
+        self.assertEqual(transform(y='left_hand')['y'], 0)
+        self.assertEqual(transform(y=['left_hand', 'right_hand'])['y'], [0, 1])
+
 
 if __name__ == '__main__':
     unittest.main()
diff --git a/test/transforms/test_numpy.py b/test/transforms/test_numpy.py
index cf16e1e..9545bdb 100644
--- a/test/transforms/test_numpy.py
+++ b/test/transforms/test_numpy.py
@@ -1,11 +1,16 @@
 import unittest
 
 import numpy as np
-from torcheeg.transforms import ToGrid, ToInterpolatedGrid, To2d, MeanStdNormalize, MinMaxNormalize, BandSignal, BandDifferentialEntropy, BandPowerSpectralDensity, BandMeanAbsoluteDeviation, BandKurtosis, BandSkewness, Concatenate, MapChunk, PickElectrode, CWTSpectrum, ARRCoefficient, PearsonCorrelation, PhaseLockingCorrelation, BandApproximateEntropy, BandSampleEntropy, BandSVDEntropy, BandDetrendedFluctuationAnalysis, BandHiguchiFractalDimension, BandHjorth, BandHurst, BandPetrosianFractalDimension, BandBinPower, BandSpectralEntropy, DWTDecomposition, Downsample, Compose
+from torcheeg.transforms import ToGrid, ToInterpolatedGrid, To2d, MeanStdNormalize, MinMaxNormalize, BandSignal, BandDifferentialEntropy, BandPowerSpectralDensity, BandMeanAbsoluteDeviation, BandKurtosis, BandSkewness, Concatenate, MapChunk, PickElectrode, CWTSpectrum, ARRCoefficient, PearsonCorrelation, PhaseLockingCorrelation, BandApproximateEntropy, BandSampleEntropy, BandSVDEntropy, BandDetrendedFluctuationAnalysis, BandHiguchiFractalDimension, BandHjorth, BandHurst, BandPetrosianFractalDimension, BandBinPower, BandSpectralEntropy, DWTDecomposition, Downsample, Compose, RearrangeElectrode, Flatten
 from torcheeg.datasets.constants import DEAP_CHANNEL_LOCATION_DICT, DEAP_CHANNEL_LIST, M3CV_CHANNEL_LOCATION_DICT
 
 
 class TestNumpyTransforms(unittest.TestCase):
+    def test_flatten(self):
+        eeg = np.random.randn(62, 5)
+        transformed_eeg = Flatten()(eeg=eeg)
+        self.assertEqual(transformed_eeg['eeg'].shape, (310,))
+
     def test_dwt_decomposition(self):
         eeg = np.random.randn(32, 1000)
         transformed_eeg = DWTDecomposition()(eeg=eeg)
@@ -215,6 +220,17 @@ def test_downsample(self):
         transformed_eeg = Downsample(num_points=32, axis=-1)(eeg=eeg)
         self.assertEqual(transformed_eeg['eeg'].shape, (32, 32))
 
+    def test_rearrange_electrode(self):
+        src_eeg = np.random.rand(3, 128)
+        tgt_eeg = RearrangeElectrode(source=['FP1', 'F3', 'F7'],
+                                 target=['F3', 'F7', 'FP1', 'AF2'],
+                                 missing='mean')(eeg=src_eeg)['eeg']
+        self.assertEqual(tgt_eeg.shape, (4, 128))
+
+        self.assertEqual(src_eeg[0, :].tolist(), tgt_eeg[2, :].tolist())
+        self.assertEqual(src_eeg[1, :].tolist(), tgt_eeg[0, :].tolist())
+        self.assertEqual(src_eeg[2, :].tolist(), tgt_eeg[1, :].tolist())
+
 
 if __name__ == '__main__':
     unittest.main()
\ No newline at end of file
diff --git a/test/transforms/test_torch.py b/test/transforms/test_torch.py
index 55ae5c4..f4d3309 100644
--- a/test/transforms/test_torch.py
+++ b/test/transforms/test_torch.py
@@ -1,7 +1,7 @@
 import unittest
 import torch
 import numpy as np
-from torcheeg.transforms import ToTensor, Resize, RandomNoise, RandomMask, RandomWindowSlice, RandomWindowWarp, RandomPCANoise, RandomFlip, RandomShift, RandomChannelShuffle, RandomFrequencyShift, RandomSignFlip, Contrastive
+from torcheeg.transforms import ToTensor, Resize, RandomNoise, RandomMask, RandomWindowSlice, RandomWindowWarp, RandomPCANoise, RandomFlip, RandomShift, RandomChannelShuffle, RandomFrequencyShift, RandomSignFlip, Contrastive, RandomMaskGrid, RandomMask2D
 
 
 class TestTorchTransforms(unittest.TestCase):
@@ -106,6 +106,16 @@ def test_random_sampling_rate_shift(self):
         transformed_eeg = RandomFrequencyShift(p=1.0, series_dim=0)(eeg=eeg)
         self.assertEqual(tuple(transformed_eeg['eeg'].shape), (128, 9, 9))
 
+    def test_random_mask_2d(self):
+        eeg = torch.randn(1, 32, 128)
+        transformed_eeg = RandomMask2D()(eeg=eeg)
+        self.assertEqual(transformed_eeg['eeg'].shape, (1, 32, 128))
+    
+    def test_random_mask_grid(self):
+        eeg = torch.randn(32, 9, 9)
+        transformed_eeg = RandomMaskGrid()(eeg=eeg)
+        self.assertEqual(transformed_eeg['eeg'].shape, (32, 9, 9))
+
 
 if __name__ == '__main__':
     unittest.main()
\ No newline at end of file