MNIST Classification

This example will take you through training your own MNIST model using Unweave. By the end, you'll have used Unweave to:

  • Launch a zepl to train an MNIST classifier,
  • Upload data from your local laptop to your uwstore,
  • Download the trained MNIST model from the uwstore back to your local laptop.

Before we begin, you should follow the Getting Started guide to install the Unweave CLI and login for the first time.

Preparing your project

We're first going to clone the starters repository from the Unweave Github account and change into the MNIST starter template directory:

git clone https://github.com/unweave/starters.git
cd starters/mnist

We then need to initialize this directory using the unweave init command. This will create a new project, which we'll name mnist-starter, and link it to the mnist directory on our local machine.

โฏ unweave init
? Enter project name [mnist]: mnist-starter
Created project mnist-starter
Path:    	/Users/markwinter/Code/starters/mnist
Project: 	mnist-starter

With the directory initialized and a new project created, it's now time to get the MNIST training dataset. We've included a download.py Python script in the starter to automate this for you. The script will download it to the local path ./uwstore/data/MNIST.

python download.py

Next, we need to upload this data to the remote uwstore. The uwstore is like an S3 bucket mounted at the base of your repository. The data in your uwstore is available to your zepls at the path ./uwstore/data relative to your project root.

โฏ unweave store upload ./uwstore/data/MNIST
Uploading file MNIST/raw/t10k-images-idx3-ubyte Uploaded!
Uploading file MNIST/raw/t10k-images-idx3-ubyte.gz Uploaded!
Uploading file MNIST/raw/t10k-labels-idx1-ubyte Uploaded!
Uploading file MNIST/raw/t10k-labels-idx1-ubyte.gz Uploaded!
Uploading file MNIST/raw/train-images-idx3-ubyte Uploaded!
Uploading file MNIST/raw/train-images-idx3-ubyte.gz Uploaded!
Uploading file MNIST/raw/train-labels-idx1-ubyte Uploaded!
Uploading file MNIST/raw/train-labels-idx1-ubyte.gz Uploaded!

If you inspect the main.py training script, you'll see that we're loading the MNIST dataset from the local path instead of downloading the dataset from the Internet.

dataset1 = datasets.MNIST('./uwstore/data', train=True, download=False, transform=transform)
dataset2 = datasets.MNIST('./uwstore/data', train=False, download=False, transform=transform)

Train Your Model

Now that we have the project setup and the MNIST dataset in the uwstore, it's time to launch a zepl to train our MNIST classifier. You can do this by pre-pending unweave in-front of the command you would normally run to execute your python script. This will build the ZeplContext from your current directory, assign a serverless compute node, and run the zepl using the command you provide - in this case python main.py.

The main.py training script loads the MNIST dataset from ./uwstore/data and trains an MNIST model for 1 epoch. It saves the final model to the ./uwstore/output/mnist_cnn.pt path. Unweave will automatically sync anything written to the ./uwstore/output directory back to the uwstore once the zepl exits.

โฏ unweave python main.py
Created zepl "apart-smell-lead-writer" with ID "e0f58038-d41a-450c-b262-28fad97ccc88"

Starting zepl
Press Ctrl+C to cancel the zepl
Press Ctrl+D to detach

๐Ÿ”„ Initializing serverless compute node
โœ… Compute node assigned to zepl
๐Ÿ”„ Syncing Unweave Store
๐Ÿ”„ Building environment
๐Ÿš€ Zepl running. Fetching logs ...

Train Epoch: 1 [0/60000 (0%)]       Loss: 2.305400
Train Epoch: 1 [640/60000 (1%)]     Loss: 1.359780
Train Epoch: 1 [1280/60000 (2%)]    Loss: 0.830670
Train Epoch: 1 [1920/60000 (3%)]    Loss: 0.605963
Train Epoch: 1 [58240/60000 (97%)]	Loss: 0.020182
Train Epoch: 1 [58880/60000 (98%)]	Loss: 0.017681
Train Epoch: 1 [59520/60000 (99%)]	Loss: 0.002070

๐Ÿงน Cleaning up
โœ… Zepl complete

The main.py script has several arguments you can configure, the number of training epochs for example. You can pass in arguments to your script by appending a -- after the script followed by the desired arguments:

unweave python main.py -- --epochs=2 --lr=0.9

Once the Python script exits, we can list the uwstore to view the MNIST dataset that we uploaded earlier, and also the new MNIST model that we just trained. The MNIST model is saved under the zepl's unique ID in the output directory.

โฏ unweave store list 
2022-08-25 10:34:46 +0900 KST	7.8 MB	data/MNIST/raw/t10k-images-idx3-ubyte
2022-08-25 10:34:47 +0900 KST	1.6 MB	data/MNIST/raw/t10k-images-idx3-ubyte.gz
2022-08-25 10:34:48 +0900 KST	10 kB	data/MNIST/raw/t10k-labels-idx1-ubyte
2022-08-25 10:34:48 +0900 KST	4.5 kB	data/MNIST/raw/t10k-labels-idx1-ubyte.gz
2022-08-25 10:34:50 +0900 KST	47 MB	data/MNIST/raw/train-images-idx3-ubyte
2022-08-25 10:35:09 +0900 KST	9.9 MB	data/MNIST/raw/train-images-idx3-ubyte.gz
2022-08-25 10:35:10 +0900 KST	60 kB	data/MNIST/raw/train-labels-idx1-ubyte
2022-08-25 10:35:10 +0900 KST	29 kB	data/MNIST/raw/train-labels-idx1-ubyte.gz
2022-08-26 09:49:54 +0900 KST	4.8 MB	output/a5825ccc-9c83-4bb8-84a0-f16801ffaba1/mnist_cnn.pt

If you want to download the model (or any other file in the uwstore) back to your local computer, you can use the unweave store download command.

$ unweave store download output/a5825ccc-9c83-4bb8-84a0-f16801ffaba1/mnist_cnn.pt
mnist_cnn.pt saved to ./uwstore/output/000f422b-a293-4f9f-be22-ce1f66d00533/mnist_cnn.pt
Unweave Store