mltrainer 0.3.0__tar.gz → 0.3.2__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.
Files changed (40) hide show
  1. mltrainer-0.3.2/.claude/settings.local.json +7 -0
  2. mltrainer-0.3.2/.remember/.gitignore +1 -0
  3. mltrainer-0.3.2/.remember/archive.md +3 -0
  4. mltrainer-0.3.2/.remember/recent.md +8 -0
  5. mltrainer-0.3.2/.remember/tmp/last-ndc.ts +1 -0
  6. mltrainer-0.3.2/.remember/tmp/last-save-ts +1 -0
  7. mltrainer-0.3.2/.remember/tmp/last-save.json +1 -0
  8. mltrainer-0.3.2/.remember/today-2026-07-20.done.md +2 -0
  9. mltrainer-0.3.2/Makefile +12 -0
  10. {mltrainer-0.3.0 → mltrainer-0.3.2}/PKG-INFO +2 -1
  11. {mltrainer-0.3.0 → mltrainer-0.3.2}/pyproject.toml +2 -2
  12. {mltrainer-0.3.0 → mltrainer-0.3.2}/src/mltrainer/__init__.py +1 -1
  13. {mltrainer-0.3.0 → mltrainer-0.3.2}/src/mltrainer/metrics.py +7 -2
  14. mltrainer-0.3.2/src/mltrainer/settings.py +61 -0
  15. {mltrainer-0.3.0 → mltrainer-0.3.2}/src/mltrainer/trainer.py +184 -67
  16. {mltrainer-0.3.0 → mltrainer-0.3.2}/src/mltrainer/transformer.py +3 -0
  17. mltrainer-0.3.2/tests/test_trainer_iteration.py +142 -0
  18. mltrainer-0.3.2/tests/utils/dummy_model.py +0 -0
  19. {mltrainer-0.3.0 → mltrainer-0.3.2}/uv.lock +1 -1
  20. mltrainer-0.3.0/Makefile +0 -7
  21. mltrainer-0.3.0/src/mltrainer/settings.py +0 -177
  22. {mltrainer-0.3.0 → mltrainer-0.3.2}/.gitignore +0 -0
  23. {mltrainer-0.3.0 → mltrainer-0.3.2}/.lefthook.yml +0 -0
  24. /mltrainer-0.3.0/tests/__init__.py → /mltrainer-0.3.2/.remember/now.md +0 -0
  25. {mltrainer-0.3.0 → mltrainer-0.3.2}/README.md +0 -0
  26. {mltrainer-0.3.0 → mltrainer-0.3.2}/src/mltrainer/imagemodels.py +0 -0
  27. {mltrainer-0.3.0 → mltrainer-0.3.2}/src/mltrainer/preprocessors.py +0 -0
  28. {mltrainer-0.3.0 → mltrainer-0.3.2}/src/mltrainer/rnn_models.py +0 -0
  29. {mltrainer-0.3.0 → mltrainer-0.3.2}/src/mltrainer/vae.py +0 -0
  30. {mltrainer-0.3.0/tests/unit → mltrainer-0.3.2/tests}/__init__.py +0 -0
  31. {mltrainer-0.3.0 → mltrainer-0.3.2}/tests/conftest.py +0 -0
  32. {mltrainer-0.3.0 → mltrainer-0.3.2}/tests/test_metrics.py +0 -0
  33. {mltrainer-0.3.0 → mltrainer-0.3.2}/tests/test_preprocessors.py +0 -0
  34. {mltrainer-0.3.0 → mltrainer-0.3.2}/tests/test_trainer.py +0 -0
  35. {mltrainer-0.3.0 → mltrainer-0.3.2}/tests/test_trainer_devices.py +0 -0
  36. {mltrainer-0.3.0/tests/utils → mltrainer-0.3.2/tests/unit}/__init__.py +0 -0
  37. {mltrainer-0.3.0 → mltrainer-0.3.2}/tests/unit/test_trainersettings.py +0 -0
  38. /mltrainer-0.3.0/tests/utils/dummy_data.py → /mltrainer-0.3.2/tests/utils/__init__.py +0 -0
  39. /mltrainer-0.3.0/tests/utils/dummy_model.py → /mltrainer-0.3.2/tests/utils/dummy_data.py +0 -0
  40. {mltrainer-0.3.0 → mltrainer-0.3.2}/tests/utils/dummy_metrics.py +0 -0
@@ -0,0 +1,7 @@
1
+ {
2
+ "permissions": {
3
+ "allow": [
4
+ "Bash(uv run *)"
5
+ ]
6
+ }
7
+ }
@@ -0,0 +1 @@
1
+ *
@@ -0,0 +1,3 @@
1
+ # Archive
2
+
3
+ ```
@@ -0,0 +1,8 @@
1
+ # Recent
2
+
3
+ ```
4
+
5
+ # Recent
6
+
7
+ ## 2026-07-20
8
+ Refactored mltrainer.report() with Reporter ABC pattern (RayReporter/MLflow/TensorBoard) and added type hints w/ builtin generators, _to_device helper, plus type checking fixes. Cleaned ~115 LOC commented code and updated Makefile (ruff/ty). 47 tests passing—ready v0.3.1.
@@ -0,0 +1 @@
1
+ 1784568708
@@ -0,0 +1 @@
1
+ 1784568694
@@ -0,0 +1 @@
1
+ {"session": "284374ad-dc46-496d-a542-5d141c54654d", "line": 204}
@@ -0,0 +1,2 @@
1
+ ## 19:31 | main
2
+ Refactored mltrainer.report() w/ Reporter ABC (RayReporter/MLflow/TensorBoard); type hints→builtin gen; added _to_device helper; fixed ty checks (transformer.py/trainer.py); Makefile: ruff/ty; cleaned ~115 commented LOC (settings.py); 47 tests pass; ready v0.3.1
@@ -0,0 +1,12 @@
1
+ format:
2
+ ruff format src
3
+ ruff check --fix src
4
+
5
+ lint:
6
+ ruff check src
7
+ ty check src
8
+
9
+ test:
10
+ pytest
11
+
12
+ .PHONY: format lint test
@@ -1,7 +1,8 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: mltrainer
3
- Version: 0.3.0
3
+ Version: 0.3.2
4
4
  Summary: toolkit for training pytorch models
5
+ Project-URL: GitHub, https://github.com/raoulg/mltrainer
5
6
  Author-email: "R.Grouls" <Raoul.Grouls@han.nl>
6
7
  License: MIT
7
8
  Classifier: Development Status :: 4 - Beta
@@ -1,7 +1,7 @@
1
1
 
2
2
  [project]
3
3
  name = "mltrainer"
4
- version = "0.3.0"
4
+ version = "0.3.2"
5
5
  description = "toolkit for training pytorch models"
6
6
 
7
7
  authors = [
@@ -44,7 +44,7 @@ build-backend = "hatchling.build"
44
44
  [tool.hatch.build.targets.wheel]
45
45
  packages = ["src/mltrainer"]
46
46
 
47
- [tools.urls]
47
+ [project.urls]
48
48
  GitHub = "https://github.com/raoulg/mltrainer"
49
49
 
50
50
  [dependency-groups]
@@ -3,4 +3,4 @@ from mltrainer.trainer import Trainer
3
3
 
4
4
  __all__ = ["Trainer", "TrainerSettings", "ReportTypes"]
5
5
 
6
- __version__ = "0.2.7"
6
+ __version__ = "0.3.1"
@@ -1,7 +1,7 @@
1
1
  from __future__ import annotations
2
2
 
3
3
  from abc import ABC, abstractmethod
4
- from typing import Iterator, Union
4
+ from typing import Iterator, Literal, Union
5
5
 
6
6
  import numpy as np
7
7
  import torch
@@ -86,7 +86,12 @@ class Accuracy(Metric):
86
86
 
87
87
 
88
88
  class F1Score(Metric):
89
- def __init__(self, average="micro", threshold=0.5, epsilon=1e-7):
89
+ def __init__(
90
+ self,
91
+ average: Literal["micro", "macro"] = "micro",
92
+ threshold: float = 0.5,
93
+ epsilon: float = 1e-7,
94
+ ) -> None:
90
95
  self.average = average
91
96
  self.threshold = threshold
92
97
  self.epsilon = epsilon
@@ -0,0 +1,61 @@
1
+ from __future__ import annotations
2
+
3
+ from enum import Enum
4
+ from pathlib import Path
5
+ from typing import Any, Callable, Dict, List, Optional
6
+
7
+ from loguru import logger
8
+ from pydantic import BaseModel, ConfigDict, field_validator
9
+
10
+
11
+ class FileTypes(Enum):
12
+ JPG = ".jpg"
13
+ PNG = ".png"
14
+ TXT = ".txt"
15
+ ZIP = ".zip"
16
+ TGZ = ".tgz"
17
+ TAR = ".tar.gz"
18
+ GZ = ".gz"
19
+
20
+
21
+ class ReportTypes(Enum):
22
+ TOML = "TOML"
23
+ TENSORBOARD = "TENSORBOARD"
24
+ MLFLOW = "MLFLOW"
25
+ RAY = "RAY"
26
+
27
+
28
+ class FormattedBase(BaseModel):
29
+ def __str__(self) -> str:
30
+ return "\n".join(f"{k}: {v}" for k, v in self.__dict__.items())
31
+
32
+ def __repr__(self) -> str:
33
+ return "\n".join(f"{k}: {v}" for k, v in self.__dict__.items())
34
+
35
+
36
+ class TrainerSettings(FormattedBase):
37
+ epochs: int
38
+ metrics: List[Callable]
39
+ logdir: Path
40
+ train_steps: int
41
+ valid_steps: int
42
+ reporttypes: List[ReportTypes]
43
+ optimizer_kwargs: Dict[str, Any] = {"lr": 1e-3, "weight_decay": 1e-5}
44
+ scheduler_kwargs: Optional[Dict[str, Any]] = {"factor": 0.1, "patience": 10}
45
+ earlystop_kwargs: Optional[Dict[str, Any]] = {
46
+ "save": False,
47
+ "verbose": True,
48
+ "patience": 10,
49
+ }
50
+
51
+ model_config = ConfigDict(arbitrary_types_allowed=True)
52
+
53
+ @field_validator("logdir")
54
+ @classmethod
55
+ def check_path(cls, logdir: Path) -> Path: # noqa: N805
56
+ if isinstance(logdir, str):
57
+ logdir = Path(logdir)
58
+ if not logdir.resolve().exists():
59
+ logdir.mkdir(parents=True)
60
+ logger.info(f"Created logdir {logdir.absolute()}")
61
+ return logdir
@@ -1,6 +1,8 @@
1
+ from abc import ABC, abstractmethod
2
+ from collections.abc import Callable, Iterable, Iterator
1
3
  from datetime import datetime
2
4
  from pathlib import Path
3
- from typing import Callable, Dict, Iterator, Optional, Tuple, Type, TypeVar
5
+ from typing import TYPE_CHECKING, TypeVar
4
6
 
5
7
  import mlflow
6
8
  import torch
@@ -11,10 +13,17 @@ from tqdm import tqdm
11
13
 
12
14
  from mltrainer import ReportTypes, TrainerSettings
13
15
 
16
+ if TYPE_CHECKING:
17
+ from torch.utils.tensorboard.writer import SummaryWriter
18
+
14
19
  OptimizerType = TypeVar("OptimizerType", bound=torch.optim.Optimizer)
15
20
 
21
+ # A batch input is either a single tensor or a tuple of tensors (e.g. packed
22
+ # sequences); the target is always a tensor.
23
+ BatchTensor = torch.Tensor | tuple[torch.Tensor, ...]
24
+
16
25
 
17
- def dir_add_timestamp(log_dir: Optional[Path] = None) -> Path:
26
+ def dir_add_timestamp(log_dir: Path | None = None) -> Path:
18
27
  if log_dir is None:
19
28
  log_dir = Path(".")
20
29
  log_dir = Path(log_dir)
@@ -26,27 +35,112 @@ def dir_add_timestamp(log_dir: Optional[Path] = None) -> Path:
26
35
  return log_dir
27
36
 
28
37
 
38
+ class Reporter(ABC):
39
+ """Reports training metrics to a single backend.
40
+
41
+ A new backend is added by subclassing (open for extension); the training
42
+ loop that iterates over reporters never changes (closed for modification).
43
+ """
44
+
45
+ @abstractmethod
46
+ def report(
47
+ self,
48
+ epoch: int,
49
+ train_loss: float,
50
+ test_loss: float,
51
+ metric_dict: dict[str, float],
52
+ ) -> None: ...
53
+
54
+
55
+ class RayReporter(Reporter):
56
+ def __init__(self) -> None:
57
+ try:
58
+ import ray
59
+ except ModuleNotFoundError:
60
+ raise ModuleNotFoundError(
61
+ "ray is required for RAY reporting. "
62
+ "Install it with: uv add 'mltrainer[tune]'"
63
+ )
64
+ self._ray = ray
65
+
66
+ def report(
67
+ self,
68
+ epoch: int,
69
+ train_loss: float,
70
+ test_loss: float,
71
+ metric_dict: dict[str, float],
72
+ ) -> None:
73
+ self._ray.train.report( # ty: ignore[possibly-missing-attribute]
74
+ {
75
+ "iterations": epoch,
76
+ "train_loss": train_loss,
77
+ "test_loss": test_loss,
78
+ **metric_dict,
79
+ }
80
+ )
81
+
82
+
83
+ class MLflowReporter(Reporter):
84
+ def __init__(self, optimizer: Optimizer) -> None:
85
+ self.optimizer = optimizer
86
+
87
+ def report(
88
+ self,
89
+ epoch: int,
90
+ train_loss: float,
91
+ test_loss: float,
92
+ metric_dict: dict[str, float],
93
+ ) -> None:
94
+ mlflow.log_metric("Loss/train", train_loss, step=epoch)
95
+ mlflow.log_metric("Loss/test", test_loss, step=epoch)
96
+ for m in metric_dict:
97
+ mlflow.log_metric(f"metric/{m}", metric_dict[m], step=epoch)
98
+ lr = self.optimizer.param_groups[0]["lr"]
99
+ mlflow.log_metric("learning_rate", lr, step=epoch)
100
+
101
+
102
+ class TensorBoardReporter(Reporter):
103
+ def __init__(self, writer: "SummaryWriter", optimizer: Optimizer) -> None:
104
+ self.writer = writer
105
+ self.optimizer = optimizer
106
+
107
+ def report(
108
+ self,
109
+ epoch: int,
110
+ train_loss: float,
111
+ test_loss: float,
112
+ metric_dict: dict[str, float],
113
+ ) -> None:
114
+ self.writer.add_scalar("Loss/train", train_loss, epoch)
115
+ self.writer.add_scalar("Loss/test", test_loss, epoch)
116
+ for m in metric_dict:
117
+ self.writer.add_scalar(f"metric/{m}", metric_dict[m], epoch)
118
+ lr = self.optimizer.param_groups[0]["lr"]
119
+ self.writer.add_scalar("learning_rate", lr, epoch)
120
+
121
+
29
122
  class Trainer:
30
123
  def __init__(
31
124
  self,
32
125
  model: torch.nn.Module,
33
126
  settings: TrainerSettings,
34
127
  loss_fn: Callable,
35
- optimizer: Type[Optimizer],
36
- traindataloader: Iterator,
37
- validdataloader: Iterator,
38
- scheduler: Optional[Callable],
39
- device: Optional[str] = None,
128
+ optimizer: type[Optimizer],
129
+ traindataloader: Iterable,
130
+ validdataloader: Iterable,
131
+ scheduler: Callable | None,
132
+ device: str | None = None,
40
133
  ) -> None:
41
134
  self.model = model
42
135
  self.settings = settings
43
136
  self.log_dir = dir_add_timestamp(settings.logdir)
44
137
  self.loss_fn = loss_fn
45
- self.optimizer = optimizer
46
138
  self.traindataloader = traindataloader
47
139
  self.validdataloader = validdataloader
140
+ self._train_iter: Iterator | None = None
141
+ self._valid_iter: Iterator | None = None
48
142
  self.device = device
49
- self.writer = None
143
+ self.writer: "SummaryWriter | None" = None
50
144
 
51
145
  if self.device:
52
146
  self.model.to(self.device)
@@ -66,7 +160,7 @@ class Trainer:
66
160
  "Found earlystop_kwargs in settings."
67
161
  "Set to None if you dont want earlystopping."
68
162
  )
69
- self.early_stopping: Optional[EarlyStopping] = EarlyStopping(
163
+ self.early_stopping: EarlyStopping | None = EarlyStopping(
70
164
  self.log_dir, **settings.earlystop_kwargs
71
165
  )
72
166
  else:
@@ -81,6 +175,61 @@ class Trainer:
81
175
  TOMLSerializer.save(model, self.log_dir / "model.toml")
82
176
  TOMLSerializer.save(settings, self.log_dir / "settings.toml")
83
177
 
178
+ self.reporters = self._build_reporters()
179
+
180
+ def _build_reporters(self) -> list[Reporter]:
181
+ """Map the requested report types onto reporter objects once, so the
182
+ per-epoch report loop stays closed for modification."""
183
+ reporters: list[Reporter] = []
184
+ if ReportTypes.RAY in self.settings.reporttypes:
185
+ reporters.append(RayReporter())
186
+ if ReportTypes.MLFLOW in self.settings.reporttypes:
187
+ reporters.append(MLflowReporter(self.optimizer))
188
+ if ReportTypes.TENSORBOARD in self.settings.reporttypes:
189
+ assert self.writer is not None
190
+ reporters.append(TensorBoardReporter(self.writer, self.optimizer))
191
+ return reporters
192
+
193
+ def _to_device(
194
+ self, x: BatchTensor, y: torch.Tensor
195
+ ) -> tuple[BatchTensor, torch.Tensor]:
196
+ """Move a batch onto self.device, handling tuple inputs (e.g. packed
197
+ sequences) as well as single tensors."""
198
+ if not self.device:
199
+ return x, y
200
+ if isinstance(x, tuple):
201
+ x = tuple(xx.to(self.device) for xx in x)
202
+ else:
203
+ x = x.to(self.device)
204
+ return x, y.to(self.device)
205
+
206
+ def _next_batch(
207
+ self, loader: Iterable, cache: str
208
+ ) -> tuple[BatchTensor, torch.Tensor]:
209
+ """Pull one batch, keeping the iterator alive across steps.
210
+
211
+ `next(iter(loader))` per step is only correct for an *iterator* -- a
212
+ streamer -- because `iter()` on an iterator returns itself. A torch
213
+ DataLoader is an *iterable*: it hands out a fresh iterator (and, with
214
+ num_workers > 0, a fresh worker pool) on every call, so a training
215
+ loader would reshuffle and replay its first batch every step, and a
216
+ validation loader would score the same first batch `valid_steps` times.
217
+
218
+ Holding the iterator and re-`iter`ing it only when it is exhausted
219
+ works for both: an infinite streamer never raises StopIteration, and a
220
+ finite DataLoader wraps around into its next epoch.
221
+ """
222
+ it = getattr(self, cache)
223
+ if it is None:
224
+ it = iter(loader)
225
+ setattr(self, cache, it)
226
+ try:
227
+ return next(it)
228
+ except StopIteration:
229
+ it = iter(loader)
230
+ setattr(self, cache, it)
231
+ return next(it)
232
+
84
233
  def __del__(self):
85
234
  """Cleanup method to ensure proper resource handling"""
86
235
  if hasattr(self, "writer") and self.writer is not None:
@@ -114,39 +263,29 @@ class Trainer:
114
263
  train_loss: float = 0.0
115
264
  train_steps = self.settings.train_steps
116
265
  for _ in tqdm(range(train_steps), colour="#1e4706"):
117
- x, y = next(iter(self.traindataloader))
118
- if self.device:
119
- if isinstance(x, tuple):
120
- x = tuple(xx.to(self.device) for xx in x)
121
- y = y.to(self.device)
122
- else:
123
- x, y = x.to(self.device), y.to(self.device)
266
+ x, y = self._next_batch(self.traindataloader, "_train_iter")
267
+ x, y = self._to_device(x, y)
124
268
  self.optimizer.zero_grad()
125
269
  yhat = self.model(x)
126
270
  loss = self.loss_fn(yhat, y)
127
271
  loss.backward()
128
- self.optimizer.step() # type: ignore
272
+ self.optimizer.step()
129
273
  train_loss += loss.cpu().detach().numpy()
130
274
  train_loss /= train_steps
131
275
 
132
276
  return train_loss
133
277
 
134
- def evalbatches(self) -> Tuple[Dict[str, float], float]:
278
+ def evalbatches(self) -> tuple[dict[str, float], float]:
135
279
  """Evaluate model on validation data with proper device handling"""
136
280
  self.model.eval()
137
281
  valid_steps = self.settings.valid_steps
138
282
  test_loss: float = 0.0
139
- metric_dict: Dict[str, float] = {}
283
+ metric_dict: dict[str, float] = {}
140
284
 
141
285
  with torch.no_grad(): # Prevent gradient computation during evaluation
142
286
  for _ in range(valid_steps):
143
- x, y = next(iter(self.validdataloader))
144
- if self.device:
145
- if isinstance(x, tuple):
146
- x = tuple(xx.to(self.device) for xx in x)
147
- y = y.to(self.device)
148
- else:
149
- x, y = x.to(self.device), y.to(self.device)
287
+ x, y = self._next_batch(self.validdataloader, "_valid_iter")
288
+ x, y = self._to_device(x, y)
150
289
 
151
290
  # Forward pass
152
291
  yhat = self.model(x)
@@ -161,7 +300,7 @@ class Trainer:
161
300
  # Average the results
162
301
  test_loss /= valid_steps
163
302
  for key in metric_dict:
164
- metric_dict[str(key)] = metric_dict[str(key)] / valid_steps
303
+ metric_dict[key] = metric_dict[key] / valid_steps
165
304
 
166
305
  # Handle scheduler
167
306
  if self.scheduler:
@@ -172,54 +311,32 @@ class Trainer:
172
311
 
173
312
  return metric_dict, test_loss
174
313
 
175
- def report(
176
- self, epoch: int, train_loss: float, test_loss: float, metric_dict: Dict
177
- ) -> None:
178
- if (self.last_epoch != 0) and (epoch == 0):
314
+ def _resolve_epoch(self, epoch: int) -> int:
315
+ """Offset the epoch counter when resuming a previous training run."""
316
+ if self.last_epoch != 0 and epoch == 0:
179
317
  self.last_epoch += 1
180
318
  logger.info(f"Resuming epochs from previous training at {self.last_epoch}")
181
319
  if self.last_epoch != 0:
182
320
  epoch += self.last_epoch
183
- reporttypes = self.settings.reporttypes
184
- self.test_loss = test_loss
185
-
186
- if ReportTypes.RAY in reporttypes:
187
- try:
188
- import ray
189
- except ModuleNotFoundError:
190
- raise ModuleNotFoundError(
191
- "ray is required for RAY reporting. "
192
- "Install it with: uv add 'mltrainer[tune]'"
193
- )
194
- ray.train.report(
195
- {
196
- "iterations": epoch,
197
- "train_loss": train_loss,
198
- "test_loss": test_loss,
199
- **metric_dict,
200
- }
201
- )
321
+ return epoch
202
322
 
203
- if ReportTypes.MLFLOW in reporttypes:
204
- mlflow.log_metric("Loss/train", train_loss, step=epoch)
205
- mlflow.log_metric("Loss/test", test_loss, step=epoch)
206
- for m in metric_dict:
207
- mlflow.log_metric(f"metric/{m}", metric_dict[m], step=epoch)
208
- lr = [group["lr"] for group in self.optimizer.param_groups][0] # type: ignore
209
- mlflow.log_metric("learning_rate", lr, step=epoch)
323
+ def report(
324
+ self,
325
+ epoch: int,
326
+ train_loss: float,
327
+ test_loss: float,
328
+ metric_dict: dict[str, float],
329
+ ) -> None:
330
+ epoch = self._resolve_epoch(epoch)
331
+ self.test_loss = test_loss
210
332
 
211
- if ReportTypes.TENSORBOARD in reporttypes:
212
- assert self.writer is not None
213
- self.writer.add_scalar("Loss/train", train_loss, epoch)
214
- self.writer.add_scalar("Loss/test", test_loss, epoch)
215
- for m in metric_dict:
216
- self.writer.add_scalar(f"metric/{m}", metric_dict[m], epoch)
217
- lr = [group["lr"] for group in self.optimizer.param_groups][0] # type: ignore
218
- self.writer.add_scalar("learning_rate", lr, epoch)
333
+ for reporter in self.reporters:
334
+ reporter.report(epoch, train_loss, test_loss, metric_dict)
219
335
 
220
336
  metric_scores = [f"{v:.4f}" for v in metric_dict.values()]
221
337
  logger.info(
222
- f"Epoch {epoch} train {train_loss:.4f} test {test_loss:.4f} metric {metric_scores}" # noqa E501
338
+ f"Epoch {epoch} train {train_loss:.4f} test {test_loss:.4f} "
339
+ f"metric {metric_scores}"
223
340
  )
224
341
 
225
342
 
@@ -21,6 +21,9 @@ class PositionalEncoding(nn.Module):
21
21
  It adds positional information to the input embeddings.
22
22
  """
23
23
 
24
+ pe: Tensor # registered as a buffer in __init__; annotated so type checkers
25
+ # know it is a Tensor rather than a generic Module attribute.
26
+
24
27
  def __init__(self, d_model: int, dropout: float = 0.1, max_seq_len: int = 5000):
25
28
  super().__init__()
26
29
  self.dropout = nn.Dropout(p=dropout)
@@ -0,0 +1,142 @@
1
+ """Regression tests for the per-step iterator (see mltrainer P0.6).
2
+
3
+ `next(iter(loader))` inside the step loop is correct for an *iterator* but not
4
+ for an *iterable* like a torch DataLoader: it rebuilds the iterator every step,
5
+ so training replays a reshuffled first batch and validation scores the same
6
+ first batch `valid_steps` times.
7
+ """
8
+
9
+ import torch
10
+ import torch.optim as optim
11
+ from torch.utils.data import DataLoader, Dataset
12
+
13
+ from mltrainer import Trainer, TrainerSettings
14
+ from tests.utils.dummy_metrics import TestMSE
15
+
16
+
17
+ class CountingDataset(Dataset):
18
+ """Each row is its own index, so a batch is identifiable by its contents."""
19
+
20
+ def __init__(self, size: int = 64) -> None:
21
+ self.size = size
22
+
23
+ def __len__(self) -> int:
24
+ return self.size
25
+
26
+ def __getitem__(self, idx: int):
27
+ x = torch.full((10,), float(idx))
28
+ return x, x
29
+
30
+
31
+ def build_trainer(model, loss_fn, temp_logdir, train_loader, valid_loader, **kwargs):
32
+ settings = TrainerSettings(
33
+ epochs=1,
34
+ metrics=[TestMSE()],
35
+ logdir=temp_logdir,
36
+ reporttypes=[],
37
+ optimizer_kwargs={"lr": 0.01},
38
+ **kwargs,
39
+ )
40
+ return Trainer(
41
+ model=model,
42
+ settings=settings,
43
+ loss_fn=loss_fn,
44
+ optimizer=optim.Adam, # type: ignore
45
+ traindataloader=train_loader,
46
+ validdataloader=valid_loader,
47
+ scheduler=None,
48
+ device="cpu",
49
+ )
50
+
51
+
52
+ def test_valid_loader_walks_the_whole_split(model, loss_fn, temp_logdir):
53
+ """A shuffle=False DataLoader must yield a *different* batch per step."""
54
+ dataset = CountingDataset(size=64)
55
+ valid_loader = DataLoader(dataset, batch_size=16, shuffle=False)
56
+ trainer = build_trainer(
57
+ model,
58
+ loss_fn,
59
+ temp_logdir,
60
+ DataLoader(dataset, batch_size=16, shuffle=True),
61
+ valid_loader,
62
+ train_steps=1,
63
+ valid_steps=4,
64
+ )
65
+
66
+ seen = [
67
+ trainer._next_batch(valid_loader, "_valid_iter")[0][0, 0].item()
68
+ for _ in range(4)
69
+ ]
70
+ assert seen == [0.0, 16.0, 32.0, 48.0]
71
+
72
+
73
+ def test_iterator_wraps_around_at_the_end(model, loss_fn, temp_logdir):
74
+ """Asking for more steps than the loader holds restarts it, not crashes."""
75
+ dataset = CountingDataset(size=32)
76
+ valid_loader = DataLoader(dataset, batch_size=16, shuffle=False)
77
+ trainer = build_trainer(
78
+ model,
79
+ loss_fn,
80
+ temp_logdir,
81
+ DataLoader(dataset, batch_size=16, shuffle=True),
82
+ valid_loader,
83
+ train_steps=1,
84
+ valid_steps=3,
85
+ )
86
+
87
+ seen = [
88
+ trainer._next_batch(valid_loader, "_valid_iter")[0][0, 0].item()
89
+ for _ in range(3)
90
+ ]
91
+ assert seen == [0.0, 16.0, 0.0]
92
+
93
+
94
+ def test_plain_iterator_still_advances(model, loss_fn, temp_logdir):
95
+ """The old streamer path: an infinite generator, not a DataLoader."""
96
+
97
+ def stream():
98
+ i = 0
99
+ while True:
100
+ yield torch.full((4, 10), float(i)), torch.full((4, 10), float(i))
101
+ i += 1
102
+
103
+ streamer = stream()
104
+ trainer = build_trainer(
105
+ model,
106
+ loss_fn,
107
+ temp_logdir,
108
+ streamer,
109
+ streamer,
110
+ train_steps=1,
111
+ valid_steps=3,
112
+ )
113
+
114
+ seen = [
115
+ trainer._next_batch(streamer, "_valid_iter")[0][0, 0].item() for _ in range(3)
116
+ ]
117
+ assert seen == [0.0, 1.0, 2.0]
118
+
119
+
120
+ def test_validation_metrics_cover_more_than_one_batch(model, loss_fn, temp_logdir):
121
+ """End to end: evalbatches over a loader whose batches differ wildly must
122
+ not report the metric of a single batch."""
123
+ dataset = CountingDataset(size=64)
124
+ valid_loader = DataLoader(dataset, batch_size=16, shuffle=False)
125
+ trainer = build_trainer(
126
+ model,
127
+ loss_fn,
128
+ temp_logdir,
129
+ DataLoader(dataset, batch_size=16, shuffle=True),
130
+ valid_loader,
131
+ train_steps=1,
132
+ valid_steps=4,
133
+ )
134
+
135
+ metric_dict, _ = trainer.evalbatches()
136
+
137
+ # what the bug reported: the same model scored on batch 0, four times over
138
+ x, y = next(iter(valid_loader))
139
+ with torch.no_grad():
140
+ first_only = float(TestMSE()(y, trainer.model(x)))
141
+
142
+ assert abs(metric_dict[str(TestMSE())] - first_only) > 1e-3
File without changes
@@ -944,7 +944,7 @@ wheels = [
944
944
 
945
945
  [[package]]
946
946
  name = "mltrainer"
947
- version = "0.3.0"
947
+ version = "0.3.2"
948
948
  source = { editable = "." }
949
949
  dependencies = [
950
950
  { name = "loguru" },
mltrainer-0.3.0/Makefile DELETED
@@ -1,7 +0,0 @@
1
- format:
2
- black mltrainer
3
- isort mltrainer
4
-
5
- lint:
6
- flake8 mltrainer
7
- mypy mltrainer
@@ -1,177 +0,0 @@
1
- from __future__ import annotations
2
-
3
- from enum import Enum
4
- from pathlib import Path
5
- from typing import Any, Callable, Dict, List, Optional
6
-
7
- from loguru import logger
8
- from pydantic import BaseModel, ConfigDict, field_validator
9
-
10
-
11
- class FileTypes(Enum):
12
- JPG = ".jpg"
13
- PNG = ".png"
14
- TXT = ".txt"
15
- ZIP = ".zip"
16
- TGZ = ".tgz"
17
- TAR = ".tar.gz"
18
- GZ = ".gz"
19
-
20
-
21
- class ReportTypes(Enum):
22
- TOML = "TOML"
23
- TENSORBOARD = "TENSORBOARD"
24
- MLFLOW = "MLFLOW"
25
- RAY = "RAY"
26
-
27
-
28
- class FormattedBase(BaseModel):
29
- def __str__(self) -> str:
30
- return "\n".join(f"{k}: {v}" for k, v in self.__dict__.items())
31
-
32
- def __repr__(self) -> str:
33
- return "\n".join(f"{k}: {v}" for k, v in self.__dict__.items())
34
-
35
-
36
- class TrainerSettings(FormattedBase):
37
- epochs: int
38
- metrics: List[Callable]
39
- logdir: Path
40
- train_steps: int
41
- valid_steps: int
42
- reporttypes: List[ReportTypes]
43
- optimizer_kwargs: Dict[str, Any] = {"lr": 1e-3, "weight_decay": 1e-5}
44
- scheduler_kwargs: Optional[Dict[str, Any]] = {"factor": 0.1, "patience": 10}
45
- earlystop_kwargs: Optional[Dict[str, Any]] = {
46
- "save": False,
47
- "verbose": True,
48
- "patience": 10,
49
- }
50
-
51
- model_config = ConfigDict(arbitrary_types_allowed=True)
52
-
53
- @field_validator("logdir")
54
- @classmethod
55
- def check_path(cls, logdir: Path) -> Path: # noqa: N805
56
- if isinstance(logdir, str):
57
- logdir = Path(logdir)
58
- if not logdir.resolve().exists():
59
- logdir.mkdir(parents=True)
60
- logger.info(f"Created logdir {logdir.absolute()}")
61
- return logdir
62
-
63
- # def to_toml_dict(self) -> dict:
64
- # """Convert the model to a TOML-compatible dictionary."""
65
- # data = self.model_dump(exclude_none=True)
66
- #
67
- # # Handle special cases
68
- # if "metrics" in data:
69
- # data["metrics"] = [
70
- # f"{metric.__module__}.{metric.__name__}"
71
- # if isinstance(metric, Callable) and hasattr(metric, "__name__")
72
- # else f"{metric.__class__.__module__}.{metric.__class__.__name__}"
73
- # for metric in data["metrics"]
74
- # ]
75
- #
76
- # # Convert Path objects to strings
77
- # for key, value in data.items():
78
- # if isinstance(value, Path):
79
- # data[key] = str(value)
80
- #
81
- # return data
82
- #
83
- # def save_toml_example(self, file_path: str | Path) -> None:
84
- # """Save an example TOML configuration file with comments."""
85
- # if isinstance(file_path, str):
86
- # file_path = Path(file_path)
87
- #
88
- # data = self.to_toml_dict()
89
- # lines = []
90
- #
91
- # # Start with just the simple scalar values we know work
92
- # lines.append("epochs = {}\n".format(data["epochs"]))
93
- # metric_paths = [
94
- # f"{metric.__module__}.{metric.__name__}" for metric in self.metrics
95
- # ]
96
- # lines.append("metrics = {}\n".format(str(metric_paths).replace("'", '"')))
97
- # lines.append('logdir = "{}"\n'.format(str(data["logdir"])))
98
- # lines.append("train_steps = {}\n".format(data["train_steps"]))
99
- # lines.append("valid_steps = {}\n".format(data["valid_steps"]))
100
- #
101
- # # Add reporttypes as list of strings (enum names)
102
- # report_types = [rt.value for rt in data["reporttypes"]]
103
- # lines.append("reporttypes = {}\n".format(str(report_types).replace("'", '"')))
104
- #
105
- # if "optimizer_kwargs" in data:
106
- # lines.append(
107
- # self.dict_to_toml("optimizer_kwargs", data["optimizer_kwargs"])
108
- # )
109
- #
110
- # if "scheduler_kwargs" in data:
111
- # lines.append(
112
- # self.dict_to_toml("scheduler_kwargs", data["scheduler_kwargs"])
113
- # )
114
- #
115
- # if "earlystop_kwargs" in data:
116
- # lines.append(
117
- # self.dict_to_toml("earlystop_kwargs", data["earlystop_kwargs"])
118
- # )
119
- #
120
- # with open(file_path, "w", encoding="utf-8") as f:
121
- # f.write("".join(lines))
122
- #
123
- # @staticmethod
124
- # def dict_to_toml(name: str, data: dict) -> str:
125
- # """
126
- # Convert a dictionary into a TOML table format.
127
- #
128
- # Args:
129
- # name (str): Name of the TOML section.
130
- # data (dict): Dictionary to convert.
131
- #
132
- # Returns:
133
- # str: TOML formatted table string.
134
- # """
135
- # toml_lines = [f"[{name}]"]
136
- # for k, v in data.items():
137
- # # Convert booleans to lowercase ('true'/'false')
138
- # if isinstance(v, bool):
139
- # v = str(v).lower()
140
- # toml_lines.append(f"{k} = {v}")
141
- # return "\n".join(toml_lines) + "\n"
142
- #
143
- # @classmethod
144
- # def from_toml(cls, file_path: str | Path) -> "TrainerSettings":
145
- # """Load the model from a TOML file."""
146
- # if isinstance(file_path, str):
147
- # file_path = Path(file_path)
148
- #
149
- # with open(file_path, "rb") as f:
150
- # data = tomllib.load(f)
151
- #
152
- # print(data)
153
- #
154
- # # Convert string paths to Path objects
155
- # if "logdir" in data and isinstance(data["logdir"], str):
156
- # data["logdir"] = Path(data["logdir"])
157
- # # Convert reporttype strings back to enums
158
- # if "reporttypes" in data:
159
- # data["reporttypes"] = [ReportTypes(rt) for rt in data["reporttypes"]]
160
- #
161
- # # Import metrics from their module paths
162
- # if "metrics" in data:
163
- # metrics = []
164
- # for metric_path in data["metrics"]:
165
- # try:
166
- # module_path, func_name = metric_path.rsplit(".", 1)
167
- # module = __import__(module_path, fromlist=[func_name])
168
- # metric = getattr(module, func_name)
169
- # metrics.append(metric)
170
- # except (ImportError, AttributeError) as e:
171
- # logger.warning(
172
- # f"Could not load metric {metric_path}. "
173
- # f"Error: {e}. Using placeholder."
174
- # )
175
- # data["metrics"] = metrics
176
- #
177
- # return cls(**data)
File without changes
File without changes
File without changes
File without changes