Small changes to end-to-end script to make it easier to use

README.md

@@ -61,9 +61,7 @@ Run `test/epoch_dnn/main.py` to generate the PyTorch DNN defined at `test/epoch_
 and evaluate it on real hardware:
 
 ```bash
-cd test/epoch_dnn
-# rm -r /tmp/miniera  # Do this if you have run main.py before
-python main.py
+cd test/epoch_dnn; python main.py
 ```
 
 This script will run through the whole code generation (putting everything under `/tmp/miniera`) and evaluation, including

hpvm/projects/torch2hpvm/torch2hpvm/compile.py

@@ -2,6 +2,7 @@ import os
 from pathlib import Path
 from tempfile import NamedTemporaryFile
 from typing import IO, Optional, Sequence, Union
+from shutil import rmtree
 
 import numpy as np
 import onnx
@@ -40,10 +41,10 @@ class ModelExporter:
         self.output_dir = output_dir = Path(output_dir).absolute()
         os.makedirs(output_dir, exist_ok=True)
         self.weight_dir = output_dir / self.weight_dir_name
-        self.weight_dir.mkdir(exist_ok=True)
+        self._remove_and_mkdir(self.weight_dir)
         self.codefile = output_dir / self.source_file_name
         self.dataset_dir = output_dir / self.dataset_dir_name
-        self.dataset_dir.mkdir(exist_ok=True)
+        self._remove_and_mkdir(self.dataset_dir)
 
         self.path_params = {}
         assert target == "hpvm_nvdla", f"Target {target} not recognized"
@@ -142,6 +143,12 @@ class ModelExporter:
         assert check, "Simplified ONNX model could not be validated"
         return onnx.shape_inference.infer_shapes(onnx_model)
 
+    @staticmethod
+    def _remove_and_mkdir(path: Path):
+        if path.exists():
+            rmtree(path)
+        os.mkdir(path)
+
 
 def check_onnx_version(model, new_version):
     try:

hpvm/test/epoch_dnn/main.py

@@ -30,7 +30,7 @@ def split_and_scp(
     child.sendline(password)
     # A rough approach to at least print something when scp is alive
     for line in child:
-        print(line.decode())
+        print(line.decode(), end="")
 
 
 def run_test_over_ssh(host: str, password: str, working_dir: str, image_dir: Path, options: str):
@@ -49,7 +49,7 @@ def run_test_over_ssh(host: str, password: str, working_dir: str, image_dir: Pat
         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("Test pass")
+        child.expect("# ")
         child.sendline("cat output.dimg")
         child.expect("# ")
         result_lines = child.before.decode().splitlines()
@@ -73,7 +73,8 @@ SCP_DST = "~/NV_NVDLA"
 # Reproducibility
 np.random.seed(42)
 
-makedirs(WORKING_DIR, exist_ok=False)
+# Create working directory
+makedirs(WORKING_DIR, exist_ok=True)
 
 # Calculate quantization scales
 ckpt = (ASSET_DIR / "miniera.pth").as_posix()
@@ -88,7 +89,9 @@ print(f"Generating NVDLA buffer into {nvdla_buffer}")
 bin_dataset = BinDataset(
     ASSET_DIR / "input.bin", ASSET_DIR / "labels.bin", (5000, 3, 32, 32)
 )
-exporter = ModelExporter(model, bin_dataset, WORKING_DIR, ASSET_DIR / "scales/calib_NONE.txt")
+# You may replace scale_output (below) with ASSET_DIR / "scales/calib_NONE.txt"
+# to use precomputed quantization scale
+exporter = ModelExporter(model, bin_dataset, WORKING_DIR, scale_output)
 exporter.generate(n_images=N_IMAGES).compile(WORKING_DIR / "miniera", WORKING_DIR)
 
 # SCP essential files to remote device