Fixed devices and demo

.gitignore

@@ -6,3 +6,4 @@
 /downloads/*
 /checkpoints/*
 /checkpoints/
+/venv/

Dockerfile

@@ -0,0 +1,3 @@
+FROM python
+WORKDIR .
+RUN pip install -r requirements.txt

app.py

@@ -0,0 +1,3 @@
+from src.demo import main
+
+main()

main.py

@@ -0,0 +1,23 @@
+import torch
+
+from models import CNN
+from dataset import DatasetMNIST, download_mnist
+from train import get_dataloaders, train_net_manually, train_net_lightning
+
+
+def main(device):
+    mnist = download_mnist("downloads/mnist/")
+    dataset, test_data = DatasetMNIST(*mnist["train"]), DatasetMNIST(*mnist["test"])
+    train_loader, validate_loader, test_loader = get_dataloaders(dataset, test_data)
+
+    # Training manually
+    net = CNN(input_channels=1, num_classes=10).to(device)
+    optim = torch.optim.Adam(net.parameters(), lr=1e-4)
+    loss_fn = torch.nn.CrossEntropyLoss()
+    max_epochs = 1
+
+    train_net_manually(net, optim, loss_fn, train_loader, validate_loader, max_epochs, device)
+
+
+if __name__ == "__main__":
+    main("cpu")

notebooks/mlruns/0/meta.yaml

@@ -0,0 +1,6 @@
+artifact_location: mlflow-artifacts:/0
+creation_time: 1671993138419
+experiment_id: '0'
+last_update_time: 1671993138419
+lifecycle_stage: active
+name: Default

notebooks/optimizers.ipynb

@@ -0,0 +1,37 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "collapsed": true,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 2
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython2",
+   "version": "2.7.6"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 0
+}

requirements.txt

@@ -0,0 +1,69 @@
+aiohttp==3.8.3
+aiosignal==1.3.1
+altair==4.2.0
+anyio==3.6.2
+async-timeout==4.0.2
+attrs==22.2.0
+certifi==2022.12.7
+charset-normalizer==2.1.1
+click==8.1.3
+contourpy==1.0.6
+cycler==0.11.0
+entrypoints==0.4
+fastapi==0.88.0
+ffmpy==0.3.0
+fonttools==4.38.0
+frozenlist==1.3.3
+fsspec==2022.11.0
+gradio==3.15.0
+h11==0.14.0
+httpcore==0.16.3
+httpx==0.23.1
+idna==3.4
+Jinja2==3.1.2
+jsonschema==4.17.3
+kiwisolver==1.4.4
+lightning-utilities==0.5.0
+linkify-it-py==1.0.3
+markdown-it-py==2.1.0
+MarkupSafe==2.1.1
+matplotlib==3.6.2
+mdit-py-plugins==0.3.3
+mdurl==0.1.2
+multidict==6.0.4
+numpy==1.24.0
+orjson==3.8.3
+packaging==22.0
+pandas==1.5.2
+Pillow==9.3.0
+plotly==5.11.0
+protobuf==3.20.1
+pycryptodome==3.16.0
+pydantic==1.10.2
+pydub==0.25.1
+pyparsing==3.0.9
+pyrsistent==0.19.2
+python-dateutil==2.8.2
+python-multipart==0.0.5
+pytorch-lightning==1.8.6
+pytz==2022.7
+PyYAML==6.0
+requests==2.28.1
+rfc3986==1.5.0
+six==1.16.0
+sniffio==1.3.0
+starlette==0.22.0
+tenacity==8.1.0
+tensorboardX==2.5.1
+toolz==0.12.0
+torch==1.13.1
+torchaudio==0.13.1
+torchmetrics==0.11.0
+torchvision==0.14.1
+tqdm==4.64.1
+typing_extensions==4.4.0
+uc-micro-py==1.0.1
+urllib3==1.26.13
+uvicorn==0.20.0
+websockets==10.4
+yarl==1.8.2

src/dataset.py

@@ -9,7 +9,7 @@ import numpy as np
 from src.downloader import download_dataset
 
 
-def load_mnist(download_dir):
+def download_mnist(download_dir):
     download_dataset("mnist", download_dir)
 
     return {"train": (download_dir + "train_images", download_dir + "train_labels"),
@@ -45,8 +45,8 @@ class DatasetMNIST(Dataset):
 
 
 if __name__ == "__main__":
-    download_dir = "../downloads/mnist/"
-    mnist = load_mnist(download_dir)
+    download_dir = "downloads/mnist/"
+    mnist = download_mnist(download_dir)
 
     dataset = DatasetMNIST(*mnist["train"])
 

src/demo.py

@@ -0,0 +1,24 @@
+import torch
+import gradio as gr
+from src.predict import predict_interval, load_torch_net
+
+
+def predict_gradio_canvas(x, net, device="cuda"):
+    if x is None:
+        return {0: 0}
+    else:
+        x = torch.from_numpy(x.reshape(1, 28, 28)).to(dtype=torch.float32, device=device)
+        return predict_interval(x, net, device)
+
+
+def main(device="cuda"):
+    net = load_torch_net("../checkpoints/pytorch/version_1.pt")
+
+    gr.Interface(fn=lambda x: predict_gradio_canvas(x, net, device),
+                 inputs="sketchpad",
+                 outputs="label",
+                 live=True).launch()
+
+
+if __name__ == "__main__":
+    main(device="cpu")

src/predict.py

@@ -7,19 +7,20 @@ import numpy as np
 import plotly.express as px
 from plotly.subplots import make_subplots
 
-from trainer import LitTrainer
-from models import CNN
-from dataset import DatasetMNIST, load_mnist
+from src.trainer import LitTrainer
+from src.models import CNN
+from src.dataset import DatasetMNIST, download_mnist
 
 
-def load_pl_net(path="../checkpoints/lightning_logs/version_26/checkpoints/epoch=9-step=1000.ckpt"):
+def load_pl_net(path="checkpoints/lightning_logs/version_26/checkpoints/epoch=9-step=1000.ckpt"):
     pl_net = LitTrainer.load_from_checkpoint(path, model=CNN(1, 10))
     return pl_net
 
 
-def load_torch_net(path="../checkpoints/pytorch/version_0.pt"):
-    net = torch.load(path)
-    net.eval()
+def load_torch_net(path="checkpoints/pytorch/version_0.pt"):
+    state_dict = torch.load(path)
+    net = CNN(1, 10)
+    net.load_state_dict(state_dict)
     return net
 
 
@@ -30,22 +31,41 @@ def get_sequence(model):
 
     while i < 10:
         x, y = dataset[j]
-        y_pred = model(x.to("cuda")).detach().cpu()
-        p = torch.max(nn.functional.softmax(y_pred, dim=0))
-        y_pred = int(np.argmax(y_pred))
-        if y_pred == i and p > 0.95:
+
+        predicted, p = predict(x, model)
+
+        if predicted == i and p > 0.95:
             img = np.flip(np.array(x.reshape(28, 28)), 0)
-            fig.add_trace(px.imshow(img).data[0], row=int(i/5)+1, col=i%5+1)
+            fig.add_trace(px.imshow(img).data[0], row=int(i/5)+1, col=i % 5+1)
             i += 1
         j += 1
     return fig
 
 
+def predict(x, model, device="cuda"):
+    y_pred = model(x.to(device)).detach().cpu()
+    predicted = int(np.argmax(y_pred))
+    p = torch.max(nn.functional.softmax(y_pred, dim=0))
+
+    return predicted, p
+
+
+def predict_interval(x, model, device="cuda"):
+    y_pred = model(x.to(device))
+
+    print(y_pred)
+
+    predicted = np.argsort(y_pred.cpu().detach().numpy())
+    p = nn.functional.softmax(y_pred, dim=0)
+
+    return {int(i): float(p[i]) for i in predicted}
+
+
 if __name__ == "__main__":
-    mnist = load_mnist("../downloads/mnist/")
+    mnist = download_mnist("downloads/mnist/")
     dataset, test_data = DatasetMNIST(*mnist["train"]), DatasetMNIST(*mnist["test"])
 
     print("PyTorch Lightning Network")
-    get_sequence(load_pl_net().to("cuda")).write_image("images/pl_net.png")
+    get_sequence(load_pl_net().to("cuda")).show()
     print("Manual Network")
-    get_sequence(load_torch_net().to("cuda")).write_image("images/pytorch_net.png")
+    get_sequence(load_torch_net().to("cuda")).show()

src/train.py

@@ -1,57 +1,69 @@
 import torch
-import numpy as np
-from torch import nn, optim
-from torch.utils.data import random_split
+from torch.utils.data import random_split, DataLoader
 import pytorch_lightning as pl
+from pytorch_lightning.loggers import MLFlowLogger
 
+from src.trainer import LitTrainer
 
-from trainer import LitTrainer
-from models import CNN
 
+def argmax(a):
+    return max(range(len(a)), key=lambda x: a[x])
 
-def main():
-    from torch.utils.data import DataLoader
-    from src.dataset import DatasetMNIST, load_mnist
-
-    mnist = load_mnist("../downloads/mnist/")
-    dataset, test_data = DatasetMNIST(*mnist["train"]), DatasetMNIST(*mnist["test"])
 
+def get_dataloaders(dataset, test_data):
     train_size = round(len(dataset) * 0.8)
     validate_size = len(dataset) - train_size
     train_data, validate_data = random_split(dataset, [train_size, validate_size])
 
-    train_dataloader = DataLoader(train_data, num_workers=6)  # My CPU has 8 cores
-    validate_dataloader = DataLoader(validate_data, num_workers=2)
-    test_dataloader = DataLoader(test_data, num_workers=8)  # My CPU has 8 cores
+    # For 8 CPU cores
+    return DataLoader(train_data, num_workers=8), \
+        DataLoader(validate_data, num_workers=8), \
+        DataLoader(test_data, num_workers=8)
+
+
+def train_loop(net, batch, loss_fn, optim, device="cuda"):
+    x, y = batch
+    x = x.to(device)
+    y = y.to(device)
+
+    y_pred = net(x).reshape(1, -1)
+    loss = loss_fn(y_pred, y)
+    truth_count = argmax(y_pred.flatten()) == y
+
+    optim.zero_grad()
+    loss.backward()
+    optim.step()
+
+    return loss.item(), truth_count
+
+
+def train_net_manually(net, optim, loss_fn, train_loader, validate_loader=None, epochs=10, device="cuda"):
+    for i in range(epochs):
 
-    net = CNN(input_channels=1, num_classes=10).to("cuda")
-    opt = optim.Adam(net.parameters(), lr=1e-4)
-    loss_fn = nn.CrossEntropyLoss()
-    max_epochs = 10
-    for i in range(max_epochs):
-        for idx, batch in enumerate(train_dataloader):
-            x, y = batch
-            x = x.to("cuda")
-            y = y.to("cuda")
+        print("Epoch: 0")
 
-            y_pred = net(x).reshape(1, -1)
-            loss = loss_fn(y_pred, y)
+        epoch_loss = 0
+        epoch_truth_count = 0
+        for idx, batch in enumerate(train_loader):
+            loss, truth_count = train_loop(net, batch, loss_fn, optim, device)
 
-            opt.zero_grad()
-            loss.backward()
-            opt.step()
+            epoch_loss += loss
+            epoch_truth_count += truth_count
 
             if idx % 1000 == 0:
-                print(f"Loss: {loss.item()} ({idx} / {len(train_dataloader)})")
+                print(f"Loss: {loss} ({idx} / {len(train_loader)} x {i})")
 
-    torch.save(net, "../checkpoints/pytorch/version_1.pt")
+        print(f"Epoch Loss: {epoch_loss}")
+        print(f"Epoch Accuracy: {epoch_truth_count / len(train_loader)}")
+    torch.save(net.state_dict(), "checkpoints/pytorch/version_1.pt")
 
-    # grayscale channels = 1, mnist num_labels = 10
-    trainer = pl.Trainer(limit_train_batches=100, max_epochs=10, default_root_dir="../checkpoints")
-    pl_net = LitTrainer(CNN(input_channels=1, num_classes=10))
-    trainer.fit(pl_net, train_dataloader, validate_dataloader)
-    trainer.test(model=pl_net, dataloaders=test_dataloader)
 
+def train_net_lightning(net, optim, loss_fn, train_loader, validate_loader=None, epochs=10):
+    logger = MLFlowLogger(experiment_name="lightning_logs", tracking_uri="file:./ml-runs")
 
-if __name__ == "__main__":
-    main()
+    pl_net = LitTrainer(net)
+    pl_net.optim = optim
+    pl_net.loss = loss_fn
+    trainer = pl.Trainer(limit_train_batches=100, max_epochs=epochs,
+                         default_root_dir="../checkpoints", logger=logger)
+    trainer.fit(pl_net, train_loader, validate_loader)

src/trainer.py

@@ -1,6 +1,5 @@
 #!/usr/bin/env python
 # coding: utf-8
-import torch
 from torch import nn, optim
 import pytorch_lightning as pl