commit 148f75498de0746d30bc8ded9ca458e0bbd936e4
parent 6ecb5e86970c22c3a69325c8fc01b9932707d32c
Author: Steven Atkinson <steven@atkinson.mn>
Date: Sun, 5 Feb 2023 12:10:36 -0800
Single-file model export (#77)
* New export pattern exports to a single file.
* Update instructions for downloading your model
* Update easy mode instructions for downloading your model
Diffstat:
5 files changed, 73 insertions(+), 11 deletions(-)
diff --git a/bin/train/colab.ipynb b/bin/train/colab.ipynb
@@ -503,6 +503,7 @@
]
},
{
+ "attachments": {},
"cell_type": "markdown",
"metadata": {
"id": "823KJ_L0Rchp"
@@ -510,9 +511,7 @@
"source": [
"## Step 7: Download your artifacts\n",
"We're done! \n",
- "Go to the file browser on the left panel ⬅ and download the contents of `exported_model` (you may need to hit the refresh button).\n",
- "\n",
- "You'll need `config.json` (the architecture) and `weights.npy` (the weights)--these are the information that the NAM plugin needs to run your model!\n",
+ "Go to the file browser on the left panel ⬅ and download `model.nam` from the `exported_model` directory (you may need to hit the refresh button).\n",
"\n",
"Additionally, if you want to continue to train this model later you can download the lightning model artifacts from `lightning_logs`. If not, that's fine.\n",
"\n",
diff --git a/bin/train/easy_colab.ipynb b/bin/train/easy_colab.ipynb
@@ -109,6 +109,7 @@
]
},
{
+ "attachments": {},
"cell_type": "markdown",
"metadata": {
"id": "823KJ_L0Rchp"
@@ -118,10 +119,8 @@
"We're done!\n",
"\n",
"Have a look at the plot above to see how your model compares to the real gear you're modeling.\n",
- "Hopefully it looks\n",
- "Go to the file browser on the left panel ⬅ and download the contents of `exported_model` (you may need to hit the refresh button).\n",
- "\n",
- "You'll need `config.json` (the architecture) and `weights.npy` (the weights)--these are the information that the NAM plugin needs to run your model!\n",
+ "Hopefully it looks good!\n",
+ "Go to the file browser on the left panel ⬅ and download `model.nam` from the `exported_model` directory (you may need to hit the refresh button).\n",
"\n",
"Additionally, if you want to continue to train this model later you can download the lightning model artifacts from `lightning_logs`. If not, that's fine too.\n",
"\n",
@@ -151,7 +150,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
- "version": "3.10.6"
+ "version": "3.10.6 | packaged by conda-forge | (main, Oct 24 2022, 16:02:16) [MSC v.1916 64 bit (AMD64)]"
},
"orig_nbformat": 4,
"vscode": {
diff --git a/nam/_version.py b/nam/_version.py
@@ -1 +1 @@
-__version__ = "0.4.0"
+__version__ = "0.5.0"
diff --git a/nam/models/_exportable.py b/nam/models/_exportable.py
@@ -30,17 +30,17 @@ class Exportable(abc.ABC):
"""
training = self.training
self.eval()
- with open(Path(outdir, "config.json"), "w") as fp:
+ with open(Path(outdir, "model.nam"), "w") as fp:
json.dump(
{
"version": __version__,
"architecture": self.__class__.__name__,
"config": self._export_config(),
+ "weights": self._export_weights().tolist(),
},
fp,
indent=4,
)
- np.save(Path(outdir, "weights.npy"), self._export_weights())
x, y = self._export_input_output()
np.save(Path(outdir, "inputs.npy"), x)
np.save(Path(outdir, "outputs.npy"), y)
diff --git a/tests/test_nam/test_models/test_exportable.py b/tests/test_nam/test_models/test_exportable.py
@@ -0,0 +1,63 @@
+# File: test_exportable.py
+# Created Date: Sunday January 29th 2023
+# Author: Steven Atkinson (steven@atkinson.mn)
+
+"""
+Test export behavior of models
+"""
+
+import json
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from typing import Tuple
+
+import numpy as np
+import pytest
+import torch
+import torch.nn as nn
+
+from nam.models import _exportable
+
+def test_export():
+ """
+ Does it work?
+ """
+ class Model(nn.Module, _exportable.Exportable):
+ def __init__(self):
+ super().__init__()
+ self._scale = nn.Parameter(torch.tensor(0.0))
+ self._bias = nn.Parameter(torch.tensor(0.0))
+
+ def forward(self, x: torch.Tensor):
+ return self._scale * x + self._bias
+
+ def export_cpp_header(self, filename: Path):
+ pass
+
+ def _export_config(self):
+ return {}
+
+ def _export_input_output(self) -> Tuple[np.ndarray, np.ndarray]:
+ x = 0.01 * np.random.randn(3,)
+ y = self(torch.Tensor(x)).detach().cpu().numpy()
+ return x, y
+
+ def _export_weights(self) -> np.ndarray:
+ return torch.stack([self._scale, self._bias]).detach().cpu().numpy()
+
+ model = Model()
+ with TemporaryDirectory() as tmpdir:
+ model.export(tmpdir)
+ model_basename = "model.nam"
+ model_path = Path(tmpdir, model_basename)
+ assert model_path.exists()
+ with open(model_path, "r") as fp:
+ model_dict = json.load(fp)
+ required_keys = {"version", "architecture", "config", "weights"}
+ for key in required_keys:
+ assert key in model_dict
+ weights_list = model_dict["weights"]
+ assert isinstance(weights_list, list)
+ assert len(weights_list) == 2
+ assert all(isinstance(w, float) for w in weights_list)
+
+\ No newline at end of file
