feature: Added prediction and model loading methods

src/predict.py

@@ -0,0 +1,51 @@
+#!/usr/bin/env python
+# coding: utf-8
+import torch
+from torch import nn
+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
+
+
+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()
+    return net
+
+
+def get_sequence(model):
+    fig = make_subplots(rows=2, cols=5)
+
+    i, j = 0, np.random.randint(0, 30000)
+
+    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:
+            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)
+            i += 1
+        j += 1
+    return fig
+
+
+if __name__ == "__main__":
+    mnist = load_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")
+    print("Manual Network")
+    get_sequence(load_torch_net().to("cuda")).write_image("images/pytorch_net.png")