Late to the Machine Learning Club

Artificial Intelligence and Machine Learning (AI/ML) have been buzzwords for a while now. Even though I appreciate the many gifts brought forth by Large Language Models (LLMs), I’ve mostly avoided the hype train and aimed not to make them the focus of my career. But as I mentioned in The Artifact Detection Chronicles, I was tasked with building my first model.

Let's be honest

I had read a book by John D. Kelleher called Deep Learning a while back, so I had a general idea of what neural networks were and how they learned. But when it came to actually implementing one in TensorFlow or PyTorch, I didn’t know where to start.

Luckily, I know how to navigate the internet and hold a productive conversation with an LLM until I have a solid game plan.

The model would need to take an epoch as input (not to be confused with an ML epoch). In this case, that’s about 500 time points for a 2-second segment at 250 Hz. The task: decide whether the segment contains an artifact or not.

Before touching any code, I decided to spend a full day just getting to know the data.

A CSV file.

The Data

I had around 300 .edf files with EEG data and channel labels, plus an equal number of .csv files containing annotation data (supposedly labeled by experts) marking where artifacts occur.

My first step was to see how the data was distributed. If I was going to classify segments as either “clean” or “raw,” the model would need to see a roughly equal amount of each.

Distribution of clean vs artifact EEG data.

I had about 92 hours of EEG data with a 60/40 split. That’s decent enough to move forward. Interestingly, the average artifact length was around 7 seconds, likely inflated by muscle movement or chewing tasks subjects were instructed to perform.

Next, I looked at the distribution of different artifact types.

Distribution of different artifact types.

The distribution wasn’t great. Some artifacts, like electrode pops, appeared over 3,000 times, while shivers appeared only 11 times. This imbalance made me realize I should stick to binary classification rather than train a multi-class artifact classifier.

Prepping

I wrote a script, generate_labeled_epochs.py, to cut the data into labeled 2-second tensors and place them into training, validation, and testing folders. Then I wrote another script, epoch_distribution.py, to check the dataset distribution.

Epoch distribution for training, validation, and testing.

Following advice from Jason in the DB club, I went with an 80/10/10 split. Now I had my first ML-dataset.

Architecture

After a few Discord DMs with Evan from the AI/ML club and some Googling, I settled on a basic architecture. My priority was getting a complete working pipeline before thinking about optimizing.

Basic 1D CNN architecture.

I chose a 1D CNN since it’s simple and works well for signal data like EEG or audio waves.

Firsts

My script train_val_test.py (yes, the all-in-one “coding crime,” SOC who?) handled the full pipeline. I also set up TensorBoard to track training.

I started with 10 ML epochs, a supposedly a safe baseline when you’re just checking if the model runs at all. I would regret this later.

TensorBoard confusion matrix.

An hour and 20 minutes later, with an ice pack under my 8 GB M1 MacBook, I got this confusion matrix:

The model was stuck predicting only one class. Since there was more clean data than raw, it just guessed “clean” every time and achieved a decent accuracy.

I balanced the dataset, but then it alternated between predicting all “clean” in one round and all “raw” in another. Like a kid trying to outsmart you.

Breaking News

My first model simply refused to learn. I kept tweaking architecture and hyperparameters. Nothing worked until FINALLY I doubled the training epochs to 20.

85% accuracy. A huge win for a first model built in a week.

Interestingly, the model consistently spent the first 10 epochs stuck on one answer, then somewhere between epochs 10 and 15 it would finally start descending the loss curve.

Time to retire

Change in performance with different epochs.

After some more runs, it was time to hand this off to Arsam for discussion.

Have you ever seen a plot of gradient descent? I have to keep looking for that global minimum.