feat: Add MNIST example (#115)

* feat: Add MNIST example

Signed-off-by: Ce Gao <cegao@tensorchord.ai>

* fix: Add a simple readme

Signed-off-by: Ce Gao <cegao@tensorchord.ai>

.gitignore

@@ -32,5 +32,6 @@ __debug_bin
 bin/
 debug-bin/
 /build.MIDI
+.ipynb_checkpoints/
 
 dist/

examples/mnist/README.md

@@ -0,0 +1,9 @@
+# MNIST Example
+
+## Quick start
+
+```bash
+$ midi up
+```
+
+Then you can open jupyter notebook at [`http://localhost:8888`](http://localhost:8888), or open vscode remote to attach to the container.

examples/mnist/build.MIDI

@@ -0,0 +1,12 @@
+vscode(plugins = [
+    "ms-python.python-2021.12.1559732655",
+])
+
+base(os="ubuntu20.04", language="python3")
+pip_package(name = [
+    "tensorflow",
+    "numpy",
+])
+cuda(version="11.6", cudnn="8")
+shell("zsh")
+jupyter(password="", port=8888)

examples/mnist/main.py

@@ -0,0 +1,79 @@
+"""
+Title: Simple MNIST convnet
+Author: [fchollet](https://twitter.com/fchollet)
+Date created: 2015/06/19
+Last modified: 2020/04/21
+Description: A simple convnet that achieves ~99% test accuracy on MNIST.
+"""
+
+"""
+## Setup
+"""
+
+import numpy as np
+from tensorflow import keras
+from tensorflow.keras import layers
+
+"""
+## Prepare the data
+"""
+
+# Model / data parameters
+num_classes = 10
+input_shape = (28, 28, 1)
+
+# the data, split between train and test sets
+(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
+
+# Scale images to the [0, 1] range
+x_train = x_train.astype("float32") / 255
+x_test = x_test.astype("float32") / 255
+# Make sure images have shape (28, 28, 1)
+x_train = np.expand_dims(x_train, -1)
+x_test = np.expand_dims(x_test, -1)
+print("x_train shape:", x_train.shape)
+print(x_train.shape[0], "train samples")
+print(x_test.shape[0], "test samples")
+
+
+# convert class vectors to binary class matrices
+y_train = keras.utils.to_categorical(y_train, num_classes)
+y_test = keras.utils.to_categorical(y_test, num_classes)
+
+"""
+## Build the model
+"""
+
+model = keras.Sequential(
+    [
+        keras.Input(shape=input_shape),
+        layers.Conv2D(32, kernel_size=(3, 3), activation="relu"),
+        layers.MaxPooling2D(pool_size=(2, 2)),
+        layers.Conv2D(64, kernel_size=(3, 3), activation="relu"),
+        layers.MaxPooling2D(pool_size=(2, 2)),
+        layers.Flatten(),
+        layers.Dropout(0.5),
+        layers.Dense(num_classes, activation="softmax"),
+    ]
+)
+
+model.summary()
+
+"""
+## Train the model
+"""
+
+batch_size = 128
+epochs = 15
+
+model.compile(loss="categorical_crossentropy", optimizer="adam", metrics=["accuracy"])
+
+model.fit(x_train, y_train, batch_size=batch_size, epochs=epochs, validation_split=0.1)
+
+"""
+## Evaluate the trained model
+"""
+
+score = model.evaluate(x_test, y_test, verbose=0)
+print("Test loss:", score[0])
+print("Test accuracy:", score[1])

examples/mnist/mnist.ipynb

@@ -0,0 +1,113 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "464baae6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\"\"\"\n",
+    "Title: Simple MNIST convnet\n",
+    "Author: [fchollet](https://twitter.com/fchollet)\n",
+    "Date created: 2015/06/19\n",
+    "Last modified: 2020/04/21\n",
+    "Description: A simple convnet that achieves ~99% test accuracy on MNIST.\n",
+    "\"\"\"\n",
+    "\n",
+    "\"\"\"\n",
+    "## Setup\n",
+    "\"\"\"\n",
+    "\n",
+    "import numpy as np\n",
+    "from tensorflow import keras\n",
+    "from tensorflow.keras import layers\n",
+    "\n",
+    "\"\"\"\n",
+    "## Prepare the data\n",
+    "\"\"\"\n",
+    "\n",
+    "# Model / data parameters\n",
+    "num_classes = 10\n",
+    "input_shape = (28, 28, 1)\n",
+    "\n",
+    "# the data, split between train and test sets\n",
+    "(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()\n",
+    "\n",
+    "# Scale images to the [0, 1] range\n",
+    "x_train = x_train.astype(\"float32\") / 255\n",
+    "x_test = x_test.astype(\"float32\") / 255\n",
+    "# Make sure images have shape (28, 28, 1)\n",
+    "x_train = np.expand_dims(x_train, -1)\n",
+    "x_test = np.expand_dims(x_test, -1)\n",
+    "print(\"x_train shape:\", x_train.shape)\n",
+    "print(x_train.shape[0], \"train samples\")\n",
+    "print(x_test.shape[0], \"test samples\")\n",
+    "\n",
+    "\n",
+    "# convert class vectors to binary class matrices\n",
+    "y_train = keras.utils.to_categorical(y_train, num_classes)\n",
+    "y_test = keras.utils.to_categorical(y_test, num_classes)\n",
+    "\n",
+    "\"\"\"\n",
+    "## Build the model\n",
+    "\"\"\"\n",
+    "\n",
+    "model = keras.Sequential(\n",
+    "    [\n",
+    "        keras.Input(shape=input_shape),\n",
+    "        layers.Conv2D(32, kernel_size=(3, 3), activation=\"relu\"),\n",
+    "        layers.MaxPooling2D(pool_size=(2, 2)),\n",
+    "        layers.Conv2D(64, kernel_size=(3, 3), activation=\"relu\"),\n",
+    "        layers.MaxPooling2D(pool_size=(2, 2)),\n",
+    "        layers.Flatten(),\n",
+    "        layers.Dropout(0.5),\n",
+    "        layers.Dense(num_classes, activation=\"softmax\"),\n",
+    "    ]\n",
+    ")\n",
+    "\n",
+    "model.summary()\n",
+    "\n",
+    "\"\"\"\n",
+    "## Train the model\n",
+    "\"\"\"\n",
+    "\n",
+    "batch_size = 128\n",
+    "epochs = 15\n",
+    "\n",
+    "model.compile(loss=\"categorical_crossentropy\", optimizer=\"adam\", metrics=[\"accuracy\"])\n",
+    "\n",
+    "model.fit(x_train, y_train, batch_size=batch_size, epochs=epochs, validation_split=0.1)\n",
+    "\n",
+    "\"\"\"\n",
+    "## Evaluate the trained model\n",
+    "\"\"\"\n",
+    "\n",
+    "score = model.evaluate(x_test, y_test, verbose=0)\n",
+    "print(\"Test loss:\", score[0])\n",
+    "print(\"Test accuracy:\", score[1])\n"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "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.8.10"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

pkg/docker/docker.go

@@ -256,7 +256,7 @@ func (c generalClient) StartMIDI(ctx context.Context, tag, name string,
 	}
 
 	if g.JupyterConfig != nil {
-		cmd, err := jupyter.GenerateCommand(*ir.DefaultGraph)
+		cmd, err := jupyter.GenerateCommand(*ir.DefaultGraph, config.WorkingDir)
 		if err != nil {
 			return "", "", errors.Wrap(err, "failed to generate jupyter command")
 		}

pkg/editor/jupyter/util.go

@@ -20,13 +20,14 @@ import (
 	"github.com/tensorchord/MIDI/pkg/lang/ir"
 )
 
-func GenerateCommand(g ir.Graph) ([]string, error) {
+func GenerateCommand(g ir.Graph, notebookDir string) ([]string, error) {
 	if g.JupyterConfig == nil {
 		return nil, nil
 	}
 
 	cmd := []string{
-		"jupyter", "notebook", "--allow-root", "--ip", "0.0.0.0",
+		"jupyter", "notebook", "--allow-root",
+		"--ip", "0.0.0.0", "--notebook-dir", notebookDir,
 	}
 	if g.JupyterConfig.Password != "" {
 		cmd = append(cmd, "--NotebookApp.password", g.JupyterConfig.Password,