Train a Convolutional Neural Network with Nvidia DIGITS and Caffe

In the previous part of this series, I introduced Nvidia DIGITS as a user-friendly interface to build Deep Learning models. In this tutorial, I will walk you through the steps involved in building a Convolutional Neural Network that can classify images. We will use the popular cat versus dog dataset to train our model. By the end of this tutorial, you will learn everything it takes to train the model with DIGITS and using it for inferencing.

You need a Linux machine with a GPU to build the model. For inferencing, any OS that can run Docker is sufficient.

Preparing the Environment

Let’s start by downloading the dataset on the host machine on which you will run Nvidia DIGITS container. The dataset is similar to the one used in Kaggle Dogs versus Cats competition. It contains about 4,000 images of cats and dogs which is just enough for the tutorial. The images are already resized to 256×256 to skip preprocessing.

Since the dataset is too big for Github, I uploaded it to Google Drive. Download the file and unzip it in a folder. You should see two directories, train and test. The first directory contains 2,500 images of each category while the test directory has about 1,500 images.

Let’s launch Nvidia DIGITS Docker container by mapping the folder that contains the train and test directories. The dataset would become available at /data directory within the container.

Training the Convolutional Neural Network

With the container up and running, we can access the web UI at port 5000.

Click on the Datasets tab, and choose Classification.

Point DIGITS to the train and test directories. Give a name to the dataset and click Create.

Within a few minutes, DIGITS will parse the directories to create three databases — train, val, and test.

Go back to the home page and select the Models tab. Under Images drop-down, select Classification.

Choose the dataset created in the previous step. Under Standard Networks, make sure you select AlexNet. Give a name to the model and click Create.

The above step kicks off the training job, which will take a few minutes to complete. You will notice that the accuracy increases with each epoch. Also, the learning rate gets adjusted with each epoch.

By the time it hits the 30th epoch, the model reaches an accuracy of 85 percent.

Let’s test our trained model with a cat and dog image from the web. Under Trained Models, browse for an image, and click on Classify One.

Our model should accurately classify dogs vs. cats.

Congratulations! You have successfully built a CNN model without writing a single line of code.

In the next step, we will use the trained model for inferencing.

Model Inferencing with Caffe

Let’s run the trained model on our local machine. Click on the Download Model button, which will start downloading the compressed model along with the weights.

Clone this repo which contains Python code for classification along with the shell script to run the Caffe Docker container for inference. It also contains a few sample images for testing the model.

Uncompress the downloaded model and copy all the files from the inference directory of the cloned repo. Your directory should look similar to the below screenshot.

Set the environment variable to the Caffe model file, and then run infer.sh script to perform inference.

This command will pull the Caffe Docker image for the CPU on the local machine.

If you try sending a cat image, you would see the below output:

Let’s take a closer look at the infer.sh

The script essentially maps the current directory with the model and weights inside the Caffe container. It then invokes classify.py with appropriate parameters such as the model name, weights, labels, and the image to be classified.

Since the container has the Caffe framework and all other dependencies, it can execute classify.py to run inference.

This tutorial covered the workflow involved in training a model through Nvidia DIGITS running on a Linux machine backed by GPU and using the same model on a Mac or Windows machine for inference.

Feature images of the cat and dog from via Pixabay.