diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index 80f2ee3..9c1ae50 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -1,8 +1,8 @@ -name: ssl_tools +name: ssl_tools documentation on: push: branches: - - main + - docs jobs: docs: @@ -17,6 +17,11 @@ jobs: with: python-version: 3.12.1 + - name: Install packages + run: | + sudo apt-get update + sudo apt-get install -y pandoc + - name: Install requirements run: | pip3 install sphinx sphinx-rtd-theme sphinx-autodoc-typehints sphinx-argparse sphinx-autoapi nbsphinx pandoc Ipython diff --git a/.gitignore b/.gitignore index 1a75871..fe06415 100644 --- a/.gitignore +++ b/.gitignore @@ -168,4 +168,5 @@ cython_debug/ # option (not recommended) you can uncomment the following to ignore the entire idea folder. #.idea/ ssl_tools/experiments/**/logs/ -data/* \ No newline at end of file +data/* +notebooks/logs/* \ No newline at end of file diff --git a/docs/contributing.rst b/docs/contributing.rst new file mode 100644 index 0000000..3fa64d7 --- /dev/null +++ b/docs/contributing.rst @@ -0,0 +1,4 @@ +Contributing +------------ + +Under construction... \ No newline at end of file diff --git a/docs/experiments.rst b/docs/experiments.rst new file mode 100644 index 0000000..eadc650 --- /dev/null +++ b/docs/experiments.rst @@ -0,0 +1,4 @@ +Running Experiments +-------------------- + +Under construction... \ No newline at end of file diff --git a/docs/index.rst b/docs/index.rst index 403e525..91a42fa 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -11,6 +11,9 @@ Welcome to SSLTools's documentation! :caption: Contents: installation + tutorials + experiments + contributing api diff --git a/docs/installation.rst b/docs/installation.rst index 66a4455..88d3618 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -1,2 +1,4 @@ Installation ------------------ \ No newline at end of file +----------------- + +Under construction... \ No newline at end of file diff --git a/docs/modules.rst b/docs/modules.rst new file mode 100644 index 0000000..973898d --- /dev/null +++ b/docs/modules.rst @@ -0,0 +1,7 @@ +ssl_tools +========= + +.. toctree:: + :maxdepth: 4 + + ssl_tools diff --git a/docs/notebooks b/docs/notebooks new file mode 120000 index 0000000..edb8f02 --- /dev/null +++ b/docs/notebooks @@ -0,0 +1 @@ +../notebooks/ \ No newline at end of file diff --git a/docs/ssl_tools.analysis.rst b/docs/ssl_tools.analysis.rst new file mode 100644 index 0000000..9bf3480 --- /dev/null +++ b/docs/ssl_tools.analysis.rst @@ -0,0 +1,21 @@ +ssl\_tools.analysis package +=========================== + +Submodules +---------- + +ssl\_tools.analysis.plot\_metrics module +---------------------------------------- + +.. automodule:: ssl_tools.analysis.plot_metrics + :members: + :undoc-members: + :show-inheritance: + +Module contents +--------------- + +.. automodule:: ssl_tools.analysis + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/tutorials.rst b/docs/tutorials.rst new file mode 100644 index 0000000..2e861ae --- /dev/null +++ b/docs/tutorials.rst @@ -0,0 +1,13 @@ +.. _tutorials: + +========================== +Tutorials +========================== + +.. toctree:: + :maxdepth: 2 + + notebooks/01_structuring_input.ipynb + notebooks/02_training_model.ipynb + notebooks/03_training_ssl_model.ipynb + notebooks/04_using_experiments.ipynb \ No newline at end of file diff --git a/notebooks/01_structuring_input.ipynb b/notebooks/01_structuring_input.ipynb new file mode 100644 index 0000000..bce7443 --- /dev/null +++ b/notebooks/01_structuring_input.ipynb @@ -0,0 +1,502 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# 1. Structuring the input\n", + "\n", + "In order to train and test models, we need to structure the input data in a way that is compatible with the model and framweork's experiments.\n", + "\n", + "This framework is designed to work with time-series data and Pytorch Lightning. \n", + "Thus, it provides the necessary tools to create a `Dataset` object and a `LightningDataModule` object.\n", + "\n", + "The `Dataset` object is responsible for loading the data. It is a Pytorch object that is used to load the data and make it available to the model. \n", + "Every `Dataset` class must implement two methods: `__len__` and `__getitem__`.\n", + "The `__len__` method returns the number of samples in the dataset, and the `__getitem__`, given an integer from 0 to `__len__` - 1, returns the corresponding sample from the dataset.\n", + "The returned type of the `__getitem__` method is not specified, but it is usually a 2-element tuple with the input and the target. The input is the data that will be used to make the predictions, and the target is the data that the model will try to predict.\n", + "\n", + "For now, this framework provide implementations for the `Dataset` objects for time-series data, where data is organized in two different ways:\n", + "\n", + "- A directory with several CSV files, where each file contains a time-series. Each row in a CSV file is a time-step, each column is a feature, and the whole file is a time-series. This is handled by the `SeriesFolderCSVDataset` class.\n", + "- A single CSV file with a windowed time-series. Each row in the CSV file is a window, and each column is a feature. This is handled by the `MultiModalSeriesCSVDataset` class.\n", + "\n", + "We explain both classes in detais nextly." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Time-series dataset implementations" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### SeriesFolderCSVDataset\n", + "\n", + "The `SeriesFolderCSVDataset` class is designed to work with a directory containing several CSV files, where each file represent a time-series. \n", + "Each row in a CSV file is a time-step, and each column is a feature. \n", + "\n", + "This class assumes that data is organized in the following way, where `my_dataset` is the path to the directory containing the CSV files:\n", + "\n", + "```bash\n", + "my_dataset/\n", + " series1.csv\n", + " series2.csv\n", + " other_series.csv\n", + " ...\n", + "```\n", + "\n", + "Where each CSV file represents a time-series. \n", + "\n", + "| accel-x | accel-y | accel-z | gyro-x | gyro-y | gyro-z | class |\n", + "|---------|---------|---------|---------|---------|---------|---------|\n", + "| 0.502123| 0.02123 | 0.12312 | 0.12312 | 0.12312 | 0.12312 | 1 |\n", + "| 0.682012| 0.02123 | 0.12312 | 0.12312 | 0.12312 | 0.12312 | 1 |\n", + "| 0.498217| 0.00001 | 0.12312 | 0.12312 | 0.12312 | 0.12312 | 1 |\n", + "\n", + "\n", + "Note that the CSV must have a header with the column names." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "To handle this kind of data, we use the `SeriesFolderCSVDataset` class. This class is a Pytorch `Dataset` object that loads the data from the CSV files and makes it available to the model.\n", + "For this class, we must specify the path to the directory containing the CSV files, the name of the columns that will be used as features, and the name of the column that will be used as the target.\n", + "Note that, each feature (column) represent a dimension of the time-series, while the rows represent the time-steps.\n", + "\n", + "Thus, the `SeriesFolderCSVDataset` class minimally requires:\n", + "\n", + "- `data_path`: the path to the directory containing the CSV files\n", + "- `features`: a list of strings with the names of the features columns, e.g. `['accel-x', 'accel-y', 'accel-z', 'gyro-x', 'gyro-y', 'gyro-z']`\n", + "- `label`: a string with the name of the label column, e.g. `'class'`" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/usr/local/lib/python3.10/dist-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n", + " from .autonotebook import tqdm as notebook_tqdm\n" + ] + }, + { + "data": { + "text/plain": [ + "SeriesFolderCSVDataset at /workspaces/hiaac-m4/ssl_tools/data/view_concatenated/KuHar_cpc/train (57 samples)" + ] + }, + "execution_count": 1, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "from ssl_tools.data.datasets import SeriesFolderCSVDataset\n", + "\n", + "# Path to the data\n", + "data_path = \"/workspaces/hiaac-m4/ssl_tools/data/view_concatenated/KuHar_cpc/train\"\n", + "\n", + "# Creating the dataset\n", + "dataset = SeriesFolderCSVDataset(\n", + " data_path=data_path,\n", + " features=[\"accel-x\", \"accel-y\", \"accel-z\", \"gyro-x\", \"gyro-y\", \"gyro-z\"],\n", + " label=\"standard activity code\",\n", + ")\n", + "\n", + "dataset" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can get the number of samples in the dataset with the `len` function, and we can retrive a sample with the `__getitem__` method, that is, using `[]`, such as `dataset[0]`.\n", + "The dataset may return:\n", + "\n", + "- A 2-element tuple, where the first element is a 2D numpy array with shape `(num_features, time_steps)`, and the second element is a 1D tensor with shape `(time_steps,)`.\n", + "- A 2D numpy array with shape `(num_features, time_steps)`, if `label` is `None`, at the time of the dataset object's creation.\n", + "\n", + "Let's check the number of samples and access the first sample and its label." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Length of dataset: 57 samples\n" + ] + } + ], + "source": [ + "# Gte the length of the dataset\n", + "length_of_dataset = len(dataset)\n", + "print(f\"Length of dataset: {length_of_dataset} samples\")" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Type of sample: tuple with 2 elements\n" + ] + } + ], + "source": [ + "# Get the first sample\n", + "sample = dataset[0]\n", + "type_of_sample = type(sample).__name__\n", + "print(f\"Type of sample: {type_of_sample} with {len(sample)} elements\")" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Shape of sample: (6, 2586), shape of label: (2586, 1)\n" + ] + } + ], + "source": [ + "# The first element of the sample is the input, while the second element is the label\n", + "# We can split the sample into input and label variables\n", + "shape_of_sample = sample[0].shape\n", + "shape_of_label = sample[1].shape\n", + "print(f\"Shape of sample: {shape_of_sample}, shape of label: {shape_of_label}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "You can see above, the sample is a 2-element tuple. The first element is a 2D numpy array with shape `(6, 2586)`, and the second element is a 1D tensor with shape `(2586,)`, that is, a label for each time-step." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### MultiModalSeriesCSVDataset\n", + "\n", + "\n", + "The `MultiModalSeriesCSVDataset` class is designed to work with a single CSV file containing a windowed time-series. \n", + "The CSV is a multi-modal time-series, where each row is a sample and each column is a feature at a given time-step. \n", + "Features are organized in a way that each group of columns represent a different modality.\n", + "\n", + "The CSV file looks like this:\n", + "\n", + "\n", + "| accel-x-0 | accel-x-1 | accel-y-0 | accel-y-1 | class |\n", + "|-----------|-----------|-----------|-----------|--------|\n", + "| 0.502123 | 0.02123 | 0.502123 | 0.502123 | 0 |\n", + "| 0.6820123 | 0.02123 | 0.502123 | 0.502123 | 1 |\n", + "| 0.498217 | 0.00001 | 1.414141 | 3.141592 | 1 |\n", + "\n", + "In the example, columns `accel-x-0` and `accel-x-1` are the `accel-x` feature at time `0` and time `1`, respectively. The same goes for the `accel-y` feature. Finally, the `class` column is the label. " + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "To handle this kind of data, we use the `MultiModalSeriesCSVDataset` class.\n", + "For this class, we must specify the path to the CSV file, the prefix of the columns that will be used as features, and the columns that will be used as label.\n", + "Note that, each feature (column) represent a dimension of the time-series, while the rows represent the samples.\n", + "\n", + "The `MultiModalSeriesCSVDataset` class minimally requires:\n", + "\n", + "- `data_path`: the path to the CSV file\n", + "- `feature_prefixes`: a list of strings with the prefixes of the feature columns, e.g. `['accel-x', 'accel-y']`. The class will look for columns with these prefixes and will consider them as features of a modality.\n", + "- `label`: a string with the name of the label column, e.g. `'class'`\n", + "- `features_as_channels`: a boolean indicating if the features should be treated as channels, that is, if each prefix will become a channel. If ``True``, the data will be returned as a vector of shape `(C, T)`, where C is the number of channels (features/prefixes) and `T` is the number of time steps. Else, the data will be returned as a vector of shape `T*C` (a single vector with all the features).\n", + "\n", + "Let's show how to read this data and create a `MultiModalSeriesCSVDataset` object." + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "MultiModalSeriesCSVDataset at /workspaces/hiaac-m4/ssl_tools/data/standartized_balanced/KuHar/train.csv (1386 samples)" + ] + }, + "execution_count": 5, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "from ssl_tools.data.datasets import MultiModalSeriesCSVDataset\n", + "\n", + "# Path to the data\n", + "data_path = \"/workspaces/hiaac-m4/ssl_tools/data/standartized_balanced/KuHar/train.csv\"\n", + "\n", + "# Instantiate the dataset\n", + "dataset = MultiModalSeriesCSVDataset(\n", + " data_path=data_path,\n", + " feature_prefixes=[\"accel-x\", \"accel-y\", \"accel-z\", \"gyro-x\", \"gyro-y\", \"gyro-z\"],\n", + " label=\"standard activity code\",\n", + " features_as_channels = True,\n", + ")\n", + "\n", + "dataset" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can get the number of samples in the dataset with the `len` function, and we can retrive a sample with the `__getitem__` method, that is, using `[]`, such as `dataset[0]`.\n", + "The dataset may return:\n", + "\n", + "- A 2-element tuple, where the first element is a 2D numpy array with shape `(num_features, time_steps)`, and the second element is a 1D tensor with shape `(time_steps,)`.\n", + "- A 2D numpy array with shape `(num_features, time_steps)`, if `label` is `None`, at the time of the dataset object's creation.\n", + "\n", + "Let's check the number of samples and access the first sample and its label." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Length of dataset: 1386 samples\n" + ] + } + ], + "source": [ + "length_of_dataset = len(dataset)\n", + "print(f\"Length of dataset: {length_of_dataset} samples\")" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Type of sample: tuple with length 2 elements\n" + ] + } + ], + "source": [ + "sample = dataset[0]\n", + "type_of_sample = type(sample).__name__\n", + "print(f\"Type of sample: {type_of_sample} with length {len(sample)} elements\")" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Shape of sample: (6, 60)\n" + ] + } + ], + "source": [ + "shape_of_sample = sample[0].shape\n", + "print(f\"Shape of sample: {shape_of_sample}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Loading batches of data\n", + "\n", + "Pytorch models are trained using batches of data. Thus, we do not feed the model with a single sample at a time, but with a batch of samples.\n", + "If we see the last example, the `MultiModalSeriesCSVDataset` object returns a single sample at a time. Each sample is a 2-element tuple, where first element is a `(6, 60)` numpy array and the second is an integer, representing the label.\n", + "\n", + "A batch of samples add an extra dimension to the data. Thus, in our case, a batch of samples is a 3D tensor, where the first dimension is the batch size (`B`), the second dimension is the number of features, or channels (`C`), and the third dimension is the number of time-steps (`T`).\n", + "Thus, if the data have the shape `(6, 60)`, a batch of samples will have the shape `(B, 6, 60)`. The same happens to `label`, which gains an extra dimension. \n", + "\n", + "The batching of samples is done using a `DataLoader` object. This object is a Pytorch object that takes a `Dataset` object and returns batches of samples. The `DataLoader` object is responsible for shuffling the data, dividing it into batches, and loading the data in parallel.\n", + "Thus, given a `Dataset` object, we can easilly create a `DataLoader` object using the `torch.utils.data.DataLoader` class." + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "" + ] + }, + "execution_count": 9, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "from torch.utils.data import DataLoader\n", + "\n", + "# Create a DataLoader\n", + "dataloader = DataLoader(dataset, batch_size=32, shuffle=True)\n", + "dataloader" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can fetch a batch of samples from the `DataLoader` object using the `__iter__` method, that is, using a `for` loop. Each iteration returns a batch of samples.\n", + "In our case, each batch is a 2-element tuple, where the first element is a 3D tensor with shape `(B, C, T)`, and the second element is a 1D tensor with shape `(B,)`." + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Inputs shape: torch.Size([32, 6, 60]), labels shape: torch.Size([32])\n" + ] + } + ], + "source": [ + "for batch in dataloader:\n", + " inputs, labels = batch\n", + " print(f\"Inputs shape: {inputs.shape}, labels shape: {labels.shape}\")\n", + " break" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Handling data splits (train, validation, and test)\n", + "\n", + "Usually, we create a `DataLoader` object for the training data, another for the validation data, and another for the test data. \n", + "A simple way to do this is to create a `LightningDataModule` object, which is a Pytorch Lightning object that is responsible for creating the `DataLoader` objects for the training, validation, and test data.\n", + "\n", + "A `LightningDataModule` object is responsible for splitting the data into training, validation, and test sets, and creating the `DataLoader` objects for each set. This object may also be responsible for setting up the data, such as downloading the data from the internet, checking the data, and add the augmentations. This module is used to encapsulate all the data loading and processing logic in a single place, and to make it easy to use the same data processing logic across different experiments.\n", + "\n", + "The `LightningDataModule` object must implement three methods: `setup`, `train_dataloader`, and `val_dataloader`. The `setup` is optional, and is responsible for splitting the data into training, validation, and test sets, and the `train_dataloader` and `val_dataloader` methods are responsible for creating the `DataLoader` objects for the training and validation sets, respectively." + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "metadata": {}, + "outputs": [], + "source": [ + "import lightning as L\n", + "\n", + "class MyDataModule(L.LightningDataModule):\n", + " def __init__(self, train_csv_path, batch_size=32):\n", + " super().__init__()\n", + " self.data_path = data_path\n", + " self.batch_size = batch_size\n", + " self.train_csv_path = train_csv_path\n", + "\n", + " def setup(self, stage=None):\n", + " if stage == \"fit\" or stage is None:\n", + " self.train_dataset = MultiModalSeriesCSVDataset(\n", + " data_path=self.train_csv_path,\n", + " feature_prefixes=[\"accel-x\", \"accel-y\", \"accel-z\", \"gyro-x\", \"gyro-y\", \"gyro-z\"],\n", + " label=\"standard activity code\",\n", + " features_as_channels=True,\n", + " )\n", + "\n", + " def train_dataloader(self):\n", + " return DataLoader(\n", + " self.train_dataset, batch_size=self.batch_size, num_workers=self.num_workers, shuffle=True\n", + " )" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "When training a Pytorch Lightning model, we pass a `LightningDataModule` object to the `Trainer` object, and the `Trainer` object is responsible calling the `setup`, `train_dataloader`, and `val_dataloader` methods, and for training the model." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Summary\n", + "\n", + "First, we need to check which is our dataset and how the data is organized.\n", + "\n", + "- If data is organized in a directory with several CSV files, we use the `SeriesFolderCSVDataset` class. \n", + "- If data is organized in a single CSV file with a windowed time-series, we use the `MultiModalSeriesCSVDataset` class.\n", + "\n", + "Then, we create a `Dataset` object, and use it to create a `DataLoader` object.\n", + "\n", + "In order to organize the creation of the `DataLoader` object for each split (train, validation and test), we encapsule this logic in a `LightningDataModule` object.\n", + "The `LightningDataModule` object is responsible for creating the `DataLoader` objects for the training (`train_dataloader`), validation (`val_dataloader`), and test (`test_dataloader`) data, and for setting up the data.\n", + "\n", + "The `LightningDataModule` object is then used to train Pytorch Lightning models, which will call the `setup`, `train_dataloader`, and `val_dataloader` methods, corretly, as needed in the training/test process." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.6" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/notebooks/02_training_model.ipynb b/notebooks/02_training_model.ipynb new file mode 100644 index 0000000..e7260ef --- /dev/null +++ b/notebooks/02_training_model.ipynb @@ -0,0 +1,468 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# 2. Training a Pytorch Lighning model\n", + "\n", + "In this notebook, we show the training of a simple CNN model using Pytorch Lightning. \n", + "We first start with data, then we define the model, and finally we train it." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Creating KuHar LightningDataModule\n", + "\n", + "In order to train a model, we must first create a LightningDataModule.\n", + "In this work, we will use the Standartized KuHar HAR data. Our data folder looks like this:\n", + "\n", + "```\n", + "KuHar/\n", + " test.csv\n", + " train.csv\n", + " validation.csv\n", + "```\n", + "\n", + "The `train.csv` file may look like this:\n", + "\n", + "| accel-x-0 | accel-x-1 | accel-y-0 | accel-y-1 | class |\n", + "|-----------|-----------|-----------|-----------|--------|\n", + "| 0.502123 | 0.02123 | 0.502123 | 0.502123 | 0 |\n", + "| 0.6820123 | 0.02123 | 0.502123 | 0.502123 | 1 |\n", + "| 0.498217 | 0.00001 | 1.414141 | 3.141592 | 1 |\n", + "\n", + "As each CSV file contains time-windows signals of two 3-axis sensors (accelerometer and gyroscope), we must use the `MultiModalSeriesCSVDataset` class. After it, we must create a LightningDataModule, that will define the data loaders for training, validation and test. \n", + "\n", + "### Faciliting the creation of the LightningDataModule with MultiModalHARSeriesDataModule\n", + "\n", + "In order to facilitate the `Dataset` and `DataLoader` creation, we will use the `MultiModalHARSeriesDataModule`. If:\n", + "\n", + "1. Your directory is organized like the one above; and \n", + "2. Each CSV file is a collection os time-windows of signals (that possibly would be used as a dataset wrapping `MultiModalSeriesCSVDataset`).\n", + "\n", + "Then, you can use the `The `train.csv` file may look like this:\n", + "\n", + "| accel-x-0 | accel-x-1 | accel-y-0 | accel-y-1 | class |\n", + "|-----------|-----------|-----------|-----------|--------|\n", + "| 0.502123 | 0.02123 | 0.502123 | 0.502123 | 0 |\n", + "| 0.6820123 | 0.02123 | 0.502123 | 0.502123 | 1 |\n", + "| 0.498217 | 0.00001 | 1.414141 | 3.141592 | 1 |\n", + "\n", + "As each CSV file contains time-windows signals of two 3-axis sensors (accelerometer and gyroscope), we must use the `MultiModalSeriesCSVDataset` class. After it, we must create a LightningDataModule, that will define the data loaders for training, validation and test. ` to create a `LightningDataModule`, easily. \n", + "The `train_dataloader` method will use `train.csv`, `val_dataloader` will use `validation.csv` and `test_dataloader` will use `test.csv`.\n", + "\n", + "To create a `MultiModalHARSeriesDataModule`, we must pass:\n", + "\n", + "- `data_path`: the path to the `KuHar` folder;\n", + "- `feature_prefixes`: the prefixes of the features in the CSV files. In this case, we have `accel-x`, `accel-y`, `accel-z`, `gyro-x`, `gyro-y` and `gyro-z`;\n", + "- `batch_size`: the batch size for the data loaders; and\n", + "- `num_workers`: the number of workers for the data loaders. Essentially, the number of parallel processes to load the data.\n", + "\n", + "All data loader will share the passed parameters, such as `batch_size`, `num_workers`, and `feature_prefixes`." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/usr/local/lib/python3.10/dist-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n", + " from .autonotebook import tqdm as notebook_tqdm\n" + ] + }, + { + "data": { + "text/plain": [ + "" + ] + }, + "execution_count": 1, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "from ssl_tools.data.data_modules.har import MultiModalHARSeriesDataModule\n", + "\n", + "data_path = \"/workspaces/hiaac-m4/ssl_tools/data/standartized_balanced/KuHar/\"\n", + "\n", + "data_module = MultiModalHARSeriesDataModule(\n", + " data_path=data_path,\n", + " feature_prefixes=(\"accel-x\", \"accel-y\", \"accel-z\", \"gyro-x\", \"gyro-y\", \"gyro-z\"),\n", + " label=\"standard activity code\",\n", + " features_as_channels=True,\n", + " batch_size=64,\n", + " num_workers=0, # Sequential, for notebook compatibility\n", + ")\n", + "data_module" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can test the dataloaders by getting the first batch of each one. Let's do it, but just for the `train_dataloader`. Note that the `.setup()` method must be called before getting the data loaders. If you don't call it, the data loaders will not be created. However, when used to train a model, the Pytorch Lightning `.fit()` method will call the `.setup()` method for you. So, we put it here just to show how to use it." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Inputs shape: torch.Size([64, 6, 60]), Targets shape: torch.Size([64])\n" + ] + } + ], + "source": [ + "data_module.setup(\"fit\") # We just put it here to test.\n", + " # When training a model, the Trainer will call this method.\n", + "train_dataloader = data_module.train_dataloader()\n", + "\n", + "# Pick the first batch to inspect. The batch size is 64, so we have 64 samples.\n", + "batch = next(iter(train_dataloader))\n", + "# Each batch is a 2-element tuple with the first element being the 64 sample input and the second the 64 sample target.\n", + "inputs, targets = batch\n", + "\n", + "print(f\"Inputs shape: {inputs.shape}, Targets shape: {targets.shape}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Training a simple model\n", + "\n", + "We will create a simple 1D CNN Pytorch Lightning model using the `Simple1DConvNetwork`. The model will be trained to classify the activities in the KuHar dataset. \n", + "\n", + "Pytorch Lightning models must implement the `forward` method, `training_step` and `configure_optimizers` methods. Also, the `__init__` method is used to define the model.\n", + "The `forward` method is the same as the Pytorch `forward` method. \n", + "The `training_step` method is the method that will be called for each batch of data during the training. \n", + "The `configure_optimizers` method is the method that will define the optimizer to be used during the training.\n", + "\n", + "The `Simple1DConvNetwork` is a simple 1D CNN model that will be used to classify the activities in the KuHar dataset. It has 3 convolutional layers and 2 fully connected layers. It is trained using the `Adam` optimizer and the `CrossEntropyLoss` loss function.\n", + "\n", + "Besides that, Lightning models implemented in this framework, usually logs the training and validation losses.\n", + "Also, the `test` usually implement common metrics, such as accuracy." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "Simple1DConvNetwork(\n", + " (loss_func): CrossEntropyLoss()\n", + " (features): Sequential(\n", + " (0): Conv1d(6, 64, kernel_size=(5,), stride=(1,))\n", + " (1): ReLU()\n", + " (2): Dropout(p=0.5, inplace=False)\n", + " (3): Conv1d(64, 64, kernel_size=(5,), stride=(1,))\n", + " (4): ReLU()\n", + " (5): Dropout(p=0.5, inplace=False)\n", + " (6): Conv1d(64, 64, kernel_size=(5,), stride=(1,))\n", + " (7): ReLU()\n", + " )\n", + " (classifier): Sequential(\n", + " (0): Dropout(p=0.5, inplace=False)\n", + " (1): Linear(in_features=3072, out_features=128, bias=True)\n", + " (2): ReLU()\n", + " (3): Dropout(p=0.5, inplace=False)\n", + " (4): Linear(in_features=128, out_features=6, bias=True)\n", + " )\n", + ")" + ] + }, + "execution_count": 3, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "from ssl_tools.models.nets.convnet import Simple1DConvNetwork\n", + "\n", + "model = Simple1DConvNetwork(\n", + " input_channels=6, # The number of input channels (accel-x, accel-y, accel-z, gyro-x, gyro-y, gyro-z)\n", + " num_classes=6, # The number of output classes\n", + " time_steps=60, # Used to automatically calculate the input size of the linear layer\n", + " learning_rate=1e-3, # The learning rate for the optimizer\n", + ")\n", + "\n", + "model" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "To train a Lightning model using Pytorch Lightning, we must create a `Trainer` and call the `fit` method. The `Trainer` is responsible for training the model. It has several parameters, such as the number of epochs, the number of GPUs to use, the number of TPU cores to use, etc. \n", + "\n", + "We will train our model using the already defined dataloader. The `fit` method will be responsible for training the model using the training and validation data loaders. After the training, we will test the model using the test data loader.\n", + "\n", + "The training will run for 300 epochs (`max_epochs`) and will use 1 (`devices`) GPU only (`accelerator`)." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "GPU available: True (cuda), used: True\n", + "TPU available: False, using: 0 TPU cores\n", + "IPU available: False, using: 0 IPUs\n", + "HPU available: False, using: 0 HPUs\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1]\n", + "\n", + " | Name | Type | Params\n", + "------------------------------------------------\n", + "0 | loss_func | CrossEntropyLoss | 0 \n", + "1 | features | Sequential | 43.1 K\n", + "2 | classifier | Sequential | 394 K \n", + "------------------------------------------------\n", + "437 K Trainable params\n", + "0 Non-trainable params\n", + "437 K Total params\n", + "1.749 Total estimated model params size (MB)\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Sanity Checking DataLoader 0: 0%| | 0/2 [00:00┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n", + "┃ Test metric ┃ DataLoader 0 ┃\n", + "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n", + "│ test_acc │ 0.9027777910232544 │\n", + "│ test_loss │ 0.626140832901001 │\n", + "└───────────────────────────┴───────────────────────────┘\n", + "\n" + ], + "text/plain": [ + "┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n", + "┃\u001b[1m \u001b[0m\u001b[1m Test metric \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1m DataLoader 0 \u001b[0m\u001b[1m \u001b[0m┃\n", + "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n", + "│\u001b[36m \u001b[0m\u001b[36m test_acc \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m 0.9027777910232544 \u001b[0m\u001b[35m \u001b[0m│\n", + "│\u001b[36m \u001b[0m\u001b[36m test_loss \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m 0.626140832901001 \u001b[0m\u001b[35m \u001b[0m│\n", + "└───────────────────────────┴───────────────────────────┘\n" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/plain": [ + "[{'test_loss': 0.626140832901001, 'test_acc': 0.9027777910232544}]" + ] + }, + "execution_count": 5, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "trainer.test(model, data_module)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "And if we want to test the model using the validation data loader, we also can use the `trainer.test` method, but passing the `val_dataloader`. " + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1]\n", + "/usr/local/lib/python3.10/dist-packages/lightning/pytorch/trainer/connectors/data_connector.py:441: The 'test_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=47` in the `DataLoader` to improve performance.\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Testing DataLoader 0: 100%|██████████| 7/7 [00:00<00:00, 167.33it/s]\n" + ] + }, + { + "data": { + "text/html": [ + "
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
+       "┃        Test metric        ┃       DataLoader 0        ┃\n",
+       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
+       "│         test_acc          │    0.5680751204490662     │\n",
+       "│         test_loss         │    13.804328918457031     │\n",
+       "└───────────────────────────┴───────────────────────────┘\n",
+       "
\n" + ], + "text/plain": [ + "┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n", + "┃\u001b[1m \u001b[0m\u001b[1m Test metric \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1m DataLoader 0 \u001b[0m\u001b[1m \u001b[0m┃\n", + "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n", + "│\u001b[36m \u001b[0m\u001b[36m test_acc \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m 0.5680751204490662 \u001b[0m\u001b[35m \u001b[0m│\n", + "│\u001b[36m \u001b[0m\u001b[36m test_loss \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m 13.804328918457031 \u001b[0m\u001b[35m \u001b[0m│\n", + "└───────────────────────────┴───────────────────────────┘\n" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/plain": [ + "[{'test_loss': 13.804328918457031, 'test_acc': 0.5680751204490662}]" + ] + }, + "execution_count": 6, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "data_module.setup(\"fit\")\n", + "trainer.test(model, data_module.val_dataloader())" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.6" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/notebooks/03_training_ssl_model.ipynb b/notebooks/03_training_ssl_model.ipynb new file mode 100644 index 0000000..88886ea --- /dev/null +++ b/notebooks/03_training_ssl_model.ipynb @@ -0,0 +1,1024 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# 3. Training a self-supervised model (CPC)\n", + "\n", + "In this notebook, we will train a self-supervised model using the Contrastive Predictive Coding (CPC) method. \n", + "This method is based on the idea of predicting future tokens in a sequence, and it has been shown to be very effective in learning useful representations for downstream tasks.\n", + "This framework already provides an implementation of CPC, so we will use it to train the model.\n", + "\n", + "We will pre-train the model using KuHar dataset, and then we will use the learned representations to train a classifier for the downstream task. \n", + "For both stages of training, as the last notebook, we will:\n", + "\n", + "1. Create a `Dataset` and then `LightningDataModule` to load the data;\n", + "2. Instantiate the CPC model; and\n", + "3. Train the model using PyTorch Lightning.\n", + "\n", + "We can instantiate the model in two ways:\n", + "\n", + "1. Instantiate each element, such as the encoder, the autoregressive model, and the CPC model, and then pass them to the CPC model; or\n", + "2. Using builder methods to instantiate the model. In this case, we do not need to instantiate each element separately, but we can still customize the model by passing the desired parameters to the builder methods. This is the approach we will use in this notebook.\n", + "\n", + "In summary, the second approach encapsulates the first one, making it easier to use and it is more convenient for our purposes." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Pre-training the model\n", + "\n", + "We will pre-train the model using the KuHar dataset. CPC is a self-supervised method, so we do not need labels to train the model. However, CPC assumes that the input data is sequential, that is, an input is a sequence of time-steps comprising different acitivities. Thus, for HAR, usually, one sample (a multi-modal time-series) correspond to the whole time-series of a single user.\n", + "\n", + "### Creating the LightningDataModule\n", + "\n", + "Our dataset must be organized in the following way:\n", + "\n", + "```\n", + "data/\n", + " train/\n", + " user1.csv\n", + " user2.csv\n", + " ...\n", + " validation/\n", + " user4.csv\n", + " user5.csv\n", + " ...\n", + " test/\n", + " user6.csv\n", + " user7.csv\n", + " ...\n", + "```\n", + "\n", + "And the content of each file should be something like:\n", + "\n", + "| timestamp | accel-x | accel-y | accel-z | gyro-x | gyro-y | gyro-z | activity |\n", + "|-----------|---------|---------|---------|--------|--------|--------|-----------|\n", + "| 0 | 0.1 | 0.2 | 0.3 | 0.4 | 0.5 | 0.6 | 0 |\n", + "| 1 | 0.2 | 0.3 | 0.4 | 0.5 | 0.6 | 0.7 | 0 |\n", + "| ... | ... | ... | ... | ... | ... | ... | ... |\n", + "\n", + "Where `timestamp` is the time-stamp of the sample, `accel-x`, `accel-y`, `accel-z`, `gyro-x`, `gyro-y`, and `gyro-z` are the features of the sample, and `activity` is the label of the time-step.\n", + "\n", + "In this way, we should use the `SeriesFolderCSVDataset` to load the data.\n", + "This will create a `Dataset` for us, where each CSV file is a sample, and each row of the CSV file is a time-step, and the columns are the features.\n", + "\n", + "> **NOTE**: The samples may have different lengths, so, for this method, the `batch_size` must be 1.\n", + "\n", + "If your data is organized as above, where inside the root folder (`data/` in this case) there are sub-folders for each split (`train/`, `validation/`, and `test/`), and inside each split folder there are the CSV files, you can use the `UserActivityFolderDataModule` to create a `LightningDataModule` for you.\n", + "This class will create `DataLoader` of `SeriesFolderCSVDataset` for each split (train, validation, and test), and will setup data correctly.\n", + "\n", + "In this notebook, we will use the `UserActivityFolderDataModule` to create the `LightningDataModule` for us. This class minimally requires:\n", + "\n", + "- `data_path`: the root directory of the data;\n", + "- `features`: the name of the features columns;\n", + "- `pad`: a boolean indicating if the samples should be padded to the same length, that is, the length of the longest sample in the dataset. The padding scheme will replicate the samples, from the beginning, until the length of the longest sample is reached. " + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "UserActivityFolderDataModule(data_path=/workspaces/hiaac-m4/data/view_concatenated/KuHar_cpc, batch_size=1)" + ] + }, + "execution_count": 1, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "import numpy as np\n", + "import torch\n", + "from ssl_tools.data.data_modules import UserActivityFolderDataModule\n", + "\n", + "data_path = \"/workspaces/hiaac-m4/data/view_concatenated/KuHar_cpc\"\n", + "\n", + "data_module = UserActivityFolderDataModule(\n", + " data_path, \n", + " features=(\"accel-x\", \"accel-y\", \"accel-z\", \"gyro-x\", \"gyro-y\", \"gyro-z\"),\n", + " batch_size=1, # We set to 1 for CPC\n", + " label=None, # We do not want to return the labels, only data.\n", + " pad=False # If you want padded data, set it to True. \n", + " # This guarantees that all data have the same length. \n", + ")\n", + "\n", + "data_module" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Pre-training the model\n", + "\n", + "Here we will use the builder method `build_cpc` to instantiate the CPC model.\n", + "This will instantiate an CPC self-supervised model, with the default encoder (`ssl_tools.models.layers.gru.GRUEncoder`), that is an GRU+Linear, and the default autoregressive model (`torch.nn.GRU`), a linear layer.\n", + "\n", + "We can parametrize the creation of the model by passing the desired parameters to the builder method. The `build_cpc` method can be parametrized the following parameters:\n", + "\n", + "- `encoding_size`: the size of the encoded representation;\n", + "- `in_channels`: number of input features;\n", + "- `gru_hidden_size`: number of features in the hidden state of the GRU;\n", + "- `gru_num_layers`: number of layers in the GRU;\n", + "- `learning_rate`: the learning rate of the optimizer;\n", + "- `window_size` : size of the input windows (`X_t`) to be fed to the encoder (GRU).\n", + "\n", + "All parameters are optional, and have default values. You may want to consult the documentation of the method to see the default values and additional parameters.\n", + "\n", + "Note that the `LightningModule` returned by the `build_cpc` method is already configured to use the `CPC` loss, and the `Adam` optimizer." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "CPC(\n", + " (encoder): GRUEncoder(\n", + " (rnn): GRU(6, 100, bidirectional=True)\n", + " (nn): Linear(in_features=200, out_features=128, bias=True)\n", + " )\n", + " (density_estimator): Linear(in_features=128, out_features=128, bias=True)\n", + " (auto_regressor): GRU(128, 128, batch_first=True)\n", + " (loss_func): CrossEntropyLoss()\n", + ")" + ] + }, + "execution_count": 2, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "from ssl_tools.models.ssl.cpc import build_cpc\n", + "encoding_size = 128 \n", + "in_channels = 6\n", + "gru_hidden_size = 100\n", + "gru_num_layers = 1\n", + "learning_rate = 1e-3\n", + "\n", + "model = build_cpc(\n", + " encoding_size=encoding_size,\n", + " in_channels=in_channels,\n", + " gru_hidden_size=gru_hidden_size,\n", + " gru_num_layers=gru_num_layers,\n", + " learning_rate=learning_rate\n", + ")\n", + "model" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We instantiate the Trainer and call the `fit` method to train the model." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Trainer will use only 1 of 2 GPUs because it is running inside an interactive / notebook environment. You may try to set `Trainer(devices=2)` but please note that multi-GPU inside interactive / notebook environments is considered experimental and unstable. Your mileage may vary.\n", + "GPU available: True (cuda), used: True\n", + "TPU available: False, using: 0 TPU cores\n", + "IPU available: False, using: 0 IPUs\n", + "HPU available: False, using: 0 HPUs\n", + "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1]\n", + "\n", + " | Name | Type | Params\n", + "-------------------------------------------------------\n", + "0 | encoder | GRUEncoder | 90.5 K\n", + "1 | density_estimator | Linear | 16.5 K\n", + "2 | auto_regressor | GRU | 99.1 K\n", + "3 | loss_func | CrossEntropyLoss | 0 \n", + "-------------------------------------------------------\n", + "206 K Trainable params\n", + "0 Non-trainable params\n", + "206 K Total params\n", + "0.824 Total estimated model params size (MB)\n" + ] + }, + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "ae5e2aa76b724bd28d4a6a30a24741b6", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "Sanity Checking: | | 0/? [00:00 **NOTE**: It is important that the SSL models implement the `forward` method to return the latent representations of the input data, so we can use these representations to train the classifier. " + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Creating the LightningDataModule\n", + "\n", + "Human acivity recognition is a supervised classification task, that usually receives multi-modal windowed time-series as input, diferently from the self-supervised task, that receives the whole time-series of a single user.\n", + "Thus, we cannot use the same `LightningDataModule` to load the data for the downstream task. \n", + "\n", + "In this notebook, we will use the windowed time-series version of the KuHar dataset, that each split is a single CSV file, containing windowed time-series of the users. The content of the file should be something like:\n", + "\n", + "```\n", + "KuHar/\n", + " train.csv\n", + " validation.csv\n", + " test.csv\n", + "```\n", + "\n", + "The `train.csv` file may look like this:\n", + "\n", + "| accel-x-0 | accel-x-1 | accel-y-0 | accel-y-1 | class |\n", + "|-----------|-----------|-----------|-----------|--------|\n", + "| 0.502123 | 0.02123 | 0.502123 | 0.502123 | 0 |\n", + "| 0.6820123 | 0.02123 | 0.502123 | 0.502123 | 1 |\n", + "| 0.498217 | 0.00001 | 1.414141 | 3.141592 | 1 |\n", + "\n", + "As each CSV file contains time-windows signals of two 3-axis sensors (accelerometer and gyroscope), we must use the `MultiModalSeriesCSVDataset` class. \n", + "\n", + "As in last notebook, we will use the `MultiModalHARSeriesDataModule` to facilitate the creation of the `LightningDataModule`. This class will create `DataLoader` of `MultiModalSeriesCSVDataset` for each split (train, validation, and test), and will setup data correctly." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "MultiModalHARSeriesDataModule(data_path=/workspaces/hiaac-m4/ssl_tools/data/standartized_balanced/KuHar, batch_size=64)" + ] + }, + "execution_count": 4, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "from ssl_tools.data.data_modules.har import MultiModalHARSeriesDataModule\n", + "\n", + "data_path = \"/workspaces/hiaac-m4/ssl_tools/data/standartized_balanced/KuHar/\"\n", + "\n", + "data_module = MultiModalHARSeriesDataModule(\n", + " data_path=data_path,\n", + " feature_prefixes=(\"accel-x\", \"accel-y\", \"accel-z\", \"gyro-x\", \"gyro-y\", \"gyro-z\"),\n", + " label=\"standard activity code\",\n", + " features_as_channels=True,\n", + " batch_size=64,\n", + " num_workers=0, # Sequential, for notebook compatibility\n", + ")\n", + "data_module" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Fine-tuning the model\n", + "\n", + "A model for a downstream task is usually composed of two parts: the backbone model, that is the model that generates the representations of the input data, *i.e.*, the encoder, and the prediction head, which is the model that receives the representations and outputs the predictions, usually, a MLP.\n", + "\n", + "To handle the fine-tune process, we can design a new model, that is composed of the pre-trained backbone and the prediction head, and then train this new model with the labeled data. \n", + "In order to facilitate this process, this framework provides the `SSLDiscriminator` class, that receives the backbone model and the prediction head, and then trains the classifier with the labeled data.\n", + "\n", + "In summary, the `SSLDiscriminator` class is a `LightningModule` that generate the representations of the input data using the backbone model, that is, using the `forward` method of the backbone model, and then uses the prediction head to output the predictions. The predictions and labels are then used to compute the loss and train the model. \n", + "By default, the `SSLDiscriminator` is trained using the `Adam` optimizer with the `learning_rate` defined by the user (1e-3 by default).\n", + "\n", + "It worth to mention that the `SSLDiscriminator` class `forward` method receives the input data and the labels, and returns the predictions. This is different from the `forward` method of the self-supervised models, that receives only the input data and returns the latent representations of the input data.\n", + "\n", + "It worth to notice that the fine-tune train process can be done in two ways: \n", + "\n", + "1. Fine-tuning the whole model, that is, backbone (encoder) and classifier, with the labeled data; or \n", + "2. Fine-tuning only the classifier, with the labeled data.\n", + "The `SSLDisriminator` class can handle both cases, with the `update_backbone` parameter. If `update_backbone` is `True`, the whole model is fine-tuned (case 1, above), otherwise, only the classifier is fine-tuned (case 2, above).\n", + "\n", + "Let's create our prediction head and `SSLDisriminator` model and train it with the labeled data. Prediction heads for most popular tasks are already implemented in the `ssl_tools.models.ssl.modules.heads` module. In this notebook, we will use the `CPCPredictionHead` prediction head, that is a MLP with 3 hidden layers and dropout." + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "CPCPredictionHead(\n", + " (layers): Sequential(\n", + " (0): Linear(in_features=128, out_features=64, bias=True)\n", + " (1): ReLU()\n", + " (2): Linear(in_features=64, out_features=64, bias=True)\n", + " (3): Sequential(\n", + " (0): ReLU()\n", + " (1): Dropout(p=0, inplace=False)\n", + " )\n", + " (4): Linear(in_features=64, out_features=6, bias=True)\n", + " (5): Softmax(dim=1)\n", + " )\n", + ")" + ] + }, + "execution_count": 5, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "from ssl_tools.models.ssl.classifier import SSLDiscriminator\n", + "from ssl_tools.models.ssl.modules.heads import CPCPredictionHead\n", + "\n", + "number_of_classes = 6\n", + "\n", + "prediction_head = CPCPredictionHead(\n", + " input_dim=encoding_size, # Size of the encoding (input)\n", + " hidden_dim1=64,\n", + " hidden_dim2=64,\n", + " output_dim=number_of_classes # Number of classes\n", + ")\n", + "\n", + "prediction_head" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We will create the `SSLDisriminator` model. \n", + "The `SSLDisriminator` minimally requires:\n", + "\n", + "- `backbone`: the backbone model, that is, the pre-trained model;\n", + "- `head`: the prediction head model;\n", + "- `loss_fn`: the loss function to be used to train the model;\n", + "\n", + "Also, we can attach metrics that will be calculated with for every batch of `validation` and `test` sets. The metrics is passed using the `metrics` parameter of the `SSLDisriminator` class, that receives a dictionary with the name of the metric as key and the `torchmetrics.Metric` as value.\n", + "\n", + "Let's create the `SSLDiscriminator` and attach the `Accuracy` metric to the model, to check the validation accuracy per epoch." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "SSLDiscriminator(\n", + " (backbone): CPC(\n", + " (encoder): GRUEncoder(\n", + " (rnn): GRU(6, 100, bidirectional=True)\n", + " (nn): Linear(in_features=200, out_features=128, bias=True)\n", + " )\n", + " (density_estimator): Linear(in_features=128, out_features=128, bias=True)\n", + " (auto_regressor): GRU(128, 128, batch_first=True)\n", + " (loss_func): CrossEntropyLoss()\n", + " )\n", + " (head): CPCPredictionHead(\n", + " (layers): Sequential(\n", + " (0): Linear(in_features=128, out_features=64, bias=True)\n", + " (1): ReLU()\n", + " (2): Linear(in_features=64, out_features=64, bias=True)\n", + " (3): Sequential(\n", + " (0): ReLU()\n", + " (1): Dropout(p=0, inplace=False)\n", + " )\n", + " (4): Linear(in_features=64, out_features=6, bias=True)\n", + " (5): Softmax(dim=1)\n", + " )\n", + " )\n", + " (loss_fn): CrossEntropyLoss()\n", + ")" + ] + }, + "execution_count": 6, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "from torchmetrics import Accuracy\n", + "from torch.nn import CrossEntropyLoss\n", + "\n", + "acc_metric = Accuracy(\n", + " task=\"multiclass\", # We are working with a multiclass\n", + " # classification, not a binary one.\n", + " num_classes=number_of_classes # Number of classes\n", + ")\n", + "\n", + "ssl_discriminator = SSLDiscriminator(\n", + " backbone=model, # The model we trained before (CPC)\n", + " head=prediction_head, # The prediction head we just created\n", + " loss_fn=CrossEntropyLoss(), # The loss function\n", + " learning_rate=1e-3,\n", + " update_backbone=False, # We do not want to update the backbone\n", + " metrics={\"acc\": acc_metric} # We want to track the accuracy\n", + ")\n", + "ssl_discriminator " + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Then we can instantiate the Trainer and call the `fit` method to train the model." + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Trainer will use only 1 of 2 GPUs because it is running inside an interactive / notebook environment. You may try to set `Trainer(devices=2)` but please note that multi-GPU inside interactive / notebook environments is considered experimental and unstable. Your mileage may vary.\n", + "GPU available: True (cuda), used: True\n", + "TPU available: False, using: 0 TPU cores\n", + "IPU available: False, using: 0 IPUs\n", + "HPU available: False, using: 0 HPUs\n", + "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1]\n", + "\n", + " | Name | Type | Params\n", + "-----------------------------------------------\n", + "0 | backbone | CPC | 206 K \n", + "1 | head | CPCPredictionHead | 12.8 K\n", + "2 | loss_fn | CrossEntropyLoss | 0 \n", + "-----------------------------------------------\n", + "12.8 K Trainable params\n", + "206 K Non-trainable params\n", + "218 K Total params\n", + "0.876 Total estimated model params size (MB)\n" + ] + }, + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "2cf48079d50145c7afc3309882108058", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "Sanity Checking: | | 0/? [00:00┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n", + "┃ Test metric ┃ DataLoader 0 ┃\n", + "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n", + "│ test_acc │ 0.5277777910232544 │\n", + "│ test_loss │ 1.5032936334609985 │\n", + "└───────────────────────────┴───────────────────────────┘\n", + "\n" + ], + "text/plain": [ + "┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n", + "┃\u001b[1m \u001b[0m\u001b[1m Test metric \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1m DataLoader 0 \u001b[0m\u001b[1m \u001b[0m┃\n", + "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n", + "│\u001b[36m \u001b[0m\u001b[36m test_acc \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m 0.5277777910232544 \u001b[0m\u001b[35m \u001b[0m│\n", + "│\u001b[36m \u001b[0m\u001b[36m test_loss \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m 1.5032936334609985 \u001b[0m\u001b[35m \u001b[0m│\n", + "└───────────────────────────┴───────────────────────────┘\n" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/plain": [ + "[{'test_loss': 1.5032936334609985, 'test_acc': 0.5277777910232544}]" + ] + }, + "execution_count": 8, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "results = trainer.test(ssl_discriminator, data_module)\n", + "results" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finally, if we want to get the predictions of the model, we can:\n", + "\n", + "1. Call the `forward` method of the model, passing the input data (iterating over all batches of the dataloader); or \n", + "2. Use the `Trainer.predict` method, passing the data module. If you use the `Trainer.predict` method, the model will be set to evaluation mode, and the predictions will be done using the `predict_dataloader` defined in the `LightningDataModule`. This is usually the test set (`test_dataloader`)." + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1]\n", + "/usr/local/lib/python3.10/dist-packages/lightning/pytorch/trainer/connectors/data_connector.py:441: The 'predict_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=47` in the `DataLoader` to improve performance.\n" + ] + }, + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "c344339416884a01b84c3bc5dace1252", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "Predicting: | | 0/? [00:00┏━━━┳━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━┳━━━━━━━━┓\n", + "┃ ┃ Name ┃ Type ┃ Params ┃\n", + "┡━━━╇━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━╇━━━━━━━━┩\n", + "│ 0 │ encoder │ GRUEncoder │ 95.0 K │\n", + "│ 1 │ density_estimator │ Linear │ 22.7 K │\n", + "│ 2 │ auto_regressor │ GRU │ 135 K │\n", + "│ 3 │ loss_func │ CrossEntropyLoss │ 0 │\n", + "└───┴───────────────────┴──────────────────┴────────┘\n", + "\n" + ], + "text/plain": [ + "┏━━━┳━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━┳━━━━━━━━┓\n", + "┃\u001b[1;35m \u001b[0m\u001b[1;35m \u001b[0m\u001b[1;35m \u001b[0m┃\u001b[1;35m \u001b[0m\u001b[1;35mName \u001b[0m\u001b[1;35m \u001b[0m┃\u001b[1;35m \u001b[0m\u001b[1;35mType \u001b[0m\u001b[1;35m \u001b[0m┃\u001b[1;35m \u001b[0m\u001b[1;35mParams\u001b[0m\u001b[1;35m \u001b[0m┃\n", + "┡━━━╇━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━╇━━━━━━━━┩\n", + "│\u001b[2m \u001b[0m\u001b[2m0\u001b[0m\u001b[2m \u001b[0m│ encoder │ GRUEncoder │ 95.0 K │\n", + "│\u001b[2m \u001b[0m\u001b[2m1\u001b[0m\u001b[2m \u001b[0m│ density_estimator │ Linear │ 22.7 K │\n", + "│\u001b[2m \u001b[0m\u001b[2m2\u001b[0m\u001b[2m \u001b[0m│ auto_regressor │ GRU │ 135 K │\n", + "│\u001b[2m \u001b[0m\u001b[2m3\u001b[0m\u001b[2m \u001b[0m│ loss_func │ CrossEntropyLoss │ 0 │\n", + "└───┴───────────────────┴──────────────────┴────────┘\n" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "
Trainable params: 253 K                                                                                            \n",
+       "Non-trainable params: 0                                                                                            \n",
+       "Total params: 253 K                                                                                                \n",
+       "Total estimated model params size (MB): 1                                                                          \n",
+       "
\n" + ], + "text/plain": [ + "\u001b[1mTrainable params\u001b[0m: 253 K \n", + "\u001b[1mNon-trainable params\u001b[0m: 0 \n", + "\u001b[1mTotal params\u001b[0m: 253 K \n", + "\u001b[1mTotal estimated model params size (MB)\u001b[0m: 1 \n" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "d3d79b4d1870487a9225f9f9afdc262c", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "Output()" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "`Trainer.fit` stopped: `max_epochs=10` reached.\n" + ] + }, + { + "data": { + "text/html": [ + "
--> Overall fit time: 19.928 seconds\n",
+       "
\n" + ], + "text/plain": [ + "--> Overall fit time: 19.928 seconds\n" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "
\n"
+      ],
+      "text/plain": []
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "
\n",
+       "
\n"
+      ],
+      "text/plain": [
+       "\n"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Training finished\n",
+      "Last checkpoint saved at: logs/pretrain/CPC/2024-02-01_23-52-39/checkpoints/last.ckpt\n",
+      "Teardown experiment: CPC...\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Executing the experiment. Result is the output of the run() method\n",
+    "result = cpc_experiment.execute()  "
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Once the experiment finished, we may have a directory structure like this:\n",
+    "\n",
+    "```\n",
+    "logs/\n",
+    "    pretrain/\n",
+    "        CPC/\n",
+    "            2024-02-01_22-01-31/\n",
+    "                checkpoints/\n",
+    "                    epoch=9-step=570.ckpt\n",
+    "                    last.ckpt\n",
+    "                hparams.yaml\n",
+    "                metrics.csv\n",
+    "```\n",
+    "\n",
+    "This is the default directory structure for experiments, where the experiment directory is `logs/pretrain/CPC/2024-02-01_22-01-31/`. The `checkpoints directory` contains the saved checkpoints and inside it we may have a `last.ckpt` file which is the last checkpoint saved.\n",
+    "The `hparams.yaml` file contains the hyperparameters, and the `metrics.csv` file contains the metrics logged during training.\n",
+    "\n",
+    "\n",
+    "We can obtain the experiment's model, data module, logger, checkpoint directory, callbacks, trianer, and hyperparameters using the `cpc_experiment.model`, `cpc_experiment.data_module`, `cpc_experiment.logger`, `cpc_experiment.checkpoint_dir`, `cpc_experiment.callbacks`, `cpc_experiment.trainer`, and `cpc_experiment.hyperparameters` attributes, respectively. \n",
+    "These objects are cached in the `cpc_experiment` object, thus, it is instantiated only once, and can be accessed multiple times.\n",
+    "Also, the `cpc_experiment.finished` attribute is a boolean indicating if the experiment has finished sucessfuly or not.\n",
+    "\n",
+    "We will need this checkpoint to load the weights of the backbone for the finetuning process.\n",
+    "Let's  obtain the checkpoint file and the experiment's model and data module, and then run the finetuning experiment."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "PosixPath('logs/pretrain/CPC/2024-02-01_23-52-39/checkpoints/last.ckpt')"
+      ]
+     },
+     "execution_count": 3,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "backbone_checkpoint_path = cpc_experiment.checkpoint_dir / \"last.ckpt\"\n",
+    "backbone_checkpoint_path"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Experiment of Fine-tune CPC\n",
+    "\n",
+    "The `CPCTrain` class also encapsuate the default code for creating models and data modules from previous notebooks into the `get_finetune_model` and `get_finetune_data_module` methods. \n",
+    "The behaviour of these methods is similar to the `get_pretrain_model` and `get_pretrain_data_module` methods, but they are used to create the model and data module for the finetuning process.\n",
+    "In fact, the `get_finetune_model` will encapsulate the CPC code inside `SSLDisriminator` class, as seen in previous notebooks.\n",
+    "\n",
+    "As we use the same class for pretrain and finetune, we just need to set the `training_mode` attribute to `finetune` and set the `load_backbone` parameter to the checkpoint file obtained in the pretrain process. \n",
+    "Then, we can call the `execute` method to run the experiment.\n",
+    "\n",
+    "However, it worth to notice that fine tune is an supervised learning process and uses windowed time-series as input. Thus, the `data` parameter must be the path to a dataset where the samples are the windows of the time-series, as in previous notebooks. In our case, we will use the standardized balanced view of the KuHar dataset."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "LightningExperiment(experiment_dir=logs/finetune/CPC/2024-02-02_00-03-12, model=CPC, run_id=2024-02-02_00-03-12, finished=False)"
+      ]
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "data_path = \"/workspaces/hiaac-m4/ssl_tools/data/standartized_balanced/KuHar/\"\n",
+    "\n",
+    "cpc_experiment = CPCTrain(\n",
+    "    # General params\n",
+    "    training_mode=\"finetune\",\n",
+    "    load_backbone=backbone_checkpoint_path,\n",
+    "    # Data Module params\n",
+    "    data=data_path,\n",
+    "    # CPC model params\n",
+    "    encoding_size=150,\n",
+    "    window_size=60,\n",
+    "    in_channel=6,\n",
+    "    num_classes=6,\n",
+    "    # Trainer params\n",
+    "    epochs=10,\n",
+    "    num_workers=12,\n",
+    "    batch_size=128,\n",
+    "    accelerator=\"gpu\",\n",
+    "    devices=1,\n",
+    ")\n",
+    "\n",
+    "cpc_experiment"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/usr/local/lib/python3.10/dist-packages/lightning/fabric/loggers/csv_logs.py:198: Experiment logs directory logs/finetune/CPC/2024-02-02_00-03-12 exists and is not empty. Previous log files in this directory will be deleted when the new ones are saved!\n",
+      "GPU available: True (cuda), used: True\n",
+      "TPU available: False, using: 0 TPU cores\n",
+      "IPU available: False, using: 0 IPUs\n",
+      "HPU available: False, using: 0 HPUs\n",
+      "`Trainer(limit_train_batches=1.0)` was configured so 100% of the batches per epoch will be used..\n",
+      "`Trainer(limit_val_batches=1.0)` was configured so 100% of the batches will be used..\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Setting up experiment: CPC...\n",
+      "Running experiment: CPC...\n",
+      "Loading model from: logs/pretrain/CPC/2024-02-01_23-52-39/checkpoints/last.ckpt...\n",
+      "Model loaded successfully\n",
+      "Training will start\n",
+      "\tExperiment path: logs/finetune/CPC/2024-02-02_00-03-12\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1]\n"
+     ]
+    },
+    {
+     "data": {
+      "text/html": [
+       "
┏━━━┳━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━┳━━━━━━━━┓\n",
+       "┃   ┃ Name     ┃ Type              ┃ Params ┃\n",
+       "┡━━━╇━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━╇━━━━━━━━┩\n",
+       "│ 0 │ backbone │ CPC               │  253 K │\n",
+       "│ 1 │ head     │ CPCPredictionHead │ 14.2 K │\n",
+       "│ 2 │ loss_fn  │ CrossEntropyLoss  │      0 │\n",
+       "└───┴──────────┴───────────────────┴────────┘\n",
+       "
\n"
+      ],
+      "text/plain": [
+       "┏━━━┳━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━┳━━━━━━━━┓\n",
+       "┃\u001b[1;35m \u001b[0m\u001b[1;35m \u001b[0m\u001b[1;35m \u001b[0m┃\u001b[1;35m \u001b[0m\u001b[1;35mName    \u001b[0m\u001b[1;35m \u001b[0m┃\u001b[1;35m \u001b[0m\u001b[1;35mType             \u001b[0m\u001b[1;35m \u001b[0m┃\u001b[1;35m \u001b[0m\u001b[1;35mParams\u001b[0m\u001b[1;35m \u001b[0m┃\n",
+       "┡━━━╇━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━╇━━━━━━━━┩\n",
+       "│\u001b[2m \u001b[0m\u001b[2m0\u001b[0m\u001b[2m \u001b[0m│ backbone │ CPC               │  253 K │\n",
+       "│\u001b[2m \u001b[0m\u001b[2m1\u001b[0m\u001b[2m \u001b[0m│ head     │ CPCPredictionHead │ 14.2 K │\n",
+       "│\u001b[2m \u001b[0m\u001b[2m2\u001b[0m\u001b[2m \u001b[0m│ loss_fn  │ CrossEntropyLoss  │      0 │\n",
+       "└───┴──────────┴───────────────────┴────────┘\n"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "
Trainable params: 14.2 K                                                                                           \n",
+       "Non-trainable params: 253 K                                                                                        \n",
+       "Total params: 267 K                                                                                                \n",
+       "Total estimated model params size (MB): 1                                                                          \n",
+       "
\n"
+      ],
+      "text/plain": [
+       "\u001b[1mTrainable params\u001b[0m: 14.2 K                                                                                           \n",
+       "\u001b[1mNon-trainable params\u001b[0m: 253 K                                                                                        \n",
+       "\u001b[1mTotal params\u001b[0m: 267 K                                                                                                \n",
+       "\u001b[1mTotal estimated model params size (MB)\u001b[0m: 1                                                                          \n"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "f478640308374e5894df3302aa127e92",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Output()"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "
/usr/local/lib/python3.10/dist-packages/lightning/pytorch/loops/fit_loop.py:293: The number of training batches (11) is smaller than the logging interval Trainer(log_every_n_steps=50). Set a lower value for log_every_n_steps if you want to see logs for the training epoch.\n",
+       "
\n"
+      ],
+      "text/plain": [
+       "/usr/local/lib/python3.10/dist-packages/lightning/pytorch/loops/fit_loop.py:293: The number of training batches (11) is smaller than the logging interval Trainer(log_every_n_steps=50). Set a lower value for log_every_n_steps if you want to see logs for the training epoch.\n"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "# Executing the experiment. Result is the output of the run() method\n",
+    "result = cpc_experiment.execute()  "
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We will pick the last checkpoint from the fine-tuning process to evaluate the model."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "PosixPath('logs/finetune/CPC/2024-02-01_22-38-32/checkpoints/last.ckpt')"
+      ]
+     },
+     "execution_count": 6,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "fine_tuned_checkpoint_path = cpc_experiment.checkpoint_dir / \"last.ckpt\"\n",
+    "fine_tuned_checkpoint_path"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### CPC performance evaluation experiment\n",
+    "\n",
+    "Finally, we can evaluate the performance of the CPC model using the `CPCTest` class. This class inherits from `LightningTest` and encapsulate the default code for creating models and data modules from previous notebooks into the `get_model` and `get_data_module` methods.\n",
+    "\n",
+    "The signature of the `CPCTest` class is very similar to the `CPCTrain` class. Also, we will use the same data module used in the finetuning process. However, differently from the train process the test process uses the `.test` method in the trainer and not the `.fit` method.\n",
+    "Also, the `load` parameter is used to load the checkpoint obtained in the finetuning process (that load the weights from `SSLDiscriminator`, backbone and prediction haad).\n",
+    "\n",
+    "Let's create experiments to test the CPC model, using the test set from different datasets besides KuHAR."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/usr/local/lib/python3.10/dist-packages/lightning/fabric/loggers/csv_logs.py:198: Experiment logs directory logs/test/CPC/2024-02-01_23-01-24 exists and is not empty. Previous log files in this directory will be deleted when the new ones are saved!\n",
+      "GPU available: True (cuda), used: True\n",
+      "TPU available: False, using: 0 TPU cores\n",
+      "IPU available: False, using: 0 IPUs\n",
+      "HPU available: False, using: 0 HPUs\n",
+      "`Trainer(limit_test_batches=1.0)` was configured so 100% of the batches will be used..\n",
+      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Dataset at: /workspaces/hiaac-m4/ssl_tools/data/standartized_balanced/KuHar\n",
+      "Loading model from logs/finetune/CPC/2024-02-01_22-38-32/checkpoints/last.ckpt and executing test using dataset at KuHar...\n",
+      "Setting up experiment: CPC...\n",
+      "Running experiment: CPC...\n",
+      "Loading model from: logs/finetune/CPC/2024-02-01_22-38-32/checkpoints/last.ckpt...\n",
+      "Model loaded successfully\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "f580fe5d3fa9441a846ad0962953d3c3",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Output()"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
+       "┃        Test metric        ┃       DataLoader 0        ┃\n",
+       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
+       "│         test_acc          │    0.4583333432674408     │\n",
+       "│         test_loss         │     1.576676845550537     │\n",
+       "└───────────────────────────┴───────────────────────────┘\n",
+       "
\n"
+      ],
+      "text/plain": [
+       "┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
+       "┃\u001b[1m \u001b[0m\u001b[1m       Test metric       \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1m      DataLoader 0       \u001b[0m\u001b[1m \u001b[0m┃\n",
+       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
+       "│\u001b[36m \u001b[0m\u001b[36m        test_acc         \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m   0.4583333432674408    \u001b[0m\u001b[35m \u001b[0m│\n",
+       "│\u001b[36m \u001b[0m\u001b[36m        test_loss        \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m    1.576676845550537    \u001b[0m\u001b[35m \u001b[0m│\n",
+       "└───────────────────────────┴───────────────────────────┘\n"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "
\n"
+      ],
+      "text/plain": []
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "
\n",
+       "
\n"
+      ],
+      "text/plain": [
+       "\n"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/usr/local/lib/python3.10/dist-packages/lightning/fabric/loggers/csv_logs.py:198: Experiment logs directory logs/test/CPC/2024-02-01_23-01-28 exists and is not empty. Previous log files in this directory will be deleted when the new ones are saved!\n",
+      "GPU available: True (cuda), used: True\n",
+      "TPU available: False, using: 0 TPU cores\n",
+      "IPU available: False, using: 0 IPUs\n",
+      "HPU available: False, using: 0 HPUs\n",
+      "`Trainer(limit_test_batches=1.0)` was configured so 100% of the batches will be used..\n",
+      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Teardown experiment: CPC...\n",
+      "Test on dataset KuHar finished !\n",
+      "Dataset at: /workspaces/hiaac-m4/ssl_tools/data/standartized_balanced/MotionSense\n",
+      "Loading model from logs/finetune/CPC/2024-02-01_22-38-32/checkpoints/last.ckpt and executing test using dataset at MotionSense...\n",
+      "Setting up experiment: CPC...\n",
+      "Running experiment: CPC...\n",
+      "Loading model from: logs/finetune/CPC/2024-02-01_22-38-32/checkpoints/last.ckpt...\n",
+      "Model loaded successfully\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "483322b4b438455c8455b6417a625b5e",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Output()"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
+       "┃        Test metric        ┃       DataLoader 0        ┃\n",
+       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
+       "│         test_acc          │    0.3860640227794647     │\n",
+       "│         test_loss         │    1.6338088512420654     │\n",
+       "└───────────────────────────┴───────────────────────────┘\n",
+       "
\n"
+      ],
+      "text/plain": [
+       "┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
+       "┃\u001b[1m \u001b[0m\u001b[1m       Test metric       \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1m      DataLoader 0       \u001b[0m\u001b[1m \u001b[0m┃\n",
+       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
+       "│\u001b[36m \u001b[0m\u001b[36m        test_acc         \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m   0.3860640227794647    \u001b[0m\u001b[35m \u001b[0m│\n",
+       "│\u001b[36m \u001b[0m\u001b[36m        test_loss        \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m   1.6338088512420654    \u001b[0m\u001b[35m \u001b[0m│\n",
+       "└───────────────────────────┴───────────────────────────┘\n"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "
\n"
+      ],
+      "text/plain": []
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "
\n",
+       "
\n"
+      ],
+      "text/plain": [
+       "\n"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/usr/local/lib/python3.10/dist-packages/lightning/fabric/loggers/csv_logs.py:198: Experiment logs directory logs/test/CPC/2024-02-01_23-01-39 exists and is not empty. Previous log files in this directory will be deleted when the new ones are saved!\n",
+      "GPU available: True (cuda), used: True\n",
+      "TPU available: False, using: 0 TPU cores\n",
+      "IPU available: False, using: 0 IPUs\n",
+      "HPU available: False, using: 0 HPUs\n",
+      "`Trainer(limit_test_batches=1.0)` was configured so 100% of the batches will be used..\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Teardown experiment: CPC...\n",
+      "Test on dataset MotionSense finished !\n",
+      "Dataset at: /workspaces/hiaac-m4/ssl_tools/data/standartized_balanced/RealWorld_thigh\n",
+      "Loading model from logs/finetune/CPC/2024-02-01_22-38-32/checkpoints/last.ckpt and executing test using dataset at RealWorld_thigh...\n",
+      "Setting up experiment: CPC...\n",
+      "Running experiment: CPC...\n",
+      "Loading model from: logs/finetune/CPC/2024-02-01_22-38-32/checkpoints/last.ckpt...\n",
+      "Model loaded successfully\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1]\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "bc9753a4e8d549daad98bc8a70185c97",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Output()"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "from pathlib import Path\n",
+    "from ssl_tools.experiments.har_classification.cpc import CPCTest\n",
+    "\n",
+    "root_datasets_path = Path(\"/workspaces/hiaac-m4/ssl_tools/data/standartized_balanced/\")\n",
+    "\n",
+    "datasets = [\n",
+    "    \"KuHar\",\n",
+    "    \"MotionSense\",\n",
+    "    \"RealWorld_thigh\",\n",
+    "    \"RealWorld_waist\",\n",
+    "    \"UCI\"\n",
+    "    \"WISDM\"\n",
+    "]\n",
+    "\n",
+    "results = dict()\n",
+    "for dataset in datasets:\n",
+    "    data_path = root_datasets_path / dataset\n",
+    "    print(f\"Dataset at: {data_path}\")\n",
+    "    cpc_experiment = CPCTest(\n",
+    "        # General params\n",
+    "        load=fine_tuned_checkpoint_path,\n",
+    "        # Data Module params\n",
+    "        data=data_path,\n",
+    "        # CPC model params\n",
+    "        encoding_size=150,\n",
+    "        window_size=60,\n",
+    "        in_channel=6,\n",
+    "        num_classes=6,\n",
+    "        # Trainer params\n",
+    "        accelerator=\"gpu\",\n",
+    "        devices=1,\n",
+    "    )\n",
+    "    print(f\"Loading model from {fine_tuned_checkpoint_path} and executing test using dataset at {dataset}...\")\n",
+    "    results[dataset] = cpc_experiment.execute()\n",
+    "    print(f\"Test on dataset {dataset} finished!\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Other advantages of using `LightningExperiment`\n",
+    "\n",
+    "The `LightningExperiment` class also provides other advantages, such as:\n",
+    "\n",
+    "* Automatically generate CLI applications for the experiments, using the `jsonargparse` library. In fact, every parameter in the class constructor is automatically converted to a command line argument. This allows the user to run the experiment from the command line, using the same parameters as in the class constructor.\n",
+    "* Default `metrics.csv` and `hparams.yaml` files are created, and the hyperparameters are logged in the `hparams.yaml` file. The `metrics.csv` file contains the metrics logged during training, and can be used to analyze the performance of the model."
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.6"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/notebooks/experiment_classes.pdf b/notebooks/experiment_classes.pdf
new file mode 100644
index 0000000..73a79b2
Binary files /dev/null and b/notebooks/experiment_classes.pdf differ
diff --git a/notebooks/experiment_classes.svg b/notebooks/experiment_classes.svg
new file mode 100644
index 0000000..7ad76d5
--- /dev/null
+++ b/notebooks/experiment_classes.svg
@@ -0,0 +1,146 @@
+
+
+
+
+
+
+classes
+
+
+
+ssl_tools.experiments.experiment.Experiment
+
+Experiment
+
+experiment_dir
+log_dir : str
+name : str
+run_id
+seed : Optional[int]
+
+execute()
+run
+(): Any
+setup
+()
+teardown
+()
+
+
+
+ssl_tools.experiments.lightning_experiment.LightningExperiment
+
+LightningExperiment
+
+accelerator : str
+batch_size : int
+callbacks
+checkpoint_dir
+data_module
+devices : int
+experiment_dir
+finished
+hyperparameters
+load : Optional[str]
+log_every_n_steps : int
+logger
+model
+num_nodes : int
+num_workers : Optional[int]
+stage_name : str
+strategy : str
+trainer
+
+get_callbacks(): List[L.Callback]
+get_data_module
+(): L.LightningDataModule
+get_logger(): Logger
+get_model
+(): L.LightningModule
+get_trainer
+(logger: Logger, callbacks: List[L.Callback]): L.Trainer
+load_checkpoint(model: L.LightningModule, path: Path): L.LightningModule
+log_hyperparams(logger: Logger): dict
+run()
+run_model
+(model: L.LightningModule, data_module: L.LightningDataModule, trainer: L.Trainer)
+setup()
+
+
+
+ssl_tools.experiments.lightning_experiment.LightningExperiment->ssl_tools.experiments.experiment.Experiment
+
+
+
+
+
+ssl_tools.experiments.lightning_experiment.LightningSSLTrain
+
+LightningSSLTrain
+
+load_backbone : Optional[str]
+training_mode : str
+
+get_data_module(): L.LightningDataModule
+get_finetune_data_module
+(): L.LightningDataModule
+get_finetune_model
+(load_backbone: str): L.LightningModule
+get_model(): L.LightningModule
+get_pretrain_data_module
+(): L.LightningDataModule
+get_pretrain_model
+(): L.LightningModule
+
+
+
+ssl_tools.experiments.lightning_experiment.LightningTrain
+
+LightningTrain
+
+checkpoint_metric : Optional[str]
+checkpoint_metric_mode : str
+epochs : int
+learning_rate : float
+limit_train_batches : Union[float, int]
+limit_val_batches : Union[float, int]
+
+get_callbacks(): List[L.Callback]
+get_trainer(logger: Logger, callbacks: List[L.Callback]): L.Trainer
+run_model(model: L.LightningModule, data_module: L.LightningDataModule, trainer: L.Trainer)
+
+
+
+ssl_tools.experiments.lightning_experiment.LightningSSLTrain->ssl_tools.experiments.lightning_experiment.LightningTrain
+
+
+
+
+
+ssl_tools.experiments.lightning_experiment.LightningTest
+
+LightningTest
+
+limit_test_batches : Union[float, int]
+
+get_callbacks(): List[L.Callback]
+get_trainer(logger: Logger, callbacks: List[L.Callback]): L.Trainer
+run_model(model: L.LightningModule, data_module: L.LightningDataModule, trainer: L.Trainer): Any
+
+
+
+ssl_tools.experiments.lightning_experiment.LightningTest->ssl_tools.experiments.lightning_experiment.LightningExperiment
+
+
+
+
+
+ssl_tools.experiments.lightning_experiment.LightningTrain->ssl_tools.experiments.lightning_experiment.LightningExperiment
+
+
+
+
+
diff --git a/ssl_tools/data/data_modules/har.py b/ssl_tools/data/data_modules/har.py
index bfc12c0..755a71a 100644
--- a/ssl_tools/data/data_modules/har.py
+++ b/ssl_tools/data/data_modules/har.py
@@ -213,6 +213,9 @@ def _load_dataset(self, split_name: str) -> SeriesFolderCSVDataset:
             "test",
             "predict",
         ], f"Invalid split_name: {split_name}"
+        
+        if split_name == "predict":
+            split_name = "test"
 
         return SeriesFolderCSVDataset(
             self.data_path / split_name,
@@ -288,6 +291,12 @@ def test_dataloader(self) -> DataLoader:
 
     def predict_dataloader(self) -> DataLoader:
         return self._get_loader("predict", shuffle=False)
+    
+    def __str__(self):
+        return f"UserActivityFolderDataModule(data_path={self.data_path}, batch_size={self.batch_size})"
+    
+    def __repr__(self) -> str:
+        return str(self)
 
 
 class TNCHARDataModule(UserActivityFolderDataModule):
@@ -556,6 +565,9 @@ def _load_dataset(self, split_name: str) -> MultiModalSeriesCSVDataset:
             "test",
             "predict",
         ], f"Invalid split_name: {split_name}"
+        
+        if split_name == "predict":
+            split_name = "test"
 
         return MultiModalSeriesCSVDataset(
             self.data_path / f"{split_name}.csv",
@@ -631,6 +643,12 @@ def test_dataloader(self) -> DataLoader:
 
     def predict_dataloader(self) -> DataLoader:
         return self._get_loader("predict", shuffle=False)
+    
+    def __str__(self):
+        return f"MultiModalHARSeriesDataModule(data_path={self.data_path}, batch_size={self.batch_size})"
+    
+    def __repr__(self) -> str:
+        return str(self)
 
 
 class TFCDataModule(L.LightningDataModule):
@@ -802,6 +820,9 @@ def _load_dataset(self, split_name: str) -> TFCDataset:
             "test",
             "predict",
         ], f"Invalid split_name: {split_name}"
+        
+        if split_name == "predict":
+            split_name = "test"
 
         path = self.data_path / f"{split_name}.csv"
         
diff --git a/ssl_tools/experiments/har_classification/cpc.py b/ssl_tools/experiments/har_classification/cpc.py
index 107e4a8..23f69c6 100755
--- a/ssl_tools/experiments/har_classification/cpc.py
+++ b/ssl_tools/experiments/har_classification/cpc.py
@@ -60,7 +60,7 @@ def __init__(
     def get_pretrain_model(self) -> L.LightningModule:
         model = build_cpc(
             encoding_size=self.encoding_size,
-            in_channel=self.in_channel,
+            in_channels=self.in_channel,
             learning_rate=self.learning_rate,
             window_size=self.window_size,
             n_size=5,
@@ -151,7 +151,7 @@ def __init__(
     def get_model(self, load_backbone: str = None) -> L.LightningModule:
         model = build_cpc(
             encoding_size=self.encoding_size,
-            in_channel=self.in_channel,
+            in_channels=self.in_channel,
             window_size=self.window_size,
             n_size=5,
         )
diff --git a/ssl_tools/models/layers/gru.py b/ssl_tools/models/layers/gru.py
index 16f8521..e0d40bb 100644
--- a/ssl_tools/models/layers/gru.py
+++ b/ssl_tools/models/layers/gru.py
@@ -5,7 +5,7 @@ class GRUEncoder(torch.nn.Module):
     def __init__(
         self,
         hidden_size: int = 100,
-        in_channel: int = 6,
+        in_channels: int = 6,
         encoding_size: int = 10,
         num_layers: int = 1,
         dropout: float = 0.0,
@@ -53,7 +53,7 @@ def __init__(
         
         # Parameters 
         self.hidden_size = hidden_size
-        self.in_channel = in_channel
+        self.in_channel = in_channels
         self.num_layers = num_layers
         self.encoding_size = encoding_size
         self.bidirectional = bidirectional
diff --git a/ssl_tools/models/nets/convnet.py b/ssl_tools/models/nets/convnet.py
new file mode 100644
index 0000000..1697378
--- /dev/null
+++ b/ssl_tools/models/nets/convnet.py
@@ -0,0 +1,313 @@
+from typing import Dict
+import torch
+import lightning as L
+from torchmetrics import Accuracy
+
+
+class Simple1DConvNetwork(L.LightningModule):
+    """Model for human-activity-recognition."""
+
+    def __init__(
+        self,
+        input_channels: int = 6,
+        num_classes: int = 6,
+        time_steps: int = 60,
+        learning_rate: float = 1e-3,
+    ):
+        super().__init__()
+        self.input_channels = input_channels
+        self.num_classes = num_classes
+        self.time_steps = time_steps
+        self.learning_rate = learning_rate
+        self.loss_func = torch.nn.CrossEntropyLoss()
+        self.metrics = {
+            "acc": Accuracy(task="multiclass", num_classes=num_classes)
+        }
+
+        # Extract features, 1D conv layers
+        self.features = torch.nn.Sequential(
+            torch.nn.Conv1d(input_channels, 64, 5),
+            torch.nn.ReLU(),
+            torch.nn.Dropout(),
+            torch.nn.Conv1d(64, 64, 5),
+            torch.nn.ReLU(),
+            torch.nn.Dropout(),
+            torch.nn.Conv1d(64, 64, 5),
+            torch.nn.ReLU(),
+        )
+        
+        self.fc_input_features = self._calculate_fc_input_features(input_channels)
+        
+        # Classify output, fully connected layers
+        self.classifier = torch.nn.Sequential(
+            torch.nn.Dropout(),
+            torch.nn.Linear(self.fc_input_features, 128),
+            torch.nn.ReLU(),
+            torch.nn.Dropout(),
+            torch.nn.Linear(128, num_classes),
+        )
+
+    def _calculate_fc_input_features(self, input_channels):
+        # Dummy input to get the output shape after conv2
+        x = torch.randn(1, input_channels, self.time_steps)
+        with torch.no_grad():
+            out = self.features(x)
+        # Return the total number of features
+        return out.view(out.size(0), -1).size(1)
+
+    def forward(self, x):
+        x = self.features(x)
+        x = x.view(x.size(0), self.fc_input_features)
+        out = self.classifier(x)
+        return out
+    
+    def loss_function(self, X, y):
+        loss = self.loss_func(X, y)
+        return loss
+
+    def _common_step(self, batch, batch_idx, prefix):
+        x, y = batch
+        y_hat = self.forward(x)
+        loss = torch.nn.functional.cross_entropy(y_hat, y)
+        self.log(
+            f"{prefix}_loss",
+            loss,
+            sync_dist=True,
+            on_epoch=True,
+            on_step=False,
+            prog_bar=True,
+            logger=True,
+        )
+        return loss, y_hat, y
+
+    def _compute_metrics(
+        self, y_hat: torch.Tensor, y: torch.Tensor, stage: str
+    ) -> Dict[str, float]:
+        """Compute the metrics.
+
+        Parameters
+        ----------
+        y_hat : torch.Tensor
+            The predictions of the model
+        y : _type_
+            The ground truth labels
+        stage : str
+            The stage of the training loop (train, val or test)
+
+        Returns
+        -------
+        Dict[str, float]
+            A dictionary containing the metrics. The keys are the names of the
+            metrics, and the values are the values of the metrics.
+        """
+        return {
+            f"{stage}_{metric_name}": metric.to(self.device)(y_hat, y)
+            for metric_name, metric in self.metrics.items()
+        }
+
+    def training_step(self, batch, batch_idx):
+        loss, y_hat, y = self._common_step(batch, batch_idx, prefix="train")
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+        loss, y_hat, y = self._common_step(batch, batch_idx, prefix="val")
+
+        if self.metrics is not None:
+            results = self._compute_metrics(y_hat, y, "val")
+            self.log_dict(
+                results,
+                on_step=False,
+                on_epoch=True,
+                prog_bar=True,
+                logger=True,
+            )
+
+        return loss
+
+    def test_step(self, batch, batch_idx):
+        loss, y_hat, y = self._common_step(batch, batch_idx, prefix="test")
+
+        if self.metrics is not None:
+            results = self._compute_metrics(y_hat, y, "test")
+            self.log_dict(
+                results,
+                on_step=False,
+                on_epoch=True,
+                prog_bar=True,
+                logger=True,
+            )
+
+        return loss
+
+    def predict_step(self, batch, batch_idx, dataloader_idx=None):
+        return self.forward(batch)
+
+    def configure_optimizers(self):
+        optimizer = torch.optim.Adam(
+            self.parameters(),
+            lr=self.learning_rate,
+        )
+        return optimizer
+
+
+
+# FROM https://github.com/jindongwang/Deep-learning-activity-recognition/blob/master/pytorch/network.py
+class Simple2DConvNetwork(L.LightningModule):
+    def __init__(
+        self,
+        input_channels: int = 10,
+        num_classes: int = 6,
+        time_steps: int = 60,
+        learning_rate: float = 1e-3,
+    ):
+        super().__init__()
+        self.input_channels = input_channels
+        self.num_classes = num_classes
+        self.time_steps = time_steps
+        self.loss_func = torch.nn.CrossEntropyLoss()
+        self.learning_rate = learning_rate
+        self.metrics = {
+            "acc": Accuracy(task="multiclass", num_classes=num_classes)
+        }
+
+        self.conv1 = torch.nn.Sequential(
+            torch.nn.Conv2d(
+                in_channels=input_channels,
+                out_channels=32,
+                kernel_size=(1, input_channels),
+            ),
+            torch.nn.ReLU(),
+            torch.nn.MaxPool2d(kernel_size=(1, 2), stride=2),
+        )
+        self.conv2 = torch.nn.Sequential(
+            torch.nn.Conv2d(
+                in_channels=32, out_channels=64, kernel_size=(1, input_channels)
+            ),
+            torch.nn.ReLU(),
+            torch.nn.MaxPool2d(kernel_size=(1, 2), stride=2),
+        )
+
+        self.fc_input_features = self._calculate_fc_input_features(
+            input_channels
+        )
+
+        self.fc1 = torch.nn.Sequential(
+            torch.nn.Linear(
+                in_features=self.fc_input_features, out_features=1000
+            ),
+            torch.nn.ReLU(),
+        )
+        self.fc2 = torch.nn.Sequential(
+            torch.nn.Linear(in_features=1000, out_features=500), torch.nn.ReLU()
+        )
+        self.fc3 = torch.nn.Sequential(
+            torch.nn.Linear(in_features=500, out_features=num_classes)
+        )
+
+    def _calculate_fc_input_features(self, input_channels):
+        # Dummy input to get the output shape after conv2
+        x = torch.randn(1, input_channels, 1, self.time_steps)
+        with torch.no_grad():
+            out = self.conv1(x)
+            out = self.conv2(out)
+        # Return the total number of features
+        return out.view(out.size(0), -1).size(1)
+
+    def loss_function(self, X, y):
+        loss = self.loss_func(X, y)
+        return loss
+
+    def forward(self, x):
+        out = self.conv1(x)
+        out = self.conv2(out)
+        out = out.view(out.size(0), -1)  # Flatten the output for fully
+        out = self.fc1(out)
+        out = self.fc2(out)
+        out = self.fc3(out)
+        return out
+
+    def _common_step(self, batch, batch_idx, prefix):
+        x, y = batch
+        if x.ndim == 3:
+            x = x.unsqueeze(2)
+        y_hat = self.forward(x)
+        loss = torch.nn.functional.cross_entropy(y_hat, y)
+        self.log(
+            f"{prefix}_loss",
+            loss,
+            sync_dist=True,
+            on_epoch=True,
+            on_step=False,
+            prog_bar=True,
+            logger=True,
+        )
+        return loss, y_hat, y
+
+    def _compute_metrics(
+        self, y_hat: torch.Tensor, y: torch.Tensor, stage: str
+    ) -> Dict[str, float]:
+        """Compute the metrics.
+
+        Parameters
+        ----------
+        y_hat : torch.Tensor
+            The predictions of the model
+        y : _type_
+            The ground truth labels
+        stage : str
+            The stage of the training loop (train, val or test)
+
+        Returns
+        -------
+        Dict[str, float]
+            A dictionary containing the metrics. The keys are the names of the
+            metrics, and the values are the values of the metrics.
+        """
+        return {
+            f"{stage}_{metric_name}": metric.to(self.device)(y_hat, y)
+            for metric_name, metric in self.metrics.items()
+        }
+
+    def training_step(self, batch, batch_idx):
+        loss, y_hat, y = self._common_step(batch, batch_idx, prefix="train")
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+        loss, y_hat, y = self._common_step(batch, batch_idx, prefix="val")
+
+        if self.metrics is not None:
+            results = self._compute_metrics(y_hat, y, "val")
+            self.log_dict(
+                results,
+                on_step=False,
+                on_epoch=True,
+                prog_bar=True,
+                logger=True,
+            )
+
+        return loss
+
+    def test_step(self, batch, batch_idx):
+        loss, y_hat, y = self._common_step(batch, batch_idx, prefix="test")
+
+        if self.metrics is not None:
+            results = self._compute_metrics(y_hat, y, "test")
+            self.log_dict(
+                results,
+                on_step=False,
+                on_epoch=True,
+                prog_bar=True,
+                logger=True,
+            )
+
+        return loss
+
+    def predict_step(self, batch, batch_idx, dataloader_idx=None):
+        return self.forward(batch)
+
+    def configure_optimizers(self):
+        optimizer = torch.optim.Adam(
+            self.parameters(),
+            lr=self.learning_rate,
+        )
+        return optimizer
diff --git a/ssl_tools/models/ssl/classifier.py b/ssl_tools/models/ssl/classifier.py
index f2656c1..4a14939 100644
--- a/ssl_tools/models/ssl/classifier.py
+++ b/ssl_tools/models/ssl/classifier.py
@@ -1,4 +1,4 @@
-from typing import Any, Callable, Dict
+from typing import Dict
 import lightning as L
 import torch
 from torchmetrics import Metric
@@ -13,6 +13,7 @@ def __init__(
         learning_rate: float = 1e-3,
         update_backbone: bool = True,
         metrics: Dict[str, Metric] = None,
+        optimizer_cls: torch.optim.Optimizer = None,
     ):
         """A generic SSL Discriminator model. It takes a backbone and a head
         and trains them jointly (or not, depending on the ``update_backbone``
@@ -50,6 +51,7 @@ def __init__(
         self.learning_rate = learning_rate
         self.update_backbone = update_backbone
         self.metrics = metrics
+        self.optimizer_cls = optimizer_cls or torch.optim.Adam
 
     def _loss_func(self, y_hat: torch.Tensor, y: torch.Tensor):
         """Calculates the loss function.
@@ -245,7 +247,7 @@ def predict_step(self, batch: torch.Tensor, batch_idx: int) -> torch.Tensor:
             The predictions of the model
         """
         x, y = batch
-        predictions = self.forwardf(x)
+        predictions = self.forward(x)
         return predictions
 
     def _freeze(self, model):
@@ -266,9 +268,9 @@ def configure_optimizers(self):
         only update the parameters of the head.
         """
         if self.update_backbone:
-            return torch.optim.Adam(self.parameters(), lr=self.learning_rate)
+            return self.optimizer_cls(self.parameters(), lr=self.learning_rate)
         else:
             self._freeze(self.backbone)
-            return torch.optim.Adam(
+            return self.optimizer_cls(
                 self.head.parameters(), lr=self.learning_rate
             )
diff --git a/ssl_tools/models/ssl/cpc.py b/ssl_tools/models/ssl/cpc.py
index 7f0a4a9..a7351ef 100644
--- a/ssl_tools/models/ssl/cpc.py
+++ b/ssl_tools/models/ssl/cpc.py
@@ -255,7 +255,7 @@ def get_config(self) -> dict:
 
 def build_cpc(
     encoding_size: int = 150,
-    in_channel: int = 6,
+    in_channels: int = 6,
     gru_hidden_size: int = 100,
     gru_num_layers: int = 1,
     gru_bidirectional: bool = True,
@@ -271,8 +271,7 @@ def build_cpc(
     Parameters
     ----------
     encoding_size : int, optional
-        Size of the encoding (output of the linear layer). This is the size of
-        the representation.
+        Size of the encoded representation (the output of the linear layer).
     in_channel : int, optional
         Number of input features (e.g. 6 for HAR data in MotionSense Dataset)
     gru_hidden_size : int, optional
@@ -303,7 +302,7 @@ def build_cpc(
     """
     encoder = GRUEncoder(
         hidden_size=gru_hidden_size,
-        in_channel=in_channel,
+        in_channels=in_channels,
         encoding_size=encoding_size,
         num_layers=gru_num_layers,
         dropout=dropout,
diff --git a/ssl_tools/models/ssl/tnc.py b/ssl_tools/models/ssl/tnc.py
index 90f3512..42b4b7f 100644
--- a/ssl_tools/models/ssl/tnc.py
+++ b/ssl_tools/models/ssl/tnc.py
@@ -264,7 +264,7 @@ def build_tnc(
 
     encoder = GRUEncoder(
         hidden_size=gru_hidden_size,
-        in_channel=in_channel,
+        in_channels=in_channel,
         encoding_size=encoding_size,
         num_layers=gru_num_layers,
         dropout=dropout,