Skip to content

Commit

Permalink
Add transform submodule, parameter compression transform (#124)
Browse files Browse the repository at this point in the history
This is the better way of compressing parameters compared to directly in
the benchmark runner, which steals responsibility of the transform that
we just introduced.

Refactors `nnbench.io.transform->nnbench.transforms`, the latter being
its own submodule. This is useful to have when adding new builtin
transforms, so that they do not have to go into a single file.
  • Loading branch information
nicholasjng authored Mar 21, 2024
1 parent aac4162 commit d09cd3c
Show file tree
Hide file tree
Showing 5 changed files with 51 additions and 2 deletions.
2 changes: 1 addition & 1 deletion docs/guides/transforms.md
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ After a successful benchmark run execution, you end up with your metrics, contex
In general, this data is a best-effort representation of the environment and configuration the benchmarks are run in.

However, in some situations, manual editing and transformation of these records is required.
nnbench exposes the `nnbench.io.transforms` module to facilitate these transforms.
nnbench exposes the `nnbench.transforms` module to facilitate these transforms.

## Types of transforms: 1->1 vs. N->1 vs. N->N

Expand Down
2 changes: 1 addition & 1 deletion examples/transforms/transforms.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,8 @@
import numpy as np

import nnbench
from nnbench.io.transforms import OneToOneTransform
from nnbench.reporter.file import FileIO
from nnbench.transforms import OneToOneTransform
from nnbench.types import BenchmarkRecord


Expand Down
1 change: 1 addition & 0 deletions src/nnbench/transforms/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
from .base import ManyToManyTransform, ManyToOneTransform, OneToOneTransform
File renamed without changes.
48 changes: 48 additions & 0 deletions src/nnbench/transforms/params.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,48 @@
from typing import Any, Sequence

from nnbench.transforms import ManyToManyTransform, OneToOneTransform
from nnbench.types import BenchmarkRecord


class CompressionMixin:
def compress(self, params: dict[str, Any]) -> dict[str, Any]:
containers = (tuple, list, set, frozenset)
natives = (float, int, str, bool, bytes, complex)
compressed: dict[str, Any] = {}

def _compress_impl(val):
if isinstance(val, natives):
# save native types without modification...
return val
else:
# ... or return the string repr.
# TODO: Allow custom representations for types with formatters.
return repr(val)

for k, v in params.items():
if isinstance(v, containers):
container_type = type(v)
compressed[k] = container_type(_compress_impl(vv) for vv in v)
elif isinstance(v, dict):
compressed[k] = self.compress(v)
else:
compressed[k] = _compress_impl(v)

return compressed


class ParameterCompression1to1(OneToOneTransform, CompressionMixin):
def apply(self, record: BenchmarkRecord) -> BenchmarkRecord:
for bm in record.benchmarks:
bm["params"] = self.compress(bm["params"])

return record


class ParameterCompressionNtoN(ManyToManyTransform, CompressionMixin):
def apply(self, record: Sequence[BenchmarkRecord]) -> Sequence[BenchmarkRecord]:
for rec in record:
for bm in rec.benchmarks:
bm["params"] = self.compress(bm["params"])

return record

0 comments on commit d09cd3c

Please sign in to comment.