Bring miniera model up to date

hpvm/test/epoch_dnn/miniera.py

+import site
+import sys
+from os import makedirs
+from pathlib import Path
+
+import numpy as np
+from torch2hpvm import ModelExporter
+
+self_folder = Path(__file__).parent.absolute()
+site.addsitedir(self_folder.as_posix())
+
+from torch_dnn import BUFFER_NAME, quantize, run_test_over_ssh, split_and_scp
+from torch_dnn.miniera import MiniEraPL
+
+# Local configs
+ASSET_DIR = self_folder / "assets/miniera"
+QUANT_STRAT = "NONE"  # Quantization method
+WORKING_DIR = Path("/tmp/miniera")
+N_IMAGES = 50
+# Remote configs
+SCP_OPTS = "-o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -P 5506"
+SSH_OPTS = "-o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -p 5506"
+SCP_HOST = "root@espgate.cs.columbia.edu"
+SCP_PWD = "openesp"
+SCP_DST = "~/NV_NVDLA"
+
+if len(sys.argv) > 1:
+    WORKING_DIR = Path(sys.argv[1])
+
+# Reproducibility
+np.random.seed(42)
+
+# Create working directory
+makedirs(WORKING_DIR, exist_ok=True)
+
+# Calculate quantization scales
+ckpt = (ASSET_DIR / "miniera.ckpt").as_posix()
+model: MiniEraPL = MiniEraPL.load_from_checkpoint(ckpt, dataset_path=ASSET_DIR)
+scale_output = quantize(model, QUANT_STRAT, WORKING_DIR)
+
+# Code generation (into /tmp/miniera/hpvm-mod.nvdla)
+nvdla_buffer = WORKING_DIR / BUFFER_NAME
+print(f"Generating NVDLA buffer into {nvdla_buffer}")
+# You may replace scale_output (below) with ASSET_DIR / "scales/calib_NONE.txt"
+# to use precomputed quantization scale
+dataset = model.test_dataloader().dataset
+exporter = ModelExporter(model, dataset, WORKING_DIR, ASSET_DIR / "scales/calib_NONE.txt")
+exporter.generate(n_images=N_IMAGES).compile(WORKING_DIR / "miniera", WORKING_DIR)
+
+# SCP essential files to remote device
+input_images = exporter.dataset_dir
+split_and_scp([nvdla_buffer, input_images], SCP_HOST, SCP_DST, SCP_PWD, SCP_OPTS)
+
+# SSH to run test remotely
+n_total, correct = 0, 0
+for image_path, output in run_test_over_ssh(SCP_HOST, SCP_PWD, SCP_DST, input_images, SSH_OPTS):
+    idx, label = [int(s) for s in image_path.stem.split("_")]
+    output = np.array(output)
+    print(idx, output.argmax(), label, output)
+    n_total += 1
+    correct += int(output.argmax() == label)
+print(f"Accuracy: {correct / n_total * 100}% ({n_total} images)")

hpvm/test/epoch_dnn/torch_dnn/__init__.py

 from .quantizer import quantize
+from ._utils import split_and_scp, run_test_over_ssh, BUFFER_NAME

hpvm/test/epoch_dnn/torch_dnn/_utils.py

+from pathlib import Path
+import pexpect
+
+# This is hardcoded in the NVDLA pass and will not change any time soon
+BUFFER_NAME = "hpvm-mod.nvdla"
+
+
+def split_and_scp(
+    local_srcs: list, host: str, remote_dst: str, password: str, options: str
+):
+    print(f"Copying files to remote host {host}...")
+    args = options.split(" ")
+    local_srcs = [str(s) for s in local_srcs]
+    args += ["-r", *local_srcs, f"{host}:{remote_dst}"]
+    child = pexpect.spawn("scp", args, timeout=None)
+    child.expect(r"password:")
+    child.sendline(password)
+    # A rough approach to at least print something when scp is alive
+    for line in child:
+        print(line.decode(), end="")
+
+
+def run_test_over_ssh(
+    host: str, password: str, working_dir: str, image_dir: Path, options: str
+):
+    print(f"Running test on remote host {host}...")
+    args = options.split(" ") + [host]
+    child = pexpect.spawn("ssh", args, timeout=None)
+    child.expect(r"password:")
+    child.sendline(password)
+    child.expect("# ")  # The bash prompt
+    child.sendline(f"cd {working_dir}")
+    child.expect("# ")
+    child.delimiter = "# "
+    for image in image_dir.glob("*"):
+        remote_path = f"{image_dir.name}/{image.name}"
+        print(f"Sending {image.name} to run")
+        child.sendline(
+            f"./nvdla_runtime --loadable {BUFFER_NAME} --image {remote_path} --rawdump"
+        )
+        child.expect("# ")
+        child.sendline("cat output.dimg")
+        child.expect("# ")
+        result_lines = child.before.decode().splitlines()
+        # Should have 2 lines. First line is the command we keyed in.
+        output = [int(s) for s in result_lines[1].strip().split()]
+        yield image, output

hpvm/test/epoch_dnn/torch_dnn/miniera/__init__.py

-from .dataset import CIFAR
-from .model import MiniERA
+from .model import MiniEraPL

hpvm/test/epoch_dnn/torch_dnn/miniera/model.py

@@ -2,27 +2,33 @@ from pathlib import Path
 from typing import Union
 
 import numpy as np
+import pytorch_lightning as pl
 import torch
-from torch.nn import Conv2d, Linear, MaxPool2d, Module, ReLU, Sequential, Softmax
+from torch import nn
+from torch.utils.data import DataLoader
 
+from .dataset import CIFAR
 
-class MiniERA(Module):
+PathLike = Union[Path, str]
+
+
+class MiniERA(nn.Module):
     def __init__(self):
         super().__init__()
-        self.convs = Sequential(
-            Conv2d(3, 32, 3),
-            ReLU(),
-            Conv2d(32, 32, 3),
-            ReLU(),
-            MaxPool2d(2, 2),
-            Conv2d(32, 64, 3),
-            ReLU(),
-            Conv2d(64, 64, 3),
-            ReLU(),
-            MaxPool2d(2, 2),
+        self.convs = nn.Sequential(
+            nn.Conv2d(3, 32, 3),
+            nn.ReLU(),
+            nn.Conv2d(32, 32, 3),
+            nn.ReLU(),
+            nn.MaxPool2d(2, 2),
+            nn.Conv2d(32, 64, 3),
+            nn.ReLU(),
+            nn.Conv2d(64, 64, 3),
+            nn.ReLU(),
+            nn.MaxPool2d(2, 2),
         )
-        self.fcs = Sequential(Linear(1600, 256), ReLU(), Linear(256, 5))
-        self.softmax = Softmax(1)
+        self.fcs = nn.Sequential(nn.Linear(1600, 256), nn.ReLU(), nn.Linear(256, 5))
+        self.softmax = nn.Softmax(1)
 
     def forward(self, input):
         outputs = self.convs(input)
@@ -34,7 +40,7 @@ class MiniERA(Module):
         # Load in model convolution weights
         count = 0
         for conv in self.convs:
-            if not isinstance(conv, Conv2d):
+            if not isinstance(conv, nn.Conv2d):
                 continue
             weight_np = np.fromfile(
                 prefix / f"conv2d_{count+1}_w.bin", dtype=np.float32
@@ -46,7 +52,7 @@ class MiniERA(Module):
         # Load in model fc weights
         count = 0
         for linear in self.fcs:
-            if not isinstance(linear, Linear):
+            if not isinstance(linear, nn.Linear):
                 continue
             weight_np = np.fromfile(prefix / f"dense_{count+1}_w.bin", dtype=np.float32)
             bias_np = np.fromfile(prefix / f"dense_{count+1}_b.bin", dtype=np.float32)
@@ -55,3 +61,47 @@ class MiniERA(Module):
             linear.bias.data = torch.tensor(bias_np).reshape(linear.bias.shape)
             count += 1
         return self
+
+
+class MiniEraPL(pl.LightningModule):
+    def __init__(self, dataset_path: PathLike):
+        super().__init__()
+        self.network = MiniERA()
+        self.dataset_path = Path(dataset_path)
+
+    def forward(self, image):
+        prediction = self.network(image)
+        return prediction
+
+    @staticmethod
+    def _get_loss(output, targets):
+        from torch.nn.functional import cross_entropy
+
+        return cross_entropy(output, targets)
+
+    @staticmethod
+    def _get_metric(output, targets):
+        predicted = torch.argmax(output, 1)
+        return (predicted == targets).sum().item() / len(targets)
+
+    def validation_step(self, val_batch, batch_idx):
+        images, target = val_batch
+        prediction = self(images)
+        loss = self._get_loss(prediction, target)
+        accuracy = self._get_metric(prediction, target)
+        self.log("val_loss", loss)
+        self.log("val_acc", accuracy)
+        return accuracy
+
+    def test_step(self, test_batch, batch_idx):
+        images, target = test_batch
+        prediction = self(images)
+        accuracy = self._get_metric(prediction, target)
+        self.log("test_acc", accuracy)
+        return accuracy
+
+    def test_dataloader(self):
+        dataset = CIFAR.from_file(
+            self.dataset_path / "input.bin", self.dataset_path / "labels.bin"
+        )
+        return DataLoader(dataset, batch_size=128)

hpvm/test/epoch_dnn/yolo.py

@@ -9,59 +9,9 @@ from torch2hpvm import ModelExporter
 self_folder = Path(__file__).parent.absolute()
 site.addsitedir(self_folder.as_posix())
 
-from torch_dnn import quantize
+from torch_dnn import BUFFER_NAME, quantize, run_test_over_ssh, split_and_scp
 from torch_dnn.yolo import TinyYoloPL
 
-# Consts (don't change)
-BUFFER_NAME = "hpvm-mod.nvdla"
-
-
-def split_and_scp(
-    local_srcs: list, host: str, remote_dst: str, password: str, options: str
-):
-    import pexpect
-
-    print(f"Copying files to remote host {host}...")
-    args = options.split(" ")
-    local_srcs = [str(s) for s in local_srcs]
-    args += ["-r", *local_srcs, f"{host}:{remote_dst}"]
-    child = pexpect.spawn("scp", args, timeout=None)
-    child.expect(r"password:")
-    child.sendline(password)
-    # A rough approach to at least print something when scp is alive
-    for line in child:
-        print(line.decode(), end="")
-
-
-def run_test_over_ssh(
-    host: str, password: str, working_dir: str, image_dir: Path, options: str
-):
-    import pexpect
-
-    print(f"Running test on remote host {host}...")
-    args = options.split(" ") + [host]
-    child = pexpect.spawn("ssh", args, timeout=None)
-    child.expect(r"password:")
-    child.sendline(password)
-    child.expect("# ")  # The bash prompt
-    child.sendline(f"cd {working_dir}")
-    child.expect("# ")
-    child.delimiter = "# "
-    for image in image_dir.glob("*"):
-        remote_path = f"{image_dir.name}/{image.name}"
-        print(f"Sending {image.name} to run")
-        child.sendline(
-            f"./nvdla_runtime --loadable {BUFFER_NAME} --image {remote_path} --rawdump"
-        )
-        child.expect("# ")
-        child.sendline("cat output.dimg")
-        child.expect("# ")
-        result_lines = child.before.decode().splitlines()
-        # Should have 2 lines. First line is the command we keyed in.
-        output = [int(s) for s in result_lines[1].strip().split()]
-        yield image, output
-
-
 # Local configs
 ASSET_DIR = self_folder / "assets/yolo"
 QUANT_STRAT = "NONE"  # Quantization method
@@ -84,17 +34,16 @@ np.random.seed(42)
 makedirs(WORKING_DIR, exist_ok=True)
 
 # Calculate quantization scales
-# model = TinyYoloPL(num_classes=12)
 ckpt = (ASSET_DIR / "yolo_atrv2.ckpt").as_posix()
 model = TinyYoloPL.load_from_checkpoint(
     ckpt, num_classes=12, dataset_path=ASSET_DIR / "atr_dataset.tar.gz"
 )
-dataset = model.test_dataloader().dataset
 scale_output = quantize(model, QUANT_STRAT, WORKING_DIR, gpus=1)
 
 # Code generation (into /tmp/miniera/hpvm-mod.nvdla)
 nvdla_buffer = WORKING_DIR / BUFFER_NAME
 print(f"Generating NVDLA buffer into {nvdla_buffer}")
+dataset = model.test_dataloader().dataset
 exporter = ModelExporter(model, dataset, WORKING_DIR, scale_output)
 exporter.generate(n_images=N_IMAGES).compile(WORKING_DIR / "miniera", WORKING_DIR)
 
@@ -103,14 +52,7 @@ input_images = exporter.dataset_dir
 split_and_scp([nvdla_buffer, input_images], SCP_HOST, SCP_DST, SCP_PWD, SCP_OPTS)
 
 # SSH to run test remotely
-# n_total, correct = 0, 0
-# image_path, output = next(
-#     run_test_over_ssh(SCP_HOST, SCP_PWD, SCP_DST, input_images, SSH_OPTS)
-# )
-# for image_path, output in run_test_over_ssh(SCP_HOST, SCP_PWD, SCP_DST, input_images, SSH_OPTS):
-#     idx, label = [int(s) for s in image_path.stem.split("_")]
-#     output = np.array(output)
-#     print(idx, output.argmax(), label, output)
-#     n_total += 1
-#     correct += int(output.argmax() == label)
-# print(f"Accuracy: {correct / n_total * 100}% ({n_total} images)")
+n_total, correct = 0, 0
+image_path, output = next(
+    run_test_over_ssh(SCP_HOST, SCP_PWD, SCP_DST, input_images, SSH_OPTS)
+)