trainlib 0.3.1__tar.gz → 0.6.1__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- trainlib-0.6.1/PKG-INFO +243 -0
- trainlib-0.6.1/README.md +206 -0
- {trainlib-0.3.1 → trainlib-0.6.1}/pyproject.toml +22 -11
- trainlib-0.6.1/test/test_cli.py +160 -0
- trainlib-0.6.1/test/test_config.py +87 -0
- trainlib-0.6.1/test/test_console_display.py +101 -0
- trainlib-0.6.1/test/test_construction.py +191 -0
- trainlib-0.6.1/test/test_experiment_runs.py +85 -0
- trainlib-0.6.1/test/test_objectives.py +56 -0
- trainlib-0.6.1/test/test_plot.py +53 -0
- trainlib-0.6.1/test/test_prepared_domain.py +86 -0
- trainlib-0.6.1/test/test_records.py +165 -0
- trainlib-0.6.1/test/test_report_outputs.py +52 -0
- trainlib-0.6.1/test/test_trainer.py +107 -0
- trainlib-0.6.1/trainlib/__main__.py +35 -0
- trainlib-0.6.1/trainlib/channel.py +78 -0
- trainlib-0.6.1/trainlib/channels/__init__.py +23 -0
- trainlib-0.6.1/trainlib/channels/duration.py +21 -0
- trainlib-0.6.1/trainlib/channels/loss.py +39 -0
- trainlib-0.6.1/trainlib/channels/regression.py +94 -0
- {trainlib-0.3.1 → trainlib-0.6.1}/trainlib/cli/__init__.py +2 -2
- trainlib-0.6.1/trainlib/cli/config.py +50 -0
- trainlib-0.6.1/trainlib/cli/experiment.py +196 -0
- trainlib-0.6.1/trainlib/cli/train.py +69 -0
- trainlib-0.6.1/trainlib/components/loss/__init__.py +26 -0
- trainlib-0.6.1/trainlib/components/optimizer/__init__.py +13 -0
- trainlib-0.6.1/trainlib/construction.py +452 -0
- {trainlib-0.3.1 → trainlib-0.6.1}/trainlib/dataloader.py +74 -26
- trainlib-0.6.1/trainlib/dataloaders/__init__.py +20 -0
- trainlib-0.6.1/trainlib/dataloaders/memory.py +48 -0
- {trainlib-0.3.1 → trainlib-0.6.1}/trainlib/dataset.py +95 -100
- trainlib-0.6.1/trainlib/datasets/__init__.py +13 -0
- {trainlib-0.3.1 → trainlib-0.6.1}/trainlib/datasets/disk.py +8 -15
- {trainlib-0.3.1 → trainlib-0.6.1}/trainlib/datasets/memory.py +29 -17
- {trainlib-0.3.1 → trainlib-0.6.1}/trainlib/domain.py +11 -26
- {trainlib-0.3.1 → trainlib-0.6.1}/trainlib/domains/disk.py +6 -7
- {trainlib-0.3.1 → trainlib-0.6.1}/trainlib/domains/functional.py +2 -6
- trainlib-0.6.1/trainlib/domains/prepared.py +212 -0
- trainlib-0.6.1/trainlib/estimator.py +148 -0
- {trainlib-0.3.1 → trainlib-0.6.1}/trainlib/estimators/__init__.py +9 -4
- {trainlib-0.3.1 → trainlib-0.6.1}/trainlib/estimators/mlp.py +16 -72
- {trainlib-0.3.1 → trainlib-0.6.1}/trainlib/estimators/rnn.py +53 -139
- {trainlib-0.3.1 → trainlib-0.6.1}/trainlib/estimators/tdnn.py +8 -12
- trainlib-0.6.1/trainlib/experiment.py +430 -0
- trainlib-0.6.1/trainlib/experiments/__init__.py +18 -0
- trainlib-0.6.1/trainlib/experiments/regression.py +16 -0
- trainlib-0.6.1/trainlib/experiments/train.py +107 -0
- trainlib-0.6.1/trainlib/loss.py +172 -0
- trainlib-0.6.1/trainlib/losses/__init__.py +17 -0
- trainlib-0.6.1/trainlib/losses/composite.py +53 -0
- trainlib-0.6.1/trainlib/losses/criteria.py +78 -0
- trainlib-0.6.1/trainlib/losses/simple.py +125 -0
- trainlib-0.6.1/trainlib/measure.py +95 -0
- trainlib-0.6.1/trainlib/measures/__init__.py +11 -0
- trainlib-0.6.1/trainlib/measures/regression.py +49 -0
- trainlib-0.6.1/trainlib/objective.py +39 -0
- trainlib-0.6.1/trainlib/optimizer.py +274 -0
- trainlib-0.6.1/trainlib/optimizers/__init__.py +15 -0
- trainlib-0.6.1/trainlib/optimizers/adversarial.py +40 -0
- trainlib-0.6.1/trainlib/optimizers/simple.py +49 -0
- {trainlib-0.3.1 → trainlib-0.6.1}/trainlib/plotter.py +187 -101
- trainlib-0.6.1/trainlib/plotters/__init__.py +16 -0
- trainlib-0.6.1/trainlib/plotters/regression.py +53 -0
- trainlib-0.6.1/trainlib/record.py +486 -0
- trainlib-0.6.1/trainlib/registry.py +103 -0
- trainlib-0.6.1/trainlib/reporter.py +251 -0
- trainlib-0.6.1/trainlib/reporters/__init__.py +20 -0
- trainlib-0.6.1/trainlib/run.py +590 -0
- trainlib-0.6.1/trainlib/tabulator.py +203 -0
- trainlib-0.6.1/trainlib/tabulators/__init__.py +8 -0
- trainlib-0.6.1/trainlib/trainer.py +820 -0
- {trainlib-0.3.1 → trainlib-0.6.1}/trainlib/transform.py +8 -9
- trainlib-0.6.1/trainlib/utils/__init__.py +0 -0
- trainlib-0.6.1/trainlib/utils/artifacts.py +108 -0
- trainlib-0.6.1/trainlib/utils/config.py +146 -0
- trainlib-0.6.1/trainlib/utils/console.py +527 -0
- trainlib-0.6.1/trainlib/utils/custom.mplstyle +50 -0
- {trainlib-0.3.1 → trainlib-0.6.1}/trainlib/utils/job.py +4 -7
- {trainlib-0.3.1 → trainlib-0.6.1}/trainlib/utils/module.py +14 -6
- {trainlib-0.3.1 → trainlib-0.6.1}/trainlib/utils/op.py +2 -2
- trainlib-0.6.1/trainlib/utils/plot.py +102 -0
- trainlib-0.6.1/trainlib/utils/statistics.py +188 -0
- trainlib-0.6.1/trainlib/utils/tables.py +476 -0
- {trainlib-0.3.1 → trainlib-0.6.1}/trainlib/utils/text.py +2 -0
- {trainlib-0.3.1 → trainlib-0.6.1}/trainlib/utils/type.py +23 -1
- trainlib-0.6.1/trainlib.egg-info/PKG-INFO +243 -0
- trainlib-0.6.1/trainlib.egg-info/SOURCES.txt +95 -0
- {trainlib-0.3.1 → trainlib-0.6.1}/trainlib.egg-info/entry_points.txt +1 -0
- {trainlib-0.3.1 → trainlib-0.6.1}/trainlib.egg-info/requires.txt +2 -0
- trainlib-0.3.1/PKG-INFO +0 -136
- trainlib-0.3.1/README.md +0 -100
- trainlib-0.3.1/trainlib/__main__.py +0 -18
- trainlib-0.3.1/trainlib/cli/train.py +0 -164
- trainlib-0.3.1/trainlib/dataloaders/__init__.py +0 -12
- trainlib-0.3.1/trainlib/dataloaders/memory.py +0 -17
- trainlib-0.3.1/trainlib/datasets/__init__.py +0 -12
- trainlib-0.3.1/trainlib/estimator.py +0 -201
- trainlib-0.3.1/trainlib/trainer.py +0 -639
- trainlib-0.3.1/trainlib/utils/plot.py +0 -38
- trainlib-0.3.1/trainlib.egg-info/PKG-INFO +0 -136
- trainlib-0.3.1/trainlib.egg-info/SOURCES.txt +0 -40
- {trainlib-0.3.1 → trainlib-0.6.1}/setup.cfg +0 -0
- {trainlib-0.3.1 → trainlib-0.6.1}/trainlib/__init__.py +0 -0
- {trainlib-0.3.1/trainlib/domains → trainlib-0.6.1/trainlib/components}/__init__.py +0 -0
- {trainlib-0.3.1/trainlib/utils → trainlib-0.6.1/trainlib/domains}/__init__.py +0 -0
- {trainlib-0.3.1 → trainlib-0.6.1}/trainlib/utils/map.py +0 -0
- {trainlib-0.3.1 → trainlib-0.6.1}/trainlib/utils/session.py +1 -1
- {trainlib-0.3.1 → trainlib-0.6.1}/trainlib.egg-info/dependency_links.txt +0 -0
- {trainlib-0.3.1 → trainlib-0.6.1}/trainlib.egg-info/top_level.txt +0 -0
trainlib-0.6.1/PKG-INFO
ADDED
|
@@ -0,0 +1,243 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: trainlib
|
|
3
|
+
Version: 0.6.1
|
|
4
|
+
Summary: Minimal framework for ML modeling. Supports advanced dataset operations and streamlined training.
|
|
5
|
+
License-Expression: MIT
|
|
6
|
+
Project-URL: Homepage, https://doc.olog.io/trainlib
|
|
7
|
+
Project-URL: Documentation, https://doc.olog.io/trainlib
|
|
8
|
+
Project-URL: Repository, https://git.olog.io/olog/trainlib
|
|
9
|
+
Project-URL: Issues, https://git.olog.io/olog/trainlib/issues
|
|
10
|
+
Keywords: machine-learning
|
|
11
|
+
Classifier: Programming Language :: Python
|
|
12
|
+
Classifier: Operating System :: OS Independent
|
|
13
|
+
Classifier: Development Status :: 3 - Alpha
|
|
14
|
+
Classifier: Intended Audience :: Developers
|
|
15
|
+
Classifier: Intended Audience :: End Users/Desktop
|
|
16
|
+
Requires-Python: >=3.12
|
|
17
|
+
Description-Content-Type: text/markdown
|
|
18
|
+
Requires-Dist: torch
|
|
19
|
+
Requires-Dist: pandas
|
|
20
|
+
Requires-Dist: colorama>=0.4.6
|
|
21
|
+
Requires-Dist: matplotlib>=3.10.8
|
|
22
|
+
Requires-Dist: numpy>=2.4.1
|
|
23
|
+
Requires-Dist: tensorboard>=2.20.0
|
|
24
|
+
Requires-Dist: tqdm>=4.67.1
|
|
25
|
+
Requires-Dist: rich>=13.7
|
|
26
|
+
Requires-Dist: setuptools<=81.0.0
|
|
27
|
+
Provides-Extra: dev
|
|
28
|
+
Requires-Dist: ipykernel; extra == "dev"
|
|
29
|
+
Provides-Extra: doc
|
|
30
|
+
Requires-Dist: furo; extra == "doc"
|
|
31
|
+
Requires-Dist: myst-parser; extra == "doc"
|
|
32
|
+
Requires-Dist: sphinx; extra == "doc"
|
|
33
|
+
Requires-Dist: sphinx-togglebutton; extra == "doc"
|
|
34
|
+
Requires-Dist: sphinx-autodoc-typehints; extra == "doc"
|
|
35
|
+
Provides-Extra: test
|
|
36
|
+
Requires-Dist: pytest; extra == "test"
|
|
37
|
+
|
|
38
|
+
# Overview
|
|
39
|
+
Minimal framework for ML modeling, supporting advanced dataset operations and
|
|
40
|
+
streamlined training workflows.
|
|
41
|
+
|
|
42
|
+
# Install
|
|
43
|
+
The `trainlib` package can be installed from PyPI:
|
|
44
|
+
|
|
45
|
+
```sh
|
|
46
|
+
pip install trainlib
|
|
47
|
+
```
|
|
48
|
+
|
|
49
|
+
# Development
|
|
50
|
+
- Initialize/synchronize the project with `uv sync`, creating a virtual
|
|
51
|
+
environment with base package dependencies.
|
|
52
|
+
- Depending on needs, install the development dependencies with `uv sync
|
|
53
|
+
--extra dev`.
|
|
54
|
+
|
|
55
|
+
## Testing
|
|
56
|
+
- To run the unit tests, make sure to first have the test dependencies
|
|
57
|
+
installed with `uv sync --extra test`, then run `make test`.
|
|
58
|
+
- For notebook testing, run `make install-kernel` to make the environment
|
|
59
|
+
available as a Jupyter kernel (to be selected when running notebooks).
|
|
60
|
+
|
|
61
|
+
## Documentation
|
|
62
|
+
- Install the documentation dependencies with `uv sync --extra doc`.
|
|
63
|
+
- Run `make docs-build` (optionally preceded by `make docs-clean`), and serve
|
|
64
|
+
locally with `make docs-serve`.
|
|
65
|
+
|
|
66
|
+
# Development remarks
|
|
67
|
+
- Across `Trainer` / `Estimator` / `Dataset`, I've considered a
|
|
68
|
+
`ParamSpec`-based typing scheme to better orchestrate alignment in the
|
|
69
|
+
`Trainer.train()` loop, e.g., so we can statically check whether a dataset
|
|
70
|
+
appears to be fulfilling the argument requirements for the estimator's
|
|
71
|
+
`forward()` / `metrics()` methods. Something like
|
|
72
|
+
|
|
73
|
+
```py
|
|
74
|
+
class Estimator[**P](nn.Module):
|
|
75
|
+
def forward(
|
|
76
|
+
self,
|
|
77
|
+
input: Tensor,
|
|
78
|
+
*args: P.args,
|
|
79
|
+
**kwargs: P.kwargs,
|
|
80
|
+
) -> Tensor:
|
|
81
|
+
...
|
|
82
|
+
|
|
83
|
+
class Trainer[**P]:
|
|
84
|
+
def __init__(
|
|
85
|
+
self,
|
|
86
|
+
estimator: Estimator[P],
|
|
87
|
+
...
|
|
88
|
+
): ...
|
|
89
|
+
```
|
|
90
|
+
|
|
91
|
+
might be how we begin threading signatures. But ensuring dataset items can
|
|
92
|
+
match `P` is challenging. You can consider a "packed" object where we
|
|
93
|
+
obfuscate passing data through `P`-signatures:
|
|
94
|
+
|
|
95
|
+
```py
|
|
96
|
+
class PackedItem[**P]:
|
|
97
|
+
def __init__(self, *args: P.args, **kwargs: P.kwargs) -> None:
|
|
98
|
+
self._args = args
|
|
99
|
+
self._kwargs = kwargs
|
|
100
|
+
|
|
101
|
+
def apply[R](self, func: Callable[P, R]) -> R:
|
|
102
|
+
return func(*self._args, **self._kwargs)
|
|
103
|
+
|
|
104
|
+
|
|
105
|
+
class BatchedDataset[U, R, I, **P](Dataset):
|
|
106
|
+
@abstractmethod
|
|
107
|
+
def _process_item_data(
|
|
108
|
+
self,
|
|
109
|
+
item_data: I,
|
|
110
|
+
item_index: int,
|
|
111
|
+
) -> PackedItem[P]:
|
|
112
|
+
...
|
|
113
|
+
|
|
114
|
+
def __iter__(self) -> Iterator[PackedItem[P]]:
|
|
115
|
+
...
|
|
116
|
+
```
|
|
117
|
+
|
|
118
|
+
Meaningfully shaping those signatures is what remains, but you can't really
|
|
119
|
+
do this, not with typical type expression flexibility. For instance, if I'm
|
|
120
|
+
trying to appropriately type my base `TupleDataset`:
|
|
121
|
+
|
|
122
|
+
```py
|
|
123
|
+
class SequenceDataset[I, **P](HomogenousDataset[int, I, I, P]):
|
|
124
|
+
...
|
|
125
|
+
|
|
126
|
+
class TupleDataset[I](SequenceDataset[tuple[I, ...], "?"]):
|
|
127
|
+
...
|
|
128
|
+
```
|
|
129
|
+
|
|
130
|
+
Here there's no way for me to shape a `ParamSpec` to indicate arbitrarily
|
|
131
|
+
many arguments of a fixed type (`I` in this case) to allow me to unpack my
|
|
132
|
+
item tuples into an appropriate `PackedItem` (since `ParamSpec` isn't exactly
|
|
133
|
+
controllable in that way; it's meant make a generic signature expressible as
|
|
134
|
+
a type variable, *not to express particular shapes* of the types in that
|
|
135
|
+
signature).
|
|
136
|
+
|
|
137
|
+
Until this (among other issues) becomes clearer, I'm setting up around a
|
|
138
|
+
simpler `TypedDict` type variable. We won't have particularly strong static
|
|
139
|
+
checks for item alignment inside `Trainer`, but this seems about as good as I
|
|
140
|
+
can get around the current infrastructure.
|
|
141
|
+
|
|
142
|
+
## Typing planning
|
|
143
|
+
There are currently two package-wide typing strategies geared toward future
|
|
144
|
+
Python version bumps and type system changes:
|
|
145
|
+
|
|
146
|
+
- Using `Unpack` with a type variable with a `TypedDict` upper bound, rather
|
|
147
|
+
than with a `TypedDict` directly. This is commented on extensively in
|
|
148
|
+
`trainlib/estimator.py`, but the short version is we want our key objects to
|
|
149
|
+
remain generic in their method signatures when they are expected to operate
|
|
150
|
+
on session-specific structured batch data. For instance, you can imagine an
|
|
151
|
+
`Estimator`'s `forward` method accepting a typed batch like
|
|
152
|
+
|
|
153
|
+
```py
|
|
154
|
+
def forward(self, batch: BatchKwargs): ...
|
|
155
|
+
|
|
156
|
+
# or
|
|
157
|
+
|
|
158
|
+
def forward(self, **kwargs: Unpack[Batchwargs]): ...
|
|
159
|
+
```
|
|
160
|
+
|
|
161
|
+
These are both currently supported patterns. Problems arise, however, when we
|
|
162
|
+
want an inherited class to specifically operate on an inherited type of
|
|
163
|
+
`BatchKwargs`
|
|
164
|
+
|
|
165
|
+
```py
|
|
166
|
+
class BatchKwargs(TypedDict):
|
|
167
|
+
inputs: Tensor
|
|
168
|
+
|
|
169
|
+
class ExtraBatchKwargs(BatchKwargs):
|
|
170
|
+
inputs: Tensor
|
|
171
|
+
labels: Tensor
|
|
172
|
+
|
|
173
|
+
class A(Estimator):
|
|
174
|
+
def forward(self, batch: BatchKwargs): ...
|
|
175
|
+
|
|
176
|
+
# want to require Estimator subclass handles both inputs and labels
|
|
177
|
+
class B(A):
|
|
178
|
+
def forward(self, batch: ExtraBatchKwargs): ...
|
|
179
|
+
```
|
|
180
|
+
|
|
181
|
+
Note how `ExtraBatchKwargs` doesn't use `NotRequired` for the new `labels`
|
|
182
|
+
attribute; the intention here is to actually require that `B` should handle
|
|
183
|
+
*more attributes* than `A`. **That breaks LSP**, so this doesn't work. You'd
|
|
184
|
+
need to either break out the `Estimator` hierarchies, or decide that `B`
|
|
185
|
+
could have inputs typed as `BatchKwargs` and possibly implement some extra
|
|
186
|
+
handling to check if batches of type `ExtraBatchKwargs` are passed in (note
|
|
187
|
+
that passing in a batch of type `ExtraBatchKwargs` to `A.forward` is
|
|
188
|
+
perfectly fine) and supply labels. **Or...** you make the class generic in
|
|
189
|
+
the batch type, and during inheritance you can establish increasingly tight
|
|
190
|
+
upper bounds on the batch type to accomplish precisely the desired effect.
|
|
191
|
+
That is, you can formally require subclasses handle new batch attributes
|
|
192
|
+
without breaking LSP. This could look like:
|
|
193
|
+
|
|
194
|
+
```py
|
|
195
|
+
class A[Kw: BatchKwargs](Estimator):
|
|
196
|
+
def forward(self, batch: Kw): ...
|
|
197
|
+
|
|
198
|
+
class B[Kw: ExtraBatchKwargs](A[Kw]):
|
|
199
|
+
def forward(self, batch: Kw): ...
|
|
200
|
+
```
|
|
201
|
+
|
|
202
|
+
This formally allows `B` to inherit from `A` while depending on new,
|
|
203
|
+
non-optional attributes in its batch type; you can achieve that with refined
|
|
204
|
+
type upper bound in a generic subtype. So this pattern works *now*, but what
|
|
205
|
+
doesn't work is further *unpacking* `Kw` in this pattern into kwargs. That
|
|
206
|
+
is,
|
|
207
|
+
|
|
208
|
+
```py
|
|
209
|
+
class A[Kw: BatchKwargs](Estimator):
|
|
210
|
+
def forward(self, **kwargs: Unpack[Kw]): ...
|
|
211
|
+
|
|
212
|
+
class B[Kw: ExtraBatchKwargs](A[Kw]):
|
|
213
|
+
def forward(self, **kwargs: Unpack[Kw]): ...
|
|
214
|
+
```
|
|
215
|
+
|
|
216
|
+
As sound as this is (and it *is* sound, described in PEP 827), `Unpack` isn't
|
|
217
|
+
supported for type variables; it has to be used directly with `TypedDict`
|
|
218
|
+
references. I have the latter in place across the package anticipating
|
|
219
|
+
eventual support for this pattern, but things are subject to change. I may
|
|
220
|
+
simply toggle back to the currently supported pattern provided it's very
|
|
221
|
+
similar, and if/when PEP 827 is integrated, the change will be as simply as
|
|
222
|
+
adding back the `Unpack` statements and keyword expansions.
|
|
223
|
+
- At the time of writing, `trainlib` attempts to remain compatible with Python
|
|
224
|
+
3.12. This means the convenient lazy evaluation of type annotations can't be
|
|
225
|
+
fully embraced across the package, but we use `__future__` where possible to
|
|
226
|
+
make the transition easy when eventually bumping to 3.14+ (since the imports
|
|
227
|
+
can just be removed at that time). For instance, class definitions that need
|
|
228
|
+
to reference the class type itself shouldn't use string annotations but
|
|
229
|
+
instead direct type references:
|
|
230
|
+
|
|
231
|
+
```py
|
|
232
|
+
# not
|
|
233
|
+
class RecordIndex:
|
|
234
|
+
@overload
|
|
235
|
+
def __getitem__(self, key: slice) -> "RecordIndex": ...
|
|
236
|
+
|
|
237
|
+
# but
|
|
238
|
+
from __future__ import annotations
|
|
239
|
+
|
|
240
|
+
class RecordIndex:
|
|
241
|
+
@overload
|
|
242
|
+
def __getitem__(self, key: slice) -> RecordIndex: ...
|
|
243
|
+
```
|
trainlib-0.6.1/README.md
ADDED
|
@@ -0,0 +1,206 @@
|
|
|
1
|
+
# Overview
|
|
2
|
+
Minimal framework for ML modeling, supporting advanced dataset operations and
|
|
3
|
+
streamlined training workflows.
|
|
4
|
+
|
|
5
|
+
# Install
|
|
6
|
+
The `trainlib` package can be installed from PyPI:
|
|
7
|
+
|
|
8
|
+
```sh
|
|
9
|
+
pip install trainlib
|
|
10
|
+
```
|
|
11
|
+
|
|
12
|
+
# Development
|
|
13
|
+
- Initialize/synchronize the project with `uv sync`, creating a virtual
|
|
14
|
+
environment with base package dependencies.
|
|
15
|
+
- Depending on needs, install the development dependencies with `uv sync
|
|
16
|
+
--extra dev`.
|
|
17
|
+
|
|
18
|
+
## Testing
|
|
19
|
+
- To run the unit tests, make sure to first have the test dependencies
|
|
20
|
+
installed with `uv sync --extra test`, then run `make test`.
|
|
21
|
+
- For notebook testing, run `make install-kernel` to make the environment
|
|
22
|
+
available as a Jupyter kernel (to be selected when running notebooks).
|
|
23
|
+
|
|
24
|
+
## Documentation
|
|
25
|
+
- Install the documentation dependencies with `uv sync --extra doc`.
|
|
26
|
+
- Run `make docs-build` (optionally preceded by `make docs-clean`), and serve
|
|
27
|
+
locally with `make docs-serve`.
|
|
28
|
+
|
|
29
|
+
# Development remarks
|
|
30
|
+
- Across `Trainer` / `Estimator` / `Dataset`, I've considered a
|
|
31
|
+
`ParamSpec`-based typing scheme to better orchestrate alignment in the
|
|
32
|
+
`Trainer.train()` loop, e.g., so we can statically check whether a dataset
|
|
33
|
+
appears to be fulfilling the argument requirements for the estimator's
|
|
34
|
+
`forward()` / `metrics()` methods. Something like
|
|
35
|
+
|
|
36
|
+
```py
|
|
37
|
+
class Estimator[**P](nn.Module):
|
|
38
|
+
def forward(
|
|
39
|
+
self,
|
|
40
|
+
input: Tensor,
|
|
41
|
+
*args: P.args,
|
|
42
|
+
**kwargs: P.kwargs,
|
|
43
|
+
) -> Tensor:
|
|
44
|
+
...
|
|
45
|
+
|
|
46
|
+
class Trainer[**P]:
|
|
47
|
+
def __init__(
|
|
48
|
+
self,
|
|
49
|
+
estimator: Estimator[P],
|
|
50
|
+
...
|
|
51
|
+
): ...
|
|
52
|
+
```
|
|
53
|
+
|
|
54
|
+
might be how we begin threading signatures. But ensuring dataset items can
|
|
55
|
+
match `P` is challenging. You can consider a "packed" object where we
|
|
56
|
+
obfuscate passing data through `P`-signatures:
|
|
57
|
+
|
|
58
|
+
```py
|
|
59
|
+
class PackedItem[**P]:
|
|
60
|
+
def __init__(self, *args: P.args, **kwargs: P.kwargs) -> None:
|
|
61
|
+
self._args = args
|
|
62
|
+
self._kwargs = kwargs
|
|
63
|
+
|
|
64
|
+
def apply[R](self, func: Callable[P, R]) -> R:
|
|
65
|
+
return func(*self._args, **self._kwargs)
|
|
66
|
+
|
|
67
|
+
|
|
68
|
+
class BatchedDataset[U, R, I, **P](Dataset):
|
|
69
|
+
@abstractmethod
|
|
70
|
+
def _process_item_data(
|
|
71
|
+
self,
|
|
72
|
+
item_data: I,
|
|
73
|
+
item_index: int,
|
|
74
|
+
) -> PackedItem[P]:
|
|
75
|
+
...
|
|
76
|
+
|
|
77
|
+
def __iter__(self) -> Iterator[PackedItem[P]]:
|
|
78
|
+
...
|
|
79
|
+
```
|
|
80
|
+
|
|
81
|
+
Meaningfully shaping those signatures is what remains, but you can't really
|
|
82
|
+
do this, not with typical type expression flexibility. For instance, if I'm
|
|
83
|
+
trying to appropriately type my base `TupleDataset`:
|
|
84
|
+
|
|
85
|
+
```py
|
|
86
|
+
class SequenceDataset[I, **P](HomogenousDataset[int, I, I, P]):
|
|
87
|
+
...
|
|
88
|
+
|
|
89
|
+
class TupleDataset[I](SequenceDataset[tuple[I, ...], "?"]):
|
|
90
|
+
...
|
|
91
|
+
```
|
|
92
|
+
|
|
93
|
+
Here there's no way for me to shape a `ParamSpec` to indicate arbitrarily
|
|
94
|
+
many arguments of a fixed type (`I` in this case) to allow me to unpack my
|
|
95
|
+
item tuples into an appropriate `PackedItem` (since `ParamSpec` isn't exactly
|
|
96
|
+
controllable in that way; it's meant make a generic signature expressible as
|
|
97
|
+
a type variable, *not to express particular shapes* of the types in that
|
|
98
|
+
signature).
|
|
99
|
+
|
|
100
|
+
Until this (among other issues) becomes clearer, I'm setting up around a
|
|
101
|
+
simpler `TypedDict` type variable. We won't have particularly strong static
|
|
102
|
+
checks for item alignment inside `Trainer`, but this seems about as good as I
|
|
103
|
+
can get around the current infrastructure.
|
|
104
|
+
|
|
105
|
+
## Typing planning
|
|
106
|
+
There are currently two package-wide typing strategies geared toward future
|
|
107
|
+
Python version bumps and type system changes:
|
|
108
|
+
|
|
109
|
+
- Using `Unpack` with a type variable with a `TypedDict` upper bound, rather
|
|
110
|
+
than with a `TypedDict` directly. This is commented on extensively in
|
|
111
|
+
`trainlib/estimator.py`, but the short version is we want our key objects to
|
|
112
|
+
remain generic in their method signatures when they are expected to operate
|
|
113
|
+
on session-specific structured batch data. For instance, you can imagine an
|
|
114
|
+
`Estimator`'s `forward` method accepting a typed batch like
|
|
115
|
+
|
|
116
|
+
```py
|
|
117
|
+
def forward(self, batch: BatchKwargs): ...
|
|
118
|
+
|
|
119
|
+
# or
|
|
120
|
+
|
|
121
|
+
def forward(self, **kwargs: Unpack[Batchwargs]): ...
|
|
122
|
+
```
|
|
123
|
+
|
|
124
|
+
These are both currently supported patterns. Problems arise, however, when we
|
|
125
|
+
want an inherited class to specifically operate on an inherited type of
|
|
126
|
+
`BatchKwargs`
|
|
127
|
+
|
|
128
|
+
```py
|
|
129
|
+
class BatchKwargs(TypedDict):
|
|
130
|
+
inputs: Tensor
|
|
131
|
+
|
|
132
|
+
class ExtraBatchKwargs(BatchKwargs):
|
|
133
|
+
inputs: Tensor
|
|
134
|
+
labels: Tensor
|
|
135
|
+
|
|
136
|
+
class A(Estimator):
|
|
137
|
+
def forward(self, batch: BatchKwargs): ...
|
|
138
|
+
|
|
139
|
+
# want to require Estimator subclass handles both inputs and labels
|
|
140
|
+
class B(A):
|
|
141
|
+
def forward(self, batch: ExtraBatchKwargs): ...
|
|
142
|
+
```
|
|
143
|
+
|
|
144
|
+
Note how `ExtraBatchKwargs` doesn't use `NotRequired` for the new `labels`
|
|
145
|
+
attribute; the intention here is to actually require that `B` should handle
|
|
146
|
+
*more attributes* than `A`. **That breaks LSP**, so this doesn't work. You'd
|
|
147
|
+
need to either break out the `Estimator` hierarchies, or decide that `B`
|
|
148
|
+
could have inputs typed as `BatchKwargs` and possibly implement some extra
|
|
149
|
+
handling to check if batches of type `ExtraBatchKwargs` are passed in (note
|
|
150
|
+
that passing in a batch of type `ExtraBatchKwargs` to `A.forward` is
|
|
151
|
+
perfectly fine) and supply labels. **Or...** you make the class generic in
|
|
152
|
+
the batch type, and during inheritance you can establish increasingly tight
|
|
153
|
+
upper bounds on the batch type to accomplish precisely the desired effect.
|
|
154
|
+
That is, you can formally require subclasses handle new batch attributes
|
|
155
|
+
without breaking LSP. This could look like:
|
|
156
|
+
|
|
157
|
+
```py
|
|
158
|
+
class A[Kw: BatchKwargs](Estimator):
|
|
159
|
+
def forward(self, batch: Kw): ...
|
|
160
|
+
|
|
161
|
+
class B[Kw: ExtraBatchKwargs](A[Kw]):
|
|
162
|
+
def forward(self, batch: Kw): ...
|
|
163
|
+
```
|
|
164
|
+
|
|
165
|
+
This formally allows `B` to inherit from `A` while depending on new,
|
|
166
|
+
non-optional attributes in its batch type; you can achieve that with refined
|
|
167
|
+
type upper bound in a generic subtype. So this pattern works *now*, but what
|
|
168
|
+
doesn't work is further *unpacking* `Kw` in this pattern into kwargs. That
|
|
169
|
+
is,
|
|
170
|
+
|
|
171
|
+
```py
|
|
172
|
+
class A[Kw: BatchKwargs](Estimator):
|
|
173
|
+
def forward(self, **kwargs: Unpack[Kw]): ...
|
|
174
|
+
|
|
175
|
+
class B[Kw: ExtraBatchKwargs](A[Kw]):
|
|
176
|
+
def forward(self, **kwargs: Unpack[Kw]): ...
|
|
177
|
+
```
|
|
178
|
+
|
|
179
|
+
As sound as this is (and it *is* sound, described in PEP 827), `Unpack` isn't
|
|
180
|
+
supported for type variables; it has to be used directly with `TypedDict`
|
|
181
|
+
references. I have the latter in place across the package anticipating
|
|
182
|
+
eventual support for this pattern, but things are subject to change. I may
|
|
183
|
+
simply toggle back to the currently supported pattern provided it's very
|
|
184
|
+
similar, and if/when PEP 827 is integrated, the change will be as simply as
|
|
185
|
+
adding back the `Unpack` statements and keyword expansions.
|
|
186
|
+
- At the time of writing, `trainlib` attempts to remain compatible with Python
|
|
187
|
+
3.12. This means the convenient lazy evaluation of type annotations can't be
|
|
188
|
+
fully embraced across the package, but we use `__future__` where possible to
|
|
189
|
+
make the transition easy when eventually bumping to 3.14+ (since the imports
|
|
190
|
+
can just be removed at that time). For instance, class definitions that need
|
|
191
|
+
to reference the class type itself shouldn't use string annotations but
|
|
192
|
+
instead direct type references:
|
|
193
|
+
|
|
194
|
+
```py
|
|
195
|
+
# not
|
|
196
|
+
class RecordIndex:
|
|
197
|
+
@overload
|
|
198
|
+
def __getitem__(self, key: slice) -> "RecordIndex": ...
|
|
199
|
+
|
|
200
|
+
# but
|
|
201
|
+
from __future__ import annotations
|
|
202
|
+
|
|
203
|
+
class RecordIndex:
|
|
204
|
+
@overload
|
|
205
|
+
def __getitem__(self, key: slice) -> RecordIndex: ...
|
|
206
|
+
```
|
|
@@ -4,12 +4,10 @@ build-backend = "setuptools.build_meta"
|
|
|
4
4
|
|
|
5
5
|
[project]
|
|
6
6
|
name = "trainlib"
|
|
7
|
-
version = "0.
|
|
7
|
+
version = "0.6.1"
|
|
8
8
|
description = "Minimal framework for ML modeling. Supports advanced dataset operations and streamlined training."
|
|
9
|
-
requires-python = ">=3.
|
|
10
|
-
authors = [
|
|
11
|
-
{ name="Sam Griesemer", email="git@olog.io" },
|
|
12
|
-
]
|
|
9
|
+
requires-python = ">=3.12"
|
|
10
|
+
authors = []
|
|
13
11
|
readme = "README.md"
|
|
14
12
|
license = "MIT"
|
|
15
13
|
keywords = [
|
|
@@ -25,15 +23,18 @@ classifiers = [
|
|
|
25
23
|
]
|
|
26
24
|
dependencies = [
|
|
27
25
|
"torch",
|
|
26
|
+
"pandas",
|
|
28
27
|
"colorama>=0.4.6",
|
|
29
28
|
"matplotlib>=3.10.8",
|
|
30
29
|
"numpy>=2.4.1",
|
|
31
30
|
"tensorboard>=2.20.0",
|
|
32
31
|
"tqdm>=4.67.1",
|
|
32
|
+
"rich>=13.7",
|
|
33
33
|
"setuptools<=81.0.0", # here currently b/c tensorboard breaks @ v82.0.0
|
|
34
34
|
]
|
|
35
35
|
|
|
36
36
|
[project.scripts]
|
|
37
|
+
trainlib = "trainlib.__main__:main"
|
|
37
38
|
symconf = "trainlib.__main__:main"
|
|
38
39
|
|
|
39
40
|
[project.optional-dependencies]
|
|
@@ -61,6 +62,10 @@ Issues = "https://git.olog.io/olog/trainlib/issues"
|
|
|
61
62
|
[tool.setuptools.packages.find]
|
|
62
63
|
include = ["trainlib*"]
|
|
63
64
|
|
|
65
|
+
[tool.pytest.ini_options]
|
|
66
|
+
testpaths = ["test"]
|
|
67
|
+
addopts = "--import-mode=importlib"
|
|
68
|
+
|
|
64
69
|
# for static data files under package root
|
|
65
70
|
# [tool.setuptools.package-data]
|
|
66
71
|
# "<package>" = ["data/*.toml"]
|
|
@@ -68,8 +73,13 @@ include = ["trainlib*"]
|
|
|
68
73
|
[tool.ruff]
|
|
69
74
|
line-length = 79
|
|
70
75
|
|
|
76
|
+
[tool.ruff.format]
|
|
77
|
+
quote-style = "double"
|
|
78
|
+
indent-style = "space"
|
|
79
|
+
docstring-code-format = true
|
|
80
|
+
|
|
71
81
|
[tool.ruff.lint]
|
|
72
|
-
select = ["ANN", "E", "F", "UP", "B", "SIM", "I", "C4", "PERF"]
|
|
82
|
+
select = ["ANN", "D", "E", "F", "UP", "B", "SIM", "I", "C4", "PERF"]
|
|
73
83
|
|
|
74
84
|
[tool.ruff.lint.isort]
|
|
75
85
|
length-sort = true
|
|
@@ -77,13 +87,14 @@ order-by-type = false
|
|
|
77
87
|
force-sort-within-sections = false
|
|
78
88
|
|
|
79
89
|
[tool.ruff.lint.per-file-ignores]
|
|
80
|
-
"
|
|
90
|
+
"test/**" = ["S101"]
|
|
81
91
|
"**/__init__.py" = ["F401"]
|
|
82
92
|
|
|
83
|
-
[tool.ruff.
|
|
84
|
-
|
|
85
|
-
|
|
86
|
-
|
|
93
|
+
[tool.ruff.lint.pydocstyle]
|
|
94
|
+
convention = "google"
|
|
95
|
+
|
|
96
|
+
[tool.setuptools.package-data]
|
|
97
|
+
"trainlib.utils" = ["custom.mplstyle"]
|
|
87
98
|
|
|
88
99
|
[tool.uv.sources]
|
|
89
100
|
torch = { index = "pytorch" }
|
|
@@ -0,0 +1,160 @@
|
|
|
1
|
+
import json
|
|
2
|
+
from copy import deepcopy
|
|
3
|
+
from pathlib import Path
|
|
4
|
+
|
|
5
|
+
import pytest
|
|
6
|
+
|
|
7
|
+
from trainlib.cli import train, experiment, create_parser
|
|
8
|
+
from trainlib.run import expand_runs
|
|
9
|
+
from trainlib.trainer import Trainer
|
|
10
|
+
from trainlib.cli.experiment import ExperimentConfig
|
|
11
|
+
from trainlib.experiments.train import TrainExperiment
|
|
12
|
+
|
|
13
|
+
EXAMPLES = Path(__file__).resolve().parents[1] / "example" / "cli"
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
def test_cli_dispatches_non_training_experiment(tmp_path, monkeypatch):
|
|
17
|
+
from trainlib.run import Run
|
|
18
|
+
from trainlib.record import RecordIndex
|
|
19
|
+
from trainlib.experiment import Experiment
|
|
20
|
+
from trainlib.experiments import experiment_map
|
|
21
|
+
|
|
22
|
+
class Evaluation(Experiment[Run]):
|
|
23
|
+
def _execute(self, spec, directory, *, progress, console, output):
|
|
24
|
+
self.runstore.save(
|
|
25
|
+
directory, RecordIndex(), {"test": {"score": 1.0}}
|
|
26
|
+
)
|
|
27
|
+
|
|
28
|
+
monkeypatch.setitem(experiment_map, "test_evaluation", Evaluation)
|
|
29
|
+
result = experiment.run(
|
|
30
|
+
{
|
|
31
|
+
"experiment": "test_evaluation",
|
|
32
|
+
"output": str(tmp_path / "eval"),
|
|
33
|
+
"groups": {"raw": {}},
|
|
34
|
+
"seeds": [7],
|
|
35
|
+
"summary": {"plots": False, "resamples": 10},
|
|
36
|
+
},
|
|
37
|
+
quiet=True,
|
|
38
|
+
)
|
|
39
|
+
assert len(result["completed"]) == 1
|
|
40
|
+
assert "summary" in result
|
|
41
|
+
|
|
42
|
+
|
|
43
|
+
def test_single_run_cli_loads_and_trains_without_experiment(
|
|
44
|
+
tmp_path: Path,
|
|
45
|
+
) -> None:
|
|
46
|
+
config = json.loads((EXAMPLES / "run.json").read_text())
|
|
47
|
+
config["data"]["dataset"]["kwargs"]["num_samples"] = 12
|
|
48
|
+
config["estimator"]["kwargs"]["verbose"] = False
|
|
49
|
+
config["reporter"]["kwargs"].update(
|
|
50
|
+
progress=False,
|
|
51
|
+
console=False,
|
|
52
|
+
tensorboard=False,
|
|
53
|
+
)
|
|
54
|
+
config.update(max_epochs=1, chkpt_every=1)
|
|
55
|
+
path = tmp_path / "train.json"
|
|
56
|
+
path.write_text(json.dumps(config))
|
|
57
|
+
output = tmp_path / "single"
|
|
58
|
+
|
|
59
|
+
trainer = train.run_from_json(
|
|
60
|
+
parameters_file=str(path),
|
|
61
|
+
load_only=True,
|
|
62
|
+
output=output,
|
|
63
|
+
)
|
|
64
|
+
assert isinstance(trainer, Trainer)
|
|
65
|
+
assert not len(trainer.record_index)
|
|
66
|
+
assert not (output / "checkpoints").exists()
|
|
67
|
+
|
|
68
|
+
args = create_parser().parse_args(
|
|
69
|
+
[
|
|
70
|
+
"train",
|
|
71
|
+
"--parameters-file",
|
|
72
|
+
str(path),
|
|
73
|
+
"--output",
|
|
74
|
+
str(output),
|
|
75
|
+
]
|
|
76
|
+
)
|
|
77
|
+
args.func(args)
|
|
78
|
+
|
|
79
|
+
assert (output / "config.json").is_file()
|
|
80
|
+
assert list((output / "checkpoints").rglob("*.pth"))
|
|
81
|
+
assert not (output / "experiment.json").exists()
|
|
82
|
+
assert not (output / "runs").exists()
|
|
83
|
+
|
|
84
|
+
|
|
85
|
+
def test_group_seed_expansion_matches_explicit_runs_and_extends_stably(
|
|
86
|
+
tmp_path: Path,
|
|
87
|
+
) -> None:
|
|
88
|
+
config = json.loads((EXAMPLES / "experiment.json").read_text())
|
|
89
|
+
original = deepcopy(config)
|
|
90
|
+
runs = expand_runs(config)
|
|
91
|
+
assert len(runs) == 9
|
|
92
|
+
assert list(runs)[:3] == ["mlp-small-1", "mlp-small-2", "mlp-small-3"]
|
|
93
|
+
assert runs["mlp-medium-2"]["group"] == "mlp-medium"
|
|
94
|
+
assert runs["mlp-medium-2"]["seed"] == 2
|
|
95
|
+
assert config == original
|
|
96
|
+
|
|
97
|
+
initial = TrainExperiment(tmp_path, runs, common=config["common"])
|
|
98
|
+
assert expand_runs({"runs": runs}) == runs
|
|
99
|
+
config["seeds"].append(4)
|
|
100
|
+
extended = TrainExperiment(
|
|
101
|
+
tmp_path,
|
|
102
|
+
expand_runs(config),
|
|
103
|
+
common=config["common"],
|
|
104
|
+
)
|
|
105
|
+
assert len(extended.runs) == 12
|
|
106
|
+
for name, spec in initial.runs.items():
|
|
107
|
+
assert extended.runs[name] == spec
|
|
108
|
+
|
|
109
|
+
runs["mlp-small-1"]["estimator"]["kwargs"]["hidden_dims"].append(99)
|
|
110
|
+
assert runs["mlp-small-2"]["estimator"]["kwargs"]["hidden_dims"] == [8]
|
|
111
|
+
assert original["groups"] == config["groups"]
|
|
112
|
+
|
|
113
|
+
|
|
114
|
+
@pytest.mark.parametrize(
|
|
115
|
+
"config",
|
|
116
|
+
[
|
|
117
|
+
{"runs": {"one": {}}, "groups": {"one": {}}, "seeds": [1]},
|
|
118
|
+
{"groups": {"one": {}}},
|
|
119
|
+
{"groups": {"one": {}}, "seeds": [1, 1]},
|
|
120
|
+
],
|
|
121
|
+
)
|
|
122
|
+
def test_experiment_rejects_ambiguous_or_incomplete_expansion(
|
|
123
|
+
config: ExperimentConfig,
|
|
124
|
+
) -> None:
|
|
125
|
+
with pytest.raises(ValueError):
|
|
126
|
+
expand_runs(config)
|
|
127
|
+
|
|
128
|
+
|
|
129
|
+
def test_compact_experiment_cli_writes_expanded_manifest_and_analysis(
|
|
130
|
+
tmp_path: Path,
|
|
131
|
+
) -> None:
|
|
132
|
+
config = json.loads((EXAMPLES / "experiment.json").read_text())
|
|
133
|
+
config["groups"] = {"small": config["groups"]["mlp-small"]}
|
|
134
|
+
config["seeds"] = [1]
|
|
135
|
+
config["comparisons"] = {}
|
|
136
|
+
common = config["common"]
|
|
137
|
+
common["data"]["dataset"]["kwargs"]["num_samples"] = 12
|
|
138
|
+
common["reporter"]["kwargs"].update(tensorboard=False)
|
|
139
|
+
common.update(max_epochs=1, trace_graph=False)
|
|
140
|
+
config["summary"].update(plots=False, resamples=10)
|
|
141
|
+
path = tmp_path / "experiment.json"
|
|
142
|
+
path.write_text(json.dumps(config))
|
|
143
|
+
output = tmp_path / "experiment"
|
|
144
|
+
|
|
145
|
+
result = experiment.run_from_json(
|
|
146
|
+
parameters_file=str(path),
|
|
147
|
+
output=output,
|
|
148
|
+
quiet=True,
|
|
149
|
+
)
|
|
150
|
+
|
|
151
|
+
assert len(result["completed"]) == 1
|
|
152
|
+
assert result["failures"] == {}
|
|
153
|
+
assert result["summary"]["rows"]
|
|
154
|
+
manifest = json.loads((output / "experiment.json").read_text())
|
|
155
|
+
(spec,) = manifest["runs"].values()
|
|
156
|
+
assert spec["name"] == "small-1"
|
|
157
|
+
assert spec["group"] == "small"
|
|
158
|
+
assert spec["seed"] == 1
|
|
159
|
+
assert list(output.glob("runs/*/attempt-*/state.pt"))
|
|
160
|
+
assert (output / "analysis" / "summary.json").is_file()
|