Supports passing param for model saving path from torchapp

predtuner/modeledapp.py

@@ -163,9 +163,9 @@ class LinearPerfModel(IPerfModel):
 
     def measure_perf(self, with_approxes: KnobsT) -> float:
         """We implement this using a weighted linear performance model."""
-        return float(sum(
-            self.cost_df.loc[layer, knob] for layer, knob in with_approxes.items()
-        ))
+        return float(
+            sum(self.cost_df.loc[layer, knob] for layer, knob in with_approxes.items())
+        )
 
 
 class QoSModelP1(IQoSModel):
@@ -212,7 +212,7 @@ class QoSModelP1(IQoSModel):
 
     def _init(self):
         dt = self.delta_tensors
-        btensor = self.baseline_tensor 
+        btensor = self.baseline_tensor
         if self.storage and self.storage.is_file():
             for op, knob, delta_tensor in self._load(self.storage):
                 dt[op][knob] = delta_tensor
@@ -239,9 +239,12 @@ class QoSModelP1(IQoSModel):
     def _try_append_save(
         path: Optional[Path], op_name: str, knob_name: str, tensor: torch.Tensor
     ):
+        import os
+
         if not path:
             return
-        path.touch(exist_ok=True)
+        if not path.parent.is_dir():
+            os.makedirs(path.parent)
         with path.open("ab") as f:
             pickle.dump((op_name, knob_name, tensor), f)
 
@@ -272,9 +275,10 @@ class QoSModelP2(IQoSModel):
 
     def measure_qos(self, with_approxes: KnobsT) -> float:
         assert self.baseline_qos is not None and self.qos_df is not None
-        delta_qoses = np.array(
-            [self.qos_df.loc[kv] for kv in with_approxes.items()]
-        ) - self.baseline_qos
+        delta_qoses = (
+            np.array([self.qos_df.loc[kv] for kv in with_approxes.items()])
+            - self.baseline_qos
+        )
         ret = delta_qoses.sum() + self.baseline_qos
         assert not np.isnan(ret)
         return float(ret)
@@ -285,10 +289,11 @@ class QoSModelP2(IQoSModel):
         else:
             knob_names = [k.name for k in self.app.knobs]
             self.qos_df = pd.DataFrame(index=self.app.ops, columns=knob_names)
+            self.qos_df = self.qos_df.where(pd.notnull(self.qos_df), None)
             self.baseline_qos = self._empirical_measure_qos({})
         df = self.qos_df
         for op, knob in barred_ravel_knobs(self.app):
-            if not np.isnan(df.loc[op, knob]):
+            if df.loc[op, knob] is not None:
                 continue
             df.loc[op, knob] = self._empirical_measure_qos({op: knob})
         if self.storage and not self.storage.is_file():
@@ -312,6 +317,10 @@ class QoSModelP2(IQoSModel):
         return df, baseline_qos
 
     def _save(self, path: Path):
+        import os
+
+        if not path.parent.is_dir():
+            os.makedirs(path.parent)
         with path.open("w") as f:
             json.dump(
                 {

predtuner/torchapp.py

 import abc
-from typing import Any, Callable, Dict, List, Set, Tuple, Union
+from pathlib import Path
+from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
 
 import numpy as np
 import torch
 from torch.nn import Module
 from torch.utils.data.dataloader import DataLoader
 
+from ._logging import PathLike
 from .approxapp import ApproxKnob, KnobsT
-from .modeledapp import (
-    IPerfModel,
-    IQoSModel,
-    LinearPerfModel,
-    ModeledApp,
-    QoSModelP1,
-    QoSModelP2,
-)
+from .modeledapp import (IPerfModel, IQoSModel, LinearPerfModel, ModeledApp,
+                         QoSModelP1, QoSModelP2)
 from .torchutil import ModuleIndexer, get_summary, move_to_device_recursively
 
 
@@ -64,6 +60,7 @@ class TorchApp(ModeledApp, abc.ABC):
         tensor_to_qos: Callable[[torch.Tensor, Any], float],
         combine_qos: Callable[[np.ndarray], float] = np.mean,
         device: Union[torch.device, str] = _default_device,
+        model_storage_folder: Optional[PathLike] = None
     ) -> None:
         self.app_name = app_name
         self.module = module
@@ -73,6 +70,7 @@ class TorchApp(ModeledApp, abc.ABC):
         self.tensor_to_qos = tensor_to_qos
         self.combine_qos = combine_qos
         self.device = device
+        self.model_storage = Path(model_storage_folder)
 
         self.module = self.module.to(device)
         self.midx = ModuleIndexer(module)
@@ -113,10 +111,12 @@ class TorchApp(ModeledApp, abc.ABC):
                 qoses.append(qos)
             return self.combine_qos(np.array(qoses))
 
+        p1_storage = self.model_storage / "p1.pkl" if self.model_storage else None
+        p2_storage = self.model_storage / "p2.json" if self.model_storage else None
         return [
             LinearPerfModel(self._op_costs, self._knob_speedups),
-            QoSModelP1(self, self._get_raw_output_valset, batched_valset_qos),
-            QoSModelP2(self),
+            QoSModelP1(self, self._get_raw_output_valset, batched_valset_qos, p1_storage),
+            QoSModelP2(self, p2_storage),
         ]
 
     @torch.no_grad()