Using this Python library to send model training updates.

Imagine this scenario — you’re working on a deep learning project and just started a time-consuming training job on a GPU. Based on your estimates, it will take you about fifteen hours for your job to finish. Obviously, you don’t want to watch your model train for that long. But you still want to know when it finishes training while you’re away from your computer or working on a different task.

Recently, HuggingFace released a Python library called knockknock that allows developers to set up and receive notifications when their models are done training. In this article, I will demonstrate how you can use knockknock to receive model training updates on a wide range of platforms in only a few lines of code!

Installing KnockKnock

You can install knockknock easily with Pip using the command below.

pip install knockknock

Please keep in mind that this library has only been tested for Python 3.6 and later versions. If you have an earlier version of Python, I would suggest upgrading to Python 3.6 or higher if you want to use this library.

Training a Simple Neural Network

For the purpose of demonstrating this library, I will define a function that trains a simple CNN on the classic MNIST Handwritten Digits Dataset. To find the full code for this tutorial, check out this GitHub repository.

Import Libraries

In order to get started, we need to import a few modules from Keras.

from keras.utils.np_utils import to_categorical
from keras.models import Sequential
from keras.layers import Conv2D, MaxPooling2D, Flatten, Dense, Dropout, Activation
from keras.datasets import mnist

Load the Data

To keep this tutorial simple, I loaded the MNIST Dataset using Keras. I performed the following standard preprocessing steps as well in the code below:

1. Added an extra dimension to the 28 x 28-pixel training and testing images.
2. Scaled the training and testing data by 255.
3. Converted the numerical targets to categorical one-hot vectors.

(X_train, y_train), (X_test, y_test) = mnist.load_data()
X_train = X_train.reshape(X_train.shape[0], 28, 28,1) # adds extra dimension
X_test = X_test.reshape(X_test.shape[0], 28, 28, 1) # adds extra dimension
input_shape = (28, 28, 1)

X_train = X_train.astype('float32')
X_test = X_test.astype('float32')
X_train /= 255
X_test /= 255

y_train = to_categorical(y_train)
y_test = to_categorical(y_test)

Training the Model

The function below creates a simple CNN, trains it on the training dataset, and returns accuracy and loss values demonstrating the model’s performance on the test dataset.

https://gist.github.com/AmolMavuduru/6d579c67801276b0f9803b15cf474608

I won’t go into detail about the architecture of the neural network in the code above because the focus of this tutorial is on sending model training notifications. I purposely used a simple neural network for this reason.

Getting Desktop Notifications

Now that we have a function that trains a neural network, we can create another function that trains the model and creates desktop notifications when the function starts and finishes executing.

from knockknock import desktop_sender

@desktop_sender(title="Knockknock Desktop Notifier")
def train_model_desktop_notify(X_train, y_train, X_test, y_test):

    return train_model(X_train, 
                       y_train, 
                       X_test, 
                       y_test)

train_model_desktop_notify(X_train, y_train, X_test, y_test)

On my Mac, running the code above produces a popup notification when the training process starts.

Once the job finishes running, it generates another Desktop notification with the results of the function as shown below.

Getting Email Notifications

We can also set up email notifications with knockknock. In order to do this, you need to set up a separate Gmail account to send the email notifications. You’ll also need to allow less secure apps to access your Gmail account in the security settings.

from knockknock import email_sender

@email_sender(recipient_emails=["amolmavuduru@gmail.com"], 
              sender_email="knockknocknotificationstest@gmail.com")
def train_model_email_notify(X_train, y_train, X_test, y_test):

    return train_model(X_train, 
                       y_train, 
                       X_test, 
                       y_test)

train_model_email_notify(X_train, y_train, X_test, y_test)

Running the code produces the following email notification in Gmail.

When the function finishes running we get another notification email shown below.

Getting Slack Notifications

Finally, if you are part of a Slack team that is working on a machine learning project, you can also set up Slack notifications in a channel when your model finishes running. To do this, you need to create a Slack workspace, add a Slack app to it, and get a Slack webhook URL that you can supply to the slack_sender function decorator. Visit the Slack API page and then follow steps 1–3 in this tutorial to create a webhook.

The code below demonstrates how to create a Slack notification given a webhook URL and a specific channel to post in. Please note that I stored my webhook URL in an environment variable for security reasons.

from knockknock import slack_sender
import os

webhook_url = os.environ['KNOCKKNOCK_WEBHOOK']
@slack_sender(webhook_url=webhook_url, channel="#general")
def train_model_slack_notify(X_train, y_train, X_test, y_test):

    return train_model(X_train, 
                       y_train, 
                       X_test, 
                       y_test)

train_model_slack_notify(X_train, y_train, X_test, y_test)

Running the code above produces the following Slack notifications.

This is a great feature if your data science team has a Slack workspace and wants to monitor your model training jobs.

More Notification Options

With knockknock, we can also send notifications on the following platforms:

• Telegram
• Microsoft Teams
• Text Message
• Discord
• Matrix
• Amazon Chime
• DingTalk
• RocketChat
• WeChat Work

For detailed documentation on how to send notifications on these platforms, be sure to check out the project repository on GitHub.

Summary

Knockknock is a useful tool that lets you keep track of your model training jobs with notifications. You can use it to send notifications on a wide variety of platforms and it is great for keeping track of experiments in data science teams.

As usual, you can find all of the code in this article on GitHub.

Join My Mailing List

Do you want to get better at data science and machine learning? Do you want to stay up to date with the latest libraries, developments, and research in the data science and machine learning community?

Join my mailing list to get updates on my data science content. You’ll also get my free Step-By-Step Guide to Solving Machine Learning Problems when you sign up!

Sources

1. Hugging Face, knockknock Github repository, (2020).