Feature: Thunks as an alternative to artifacts #120
Merged
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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.
Adds two conditional branches to disable typechecks in the _check() method. This is nice to have when prototyping new features and inputs to benchmarks might not exactly be of the requested types.
Also stop the practice of binding partial parametrizations directly to benchmarks. This has the effect that we can manipulate benchmark function parameters if need be (for example by lazy-loading thunk parameters). Changes the interface construction slightly to inject the partial parametrization as defaults over the `inspect.Parameter` default values.
Slightly changes parameter construction and adds a dethunking step right before the benchmark loop. This means that the thunk values are accessed at the latest possible time, which is just before benchmark execution. Moves the context construction ahead of the empty collection check, so that we give back a constructed context even in the case of no found benchmarks. Adds two C++-style thunk helpers, `is_thunk` for deciding if a value is a thunk, and `is_thunk_type` to decide if a value type is a thunk type annotation. The whole thunk facility is designed to work both with the `nnbench.types.Thunk` type as well as with general anonymous functions.
In the current setup, properly typed memos and callables pass the type checker. Factors out the types into their own submodule, to be refactored later into their biggest constituents.
Showcases memo subclassing, parametrization, and trivializes the run() command again. As a downside, only per-class benchmarks and aggregates can be run in a single run, not side-by-side (that would require `params` injection).
Partial parametrizations are not bound eagerly to the benchmark functions anymore, which makes it simpler to inject memos and de-memoize variables just in time for execution. What is left is validation that a subsequent benchmark of models with intermittent garbage collection actually reaps each model after the benchmark is done.
nicholasjng
force-pushed
the
modelthunks
branch
from
c789d72 to
98c2909
Compare
This was referenced Mar 21, 2024
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This branch contains
Thunks as an alternative to Artifacts. Wiki -> https://en.wikipedia.org/wiki/Thunk.This is basically the same as React's
useMemo. Thunks are lazy-loadable similarly to artifacts, customizable as a class if state is required (with the genericnnbench.types.Thunktype), and usable as a drop-in replacement for a normal e.g.torch.Modulevalue to be computed at benchmark runtime.The strategy here is to bootstrap each value in the
Thunk.__call__()method, relying on caching as much as possible to avoid recomputing values.Another important addition is the
BenchmarkRunner.typecheckflag, which can be turned off to skip typechecks and avoid shortcomings / bugs in our current typechecking logic.Still to do: