attr-eomt 1.1.2__tar.gz → 1.1.4__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.
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/PKG-INFO +266 -10
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/README.md +265 -9
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/eomt/__init__.py +1 -1
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/eomt/api.py +14 -5
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/eomt/aux_cls.py +70 -30
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/eomt/box_loss.py +30 -0
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/eomt/config.py +77 -14
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/eomt/data/coco.py +171 -28
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/eomt/data/transforms.py +119 -41
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/eomt/device.py +13 -1
- attr_eomt-1.1.4/eomt/engine/distributed.py +103 -0
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/eomt/engine/predict.py +64 -9
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/eomt/engine/track.py +2 -1
- attr_eomt-1.1.4/eomt/engine/train.py +1295 -0
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/eomt/engine/validate.py +168 -27
- attr_eomt-1.1.4/eomt/loss.py +323 -0
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/eomt/model.py +158 -53
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/eomt/plotting.py +41 -25
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/eomt/postprocess.py +81 -28
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/eomt/reid.py +3 -0
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/eomt/serialization.py +33 -1
- attr_eomt-1.1.4/eomt/tta.py +419 -0
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/eomt/visualize.py +5 -3
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/pyproject.toml +1 -1
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/scripts/predict.py +5 -1
- attr_eomt-1.1.4/scripts/train.py +144 -0
- attr_eomt-1.1.4/scripts/val.py +88 -0
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/site/index.html +8 -5
- attr_eomt-1.1.4/tests/test_architecture.py +457 -0
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/tests/test_augment.py +82 -12
- attr_eomt-1.1.4/tests/test_ddp.py +74 -0
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/tests/test_parity.py +28 -6
- attr_eomt-1.1.4/tests/test_recipe.py +160 -0
- attr_eomt-1.1.4/tests/test_regression_attr.py +236 -0
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/tests/test_smoke.py +34 -4
- attr_eomt-1.1.4/tests/test_tta.py +288 -0
- attr_eomt-1.1.2/eomt/engine/train.py +0 -1014
- attr_eomt-1.1.2/eomt/loss.py +0 -279
- attr_eomt-1.1.2/scripts/train.py +0 -96
- attr_eomt-1.1.2/scripts/val.py +0 -30
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/.github/workflows/pages.yml +0 -0
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/.github/workflows/publish.yml +0 -0
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/.github/workflows/release.yml +0 -0
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/.github/workflows/tests.yml +0 -0
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/.gitignore +0 -0
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/LICENSE +0 -0
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/configs/coco.yaml +0 -0
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/docs/assets/00-hero-banner.png +0 -0
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/docs/assets/00-hero-banner.svg +0 -0
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/docs/examples/billie_eilish.jpg +0 -0
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/docs/examples/rosalia.jpg +0 -0
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/docs/examples/sample.jpg +0 -0
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/docs/examples/taylor_swift.jpg +0 -0
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/eomt/compress.py +0 -0
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/eomt/data/__init__.py +0 -0
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/eomt/data/config.py +0 -0
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/eomt/data/download.py +0 -0
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/eomt/ema.py +0 -0
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/eomt/engine/__init__.py +0 -0
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/eomt/engine/match.py +0 -0
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/eomt/preprocess.py +0 -0
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/sample_data/README.md +0 -0
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/sample_data/annotations/instances_train.json +0 -0
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/sample_data/annotations/instances_val.json +0 -0
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/sample_data/data.yaml +0 -0
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/sample_data/images/train/.gitkeep +0 -0
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/sample_data/images/val/.gitkeep +0 -0
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/scripts/info.py +0 -0
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/scripts/match.py +0 -0
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/scripts/push_to_hub.py +0 -0
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/scripts/track.py +0 -0
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/scripts/train_detect_coco.sh +0 -0
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/tests/test_api.py +0 -0
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/tests/test_compress.py +0 -0
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/tests/test_reid.py +0 -0
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/tests/test_sweep.py +0 -0
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/tools/attr_tagger/README.md +0 -0
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/tools/attr_tagger/tag.py +0 -0
- {attr_eomt-1.1.2 → attr_eomt-1.1.4}/uv.lock +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.5
|
|
2
2
|
Name: attr-eomt
|
|
3
|
-
Version: 1.1.
|
|
3
|
+
Version: 1.1.4
|
|
4
4
|
Summary: Standalone EoMT (Encoder-only Mask Transformer) for instance segmentation, with DINOv2 init, COCO training/validation and inference.
|
|
5
5
|
Project-URL: Homepage, https://github.com/imagra93/attr-eomt
|
|
6
6
|
Project-URL: Documentation, https://imagra93.github.io/attr-eomt
|
|
@@ -72,7 +72,7 @@ It's a clean-room, Apache-2.0 reimplementation: the weights you train are yours
|
|
|
72
72
|
from eomt import EoMT
|
|
73
73
|
|
|
74
74
|
model = EoMT("l") # fresh large model (DINOv2 backbone)
|
|
75
|
-
model.train(data="coco"
|
|
75
|
+
model.train(data="coco") # COCO 2017 auto-downloads if missing; 12 epochs, the official recipe
|
|
76
76
|
|
|
77
77
|
model = EoMT("runs/train/eomt-l") # reload a run — size/classes/heads auto-detected
|
|
78
78
|
model.predict("images/", plot=True) # render masks/boxes + per-instance attributes
|
|
@@ -206,14 +206,22 @@ simply reuse that same query→GT assignment and read the answer off the matched
|
|
|
206
206
|
query→GT assignment the detection loss used; the attribute is read *after* matching.
|
|
207
207
|
- **Gate.** An optional IoU gate drops barely-overlapping matched pairs (common early in
|
|
208
208
|
training) so attributes only learn from queries that actually localize the object.
|
|
209
|
-
- **Loss.** Cross-entropy per head over matched queries
|
|
209
|
+
- **Loss.** Cross-entropy per head over matched queries (a regression head: Huber on its
|
|
210
|
+
standardized value, see below), summed across heads and scaled
|
|
210
211
|
by `aux_w` (default `1.0`), added to the detector loss. Empty-match batches contribute
|
|
211
212
|
a graph-preserving zero, and missing labels use `ignore_index` and contribute nothing.
|
|
213
|
+
- **Balancing.** Both losses are O(1), so equal weights are a sane start. `aux_w_per_head`
|
|
214
|
+
weights each head, `aux_class_weights=True` re-weights rare classes (rare values of a
|
|
215
|
+
regression head, by smoothed density), and `aux_balance="uncertainty"` learns one weight
|
|
216
|
+
per head (Kendall et al. 2018) with the detector loss as the fixed anchor. Each epoch
|
|
217
|
+
prints every head's pull on the shared queries (its gradient norm over the detector's,
|
|
218
|
+
`train/aux_grad/<head>`): well above 1 a head steers the encoder, far below it only reads it.
|
|
212
219
|
- **Checkpoint selection is unchanged.** The attribute "rides along": its per-head
|
|
213
220
|
matched-query accuracy is shown live and written to `metrics.csv`, but never drives
|
|
214
221
|
`best.pt` (still `segm/mAP` or `bbox/mAP`).
|
|
215
222
|
- **Inference.** Each result attaches `aux = {head: {"ids", "probs"}}` for the kept
|
|
216
|
-
detections
|
|
223
|
+
detections (and `aux_values = {head: values}` in real units for the regression heads), and
|
|
224
|
+
`predict(plot=True)` renders each attribute next to the class label
|
|
217
225
|
using names stored in the checkpoint.
|
|
218
226
|
|
|
219
227
|
---
|
|
@@ -282,8 +290,76 @@ label set, in one model:
|
|
|
282
290
|
At inference a scoped head reports `ids = -1` ("not applicable") for detections whose class it
|
|
283
291
|
does not cover. The scope is stored in the checkpoint, so it survives reload.
|
|
284
292
|
|
|
293
|
+
**Flip-aware attributes.** A horizontal flip mirrors the photo but not the labels, so an attribute
|
|
294
|
+
that encodes a side (a left door shows up as a right one) is corrupted by `flip_prob > 0`. Give
|
|
295
|
+
its definition `flip_aware`, a map of category names to the ones they become on a horizontal
|
|
296
|
+
flip, and the training augmentation swaps them for every instance it flips (each pair swaps both
|
|
297
|
+
ways, so `{"left": "right"}` is enough; unlisted categories stay as they are):
|
|
298
|
+
|
|
299
|
+
```jsonc
|
|
300
|
+
{"name": "laterality", "categories": [{"id": 0, "name": "left"}, {"id": 1, "name": "center"}, {"id": 2, "name": "right"}],
|
|
301
|
+
"flip_aware": {"left": "right"}}
|
|
302
|
+
```
|
|
303
|
+
|
|
304
|
+
Mosaic, mixup and copy-paste flip each source image independently, and the swap follows each
|
|
305
|
+
instance. Only the horizontal flip is handled: keep `vflip_prob` and `rot90_prob` off for such an
|
|
306
|
+
attribute. It is categorical only. The mapping is stored in the checkpoint, and test-time augmentation with
|
|
307
|
+
`flip` swaps it back in each mirrored view's prediction before the views are merged (a checkpoint saved before the
|
|
308
|
+
mapping was stored takes it from the val JSON, with a warning). The training banner lists each head's swapped pairs. A
|
|
309
|
+
regression head needs no mapping: its value must not change under a flip (see below).
|
|
310
|
+
|
|
311
|
+
**Regression attributes.** `"type": "regression"` makes a head regress one real number per
|
|
312
|
+
instance instead of choosing a category; categorical and regression heads mix freely in one
|
|
313
|
+
dataset. `log` (regress `log1p` of the value, for heavy-tailed positives such as costs) and
|
|
314
|
+
`range` (valid values; predictions are clipped to it) are optional, and `applies_to` works as above:
|
|
315
|
+
|
|
316
|
+
```jsonc
|
|
317
|
+
"attributes": [
|
|
318
|
+
{"name": "severity", "categories": [...]},
|
|
319
|
+
{"name": "repair_cost", "type": "regression", "log": true, "range": [0, 20000], "applies_to": ["bumper"]}
|
|
320
|
+
]
|
|
321
|
+
// per annotation:
|
|
322
|
+
"attributes": {"severity": 3, "repair_cost": 812.5}
|
|
323
|
+
```
|
|
324
|
+
|
|
325
|
+
A missing, non-numeric or out-of-range value is ignored. The head is trained on the value
|
|
326
|
+
standardized by the train split's mean and std (stored in the checkpoint), validated as the
|
|
327
|
+
mean absolute error in real units (`val/aux_mae/<head>`; `metrics.png` gives each regression
|
|
328
|
+
head its own MAE panel, `aux_per_class.png` its MAE per primary class), and predicted in real units
|
|
329
|
+
(`aux_values`). The value must not change under the augmentations (scale, crop, flip,
|
|
330
|
+
rotation): a cost or a score is fine, a pixel area or an angle is not. Tracking votes on the
|
|
331
|
+
categorical heads only, and re-identification cannot gate on a regression head.
|
|
332
|
+
|
|
333
|
+
A value that depends on how much of the instance is in view can say so with `crop`: the
|
|
334
|
+
training augmentation measures, per instance, the fraction it keeps in view (mask mass, or box
|
|
335
|
+
area, inside the crop window, after copy-paste occlusion), and the value follows it.
|
|
336
|
+
`"crop": "scale"` multiplies the value by that fraction — for an extensive quantity such as the
|
|
337
|
+
object's area in cm² (half the damage cropped away ⇒ half the area). `"crop": "ignore"` drops the
|
|
338
|
+
value when less than `min_visible` (default `0.9`) is kept — for a property of the whole object
|
|
339
|
+
that a cut view no longer shows. Validation and inference see whole images, so the head predicts
|
|
340
|
+
the full value there:
|
|
341
|
+
|
|
342
|
+
```jsonc
|
|
343
|
+
{"name": "damage_area", "type": "regression", "log": true, "crop": "scale"},
|
|
344
|
+
{"name": "part_length", "type": "regression", "crop": "ignore", "min_visible": 0.95}
|
|
345
|
+
```
|
|
346
|
+
|
|
347
|
+
How the fraction is measured: for masks, the share of the instance's mask mass whose source pixels
|
|
348
|
+
the crop (and rotation / shear / perspective warp) lands inside the window: exact, so an instance
|
|
349
|
+
the window does not cut keeps exactly 1 (a warp alone is not a cut); copy-paste occlusion counts as
|
|
350
|
+
hidden, and instances brought in by mosaic, mixup or copy-paste are measured against their own crops. For boxes (detect family), the ratio of
|
|
351
|
+
the box areas: exact for a crop, approximate under rotation. Each call of the built-in
|
|
352
|
+
augmentation exposes it as `TrainAugment.last_info["visible"]`.
|
|
353
|
+
|
|
354
|
+
Limits: the head's mean / std and its `aux_class_weights` density come from the unscaled values,
|
|
355
|
+
so with `"scale"` the values trained on sit somewhat lower than those statistics. Only the visible
|
|
356
|
+
fraction is tracked, not the scale jitter, so a pixel area is still not supported. A custom
|
|
357
|
+
`transform=` gets no adjustment. `crop` is rejected on a categorical attribute, and it is a
|
|
358
|
+
training-only setting (not stored in the checkpoint). The training banner prints each head's crop
|
|
359
|
+
mode.
|
|
360
|
+
|
|
285
361
|
**Sidecar format (optional).** You can keep the COCO JSON as plain, standard COCO and put the
|
|
286
|
-
attributes beside it instead of inside it: an `attributes.yaml` schema in the dataset root plus
|
|
362
|
+
attributes beside it instead of inside it: an `attributes.yaml` schema (same keys, regression heads included) in the dataset root plus
|
|
287
363
|
`attributes/<split>.json` values keyed by annotation id (`{ann_id: {head: value}}`). If present
|
|
288
364
|
(and the JSON has no embedded `attributes`), it is merged in memory at load — so a plain COCO
|
|
289
365
|
dataset always works and the sidecar is picked up automatically when you add it. Embedded
|
|
@@ -406,7 +482,7 @@ size, normalization and any auxiliary heads are auto-detected from the `.pt`):
|
|
|
406
482
|
from eomt import EoMT
|
|
407
483
|
|
|
408
484
|
# Train on COCO 2017 (auto-downloaded on first run):
|
|
409
|
-
EoMT("l").train(data="coco",
|
|
485
|
+
EoMT("l").train(data="coco", batch=4) # 12 epochs, validated every 2: the official recipe
|
|
410
486
|
|
|
411
487
|
# ...or any COCO-format dataset (point at its data.yaml):
|
|
412
488
|
EoMT("s").train(data="sample_data/data.yaml", epochs=1, batch=1)
|
|
@@ -416,9 +492,179 @@ EoMT("runs/train/eomt-l").val(data="coco")
|
|
|
416
492
|
EoMT("runs/train/eomt-l").predict("images/", plot=True) # writes annotated images
|
|
417
493
|
```
|
|
418
494
|
|
|
495
|
+
**Several GPUs.** `device="0,1"` trains on both GPUs, one process each (DistributedDataParallel); `"auto"` is GPU 0.
|
|
496
|
+
`batch` is per GPU and `nominal_batch` (16) stays the global effective batch, so the optimizer steps, LR schedule and
|
|
497
|
+
masked-attention annealing are exactly those of one GPU, with an epoch 1.9× faster on two RTX 5090s (1.8× end to end with
|
|
498
|
+
validation and checkpoints). Validation splits the val images between the
|
|
499
|
+
GPUs; the main process scores them and alone logs and writes checkpoints. The processes are spawned, so every `train()`
|
|
500
|
+
argument must be picklable.
|
|
501
|
+
|
|
502
|
+
```python
|
|
503
|
+
EoMT("l", device="0,1").train(data="coco", batch=2) # 2 per GPU x accum 4 x 2 GPUs = effective 16
|
|
504
|
+
```
|
|
505
|
+
|
|
506
|
+
**Test-time augmentation and tiled inference** (instance family). `predict()` and `val()` take `augment=True` (the
|
|
507
|
+
image at `imgsz` and zoomed to 1.5x `imgsz`) and `tiles=True` (overlapping tiles of `imgsz` px of the original image, i.e. native resolution,
|
|
508
|
+
25 % overlap, plus the whole image: small objects in large images). The views' instances are merged: same class and
|
|
509
|
+
mask IoU >= 0.5 measured where both views can see, so an object cut by a tile border is put back together; the mask
|
|
510
|
+
logits are averaged and the score is the mean over the views that could see the object (one that missed it votes 0).
|
|
511
|
+
Each takes a dict of overrides; see
|
|
512
|
+
[`eomt/tta.py`](eomt/tta.py). `augment={"flip": True}` also merges mirrored views. Attributes marked `flip_aware` are swapped
|
|
513
|
+
back in a mirrored view; left / right *primary classes* are not (a mirrored view predicts them swapped), so keep `flip` off
|
|
514
|
+
for those, and for attributes that encode a side without being flip-aware.
|
|
515
|
+
|
|
516
|
+
```python
|
|
517
|
+
m = EoMT("runs/train/eomt-l")
|
|
518
|
+
m.predict("images/", augment=True) # + the image at 1.5x imgsz
|
|
519
|
+
m.val(data="data.yaml", tiles=True) # native tiles + the whole image
|
|
520
|
+
m.predict("images/", augment={"flip": True}, tiles=1024) # 1x / 1.5x and mirrored; 1024 px tiles
|
|
521
|
+
```
|
|
522
|
+
|
|
523
|
+
`scripts/val.py` and `scripts/predict.py` take `--augment` and `--tiles [SIZE]` (`scripts/val.py` also `--flip`; its arguments are
|
|
524
|
+
in [Command-line scripts](#command-line-scripts-arguments-and-things-to-take-into-account)). Each view is a forward pass: `augment`
|
|
525
|
+
costs ~3.5x one prediction (the 1.5x view has 2.25x the pixels); a 1280 x 960 photo at `imgsz` 644 makes 6 tiles + the
|
|
526
|
+
whole image.
|
|
527
|
+
|
|
419
528
|
For the full training recipe, every `train()` knob, and int8 compression, see the
|
|
420
529
|
**[annotated explainer →](https://imagra93.github.io/attr-eomt)** — it's the deep dive.
|
|
421
530
|
|
|
531
|
+
### Optimisation recipe
|
|
532
|
+
|
|
533
|
+
`train()` follows the reference EoMT recipe (checked numerically against the official code: every parameter group's
|
|
534
|
+
LR, every step of the LR schedule, and the masked-attention annealing agree exactly):
|
|
535
|
+
|
|
536
|
+
| | default |
|
|
537
|
+
|---|---|
|
|
538
|
+
| optimizer | AdamW, `lr0=1e-4`, `weight_decay=0.05` (not applied to norms, biases and embeddings) |
|
|
539
|
+
| head LR | `lr0` |
|
|
540
|
+
| encoder LR | layer-wise decay `llrd=0.8`, `backbone_lr_mult=1.0`: the **top ViT block is at the full `lr0`** and each block below is 0.8× the one above; embeddings follow the first block, the final norm is at `lr0` |
|
|
541
|
+
| schedule | polynomial decay (`poly_power=0.9`) to 0 after a two-stage warmup in optimizer steps `warmup_steps=(500, 1000)`: the head warms up while the ViT LR is held at 0, then the ViT ramps up |
|
|
542
|
+
| masked attention | annealed block by block (`mask_anneal_start/end`, fractions of training: 1/6–2/6, 2/6–3/6, 3/6–4/6, 4/6–5/6, polynomial); the last sixth of training is mask-free |
|
|
543
|
+
| length | `epochs=12`, validated every `val_interval=2` epochs and after the last (official COCO instance recipe, about 89k optimizer steps at batch 16) |
|
|
544
|
+
|
|
545
|
+
In EoMT the last four transformer blocks are the ones that process the queries and predict the masks, so they must learn at
|
|
546
|
+
close to the head's speed: scaling the whole encoder by 0.1 (as detector-style recipes do for a backbone) starves them. The
|
|
547
|
+
LR decay and the annealing span `epochs`.
|
|
548
|
+
|
|
549
|
+
**Validation follows the training regime**, as in the official EoMT: masked attention at the annealing probabilities of
|
|
550
|
+
the epoch's end. A model that still masks is far worse without its masks (part damage, ViT-L at epoch 11 of 50: 0.034
|
|
551
|
+
segm mAP mask-free against 0.114 masked), so a mask-free curve before the last sixth follows the annealing and jumps as each
|
|
552
|
+
block finishes it, not the model. `best.pt` is deployed mask-free, so it is only chosen once every block's annealing is
|
|
553
|
+
over (a run stopped earlier, e.g. with `stop_after_epochs`, has `last.pt` alone), and every checkpoint is saved mask-free.
|
|
554
|
+
Training prints the windows in epochs when it starts.
|
|
555
|
+
|
|
556
|
+
### Command-line scripts: arguments and things to take into account
|
|
557
|
+
|
|
558
|
+
**`scripts/train.py`.** Every flag is a keyword of `train()`; a flag left out keeps `train()`'s default.
|
|
559
|
+
|
|
560
|
+
| flag | `train()` keyword | what it does |
|
|
561
|
+
|---|---|---|
|
|
562
|
+
| `--data` | `data` | dataset YAML, or `coco` (auto-downloads) |
|
|
563
|
+
| `--size` / `--task` | | model size `s` `b` `l`; head family `instance` (masks) or `detect` (boxes) |
|
|
564
|
+
| `--epochs` | `epochs` | total epochs: the LR schedule and the masked-attention annealing span them |
|
|
565
|
+
| `--batch` / `--nominal-batch` | `batch` / `nominal_batch` | micro-batch **per GPU** / global effective batch (16); the difference is gradient accumulation |
|
|
566
|
+
| `--device` | `device` | `auto`, `cpu`, `cuda:N`, or `0,1` for several GPUs (one process each) |
|
|
567
|
+
| `--weights` / `--resume` | `init_weights` / `resume` | warm-start a fresh run from a checkpoint / continue a run. **One of the two, not both** |
|
|
568
|
+
| `--lr`, `--warmup HEAD,VIT` | `lr0`, `warmup_steps` | peak head LR (the encoder follows with layer-wise decay) and the two-stage warmup in optimizer steps |
|
|
569
|
+
| `--ema-decay`, `--ema-tau` | `ema_decay`, `ema_tau` | EMA of the weights (what is validated and exported as `best.pt`) |
|
|
570
|
+
| `--mask-free` | `mask_anneal_start/end` | no masked attention from step 0 (annealing windows 0 -> 0) |
|
|
571
|
+
| `--aux-dropout` | `aux_head_dropout` | dropout in the attribute heads |
|
|
572
|
+
| `--val-interval`, `--workers`, `--seed` | same | validate every N epochs (and after the last); data-loader workers per GPU; seed |
|
|
573
|
+
| `--aug KEY=VALUE`, `--aug-preset` | `aug` | augmentation overrides (see [Augmentation](#augmentation)) |
|
|
574
|
+
| `--fpn-scales`, `--box-head`, `--no-deep-supervision`, `--upscale-blocks`, `--iou-aware-cls`, `--quality-weight`, `--val-imgsz` | same | [architecture options](#architecture-options) |
|
|
575
|
+
|
|
576
|
+
Things to take into account when training:
|
|
577
|
+
|
|
578
|
+
- **Resume or warm start.** `--resume` continues a run exactly (epoch, optimizer, schedule, EMA, best metric) and re-reads the
|
|
579
|
+
architecture from the checkpoint: keep the other flags as they were, and the seed is ignored so the shuffle is not replayed.
|
|
580
|
+
`--weights` loads the weights only and starts a new run at epoch 0. Start from a `best.pt` (the EMA, saved mask-free); a
|
|
581
|
+
`last.pt` holds the live weights.
|
|
582
|
+
- **A short fine-tune needs its own recipe.** Lower `--lr`, shorten `--warmup`, and lower `--ema-decay` (at 0.9999 the EMA
|
|
583
|
+
averages over ~10k steps and barely leaves its starting point in a few epochs). If the model is deployed mask-free, pass
|
|
584
|
+
`--mask-free`: with the default annealing a warm-started model is trained *masked* for the first sixth of the run, a regime it
|
|
585
|
+
no longer runs in.
|
|
586
|
+
- **`best.pt` is only chosen mask-free.** With the default annealing that means from 5/6 of the schedule on; with `--mask-free`
|
|
587
|
+
every validation is eligible.
|
|
588
|
+
- **Validated epochs only.** With `--val-interval N` the other epochs have empty validation cells in `metrics.csv`; the curves in
|
|
589
|
+
`metrics.png` join the validated points.
|
|
590
|
+
- **Attribute heads follow the dataset's definitions.** Categories, `applies_to`, `flip_aware` and `crop` / zoom settings come
|
|
591
|
+
from the dataset (the training banner lists them). The flip mapping, scope and regression statistics are stored in the
|
|
592
|
+
checkpoint; `crop` and the zoom settings are training-only.
|
|
593
|
+
|
|
594
|
+
**`scripts/val.py`.**
|
|
595
|
+
|
|
596
|
+
| flag | what it does |
|
|
597
|
+
|---|---|
|
|
598
|
+
| `weights` | checkpoint `.pt` or a run / weights folder (`best.pt` preferred) |
|
|
599
|
+
| `--data`, `--batch`, `--workers`, `--device` | dataset YAML (`coco`), batch, data-loader workers, device |
|
|
600
|
+
| `--imgsz` | evaluation input size (default: the trained one; the positional embeddings are interpolated) |
|
|
601
|
+
| `--augment` | add the image at 1.5x `imgsz` |
|
|
602
|
+
| `--tiles [SIZE]` | overlapping tiles (default `imgsz` px, native resolution) plus the whole image |
|
|
603
|
+
| `--flip` | also run every view mirrored (alone: `--augment` with mirrored views) |
|
|
604
|
+
| `--no-attributes` | mAP only |
|
|
605
|
+
| `--ema`, `--mask annealed\|free` | validate the EMA weights of a `last.pt` (see below) |
|
|
606
|
+
|
|
607
|
+
Things to take into account when validating:
|
|
608
|
+
|
|
609
|
+
- **What is reported.** `segm/*` and `bbox/*` mAP and, for a model with secondary heads, `aux_acc/<head>` (accuracy) and
|
|
610
|
+
`aux_mae/<head>` (regression heads, mean absolute error in real units). They are scored on the same predictions as the mAP, so
|
|
611
|
+
they also work with `--augment`, `--tiles` and `--flip`.
|
|
612
|
+
- **The attribute population is the deployed one.** Each GT instance is matched to the kept detection of its class with the
|
|
613
|
+
highest mask IoU (one to one, IoU >= 0.5). The trainer's `val/aux_*` columns score matched *queries* instead, so the two differ
|
|
614
|
+
slightly: compare TTA against plain within the script's numbers, not against `metrics.csv`. A regression head's MAE is
|
|
615
|
+
dominated by its large values, so also look at a relative error before concluding.
|
|
616
|
+
- **Regime.** A model that is still annealing is only meaningful with its masks: `--ema --mask annealed` reproduces the trainer's
|
|
617
|
+
validation of a `last.pt` (EMA weights, the masked-attention probabilities of that epoch, read from the run's `args.yaml`),
|
|
618
|
+
`--ema --mask free` is the deployed regime, and `best.pt` is already the EMA, mask-free.
|
|
619
|
+
- **TTA is not free and not uniform.** It costs one forward pass per view and mostly helps small objects and strict-IoU mAP; the
|
|
620
|
+
classification heads often do not move. A regression head's prediction is *averaged over views at different zoom*, so a value
|
|
621
|
+
that depends on the object's apparent size can drift under TTA: check it against the plain run before relying on it.
|
|
622
|
+
- **`--flip` and the heads.** Heads the checkpoint stores as flip-aware are swapped back in mirrored views; an older checkpoint
|
|
623
|
+
takes the mapping from the dataset's JSON, with a warning. Left / right *primary classes* are not handled: keep `--flip` off for
|
|
624
|
+
them. The val categories of every head must match the model's, or the run stops instead of scoring against the wrong ids.
|
|
625
|
+
- **One process.** `scripts/val.py` runs on one device; the trainer's own validation uses all the GPUs.
|
|
626
|
+
|
|
627
|
+
### Architecture options
|
|
628
|
+
|
|
629
|
+
All are keyword arguments of `train()` (and flags of `scripts/train.py`); checkpoints record them, so `EoMT(path)` rebuilds
|
|
630
|
+
the model with the same modules, and checkpoints without them load as before.
|
|
631
|
+
|
|
632
|
+
| option | what it does | cost, ViT-L at 644 px, batch 2, bf16 (measured, relative to the default) |
|
|
633
|
+
|---|---|---|
|
|
634
|
+
| `fpn_scales=None` | drop the multi-scale FPN (default `(2, 1, 0.5)`) | +12 % throughput, −0.5 GB |
|
|
635
|
+
| `box_head=True` | auxiliary box head: L1 + GIoU loss on the matched queries, and the same terms in the matching cost (a localisation signal that does not depend on the mask grid). Predicted boxes are returned as `aux_boxes` / `head_boxes` | −7 % |
|
|
636
|
+
| `deep_supervision=False` | stop supervising a query block's prediction once that block's masked attention is annealed away (while it masks, the prediction builds its attention mask and stays supervised: unsupervised, it made noise masks). Blocks drop out at 2/6, 3/6, 4/6, 5/6 of training | none while blocks mask; +26 % per step once all are annealed (last sixth) |
|
|
637
|
+
| `num_upscale_blocks=3` | 368² mask logits at 644 px (small masks); the block predictions stay at 184² (they only build 46² attention masks), only the output is 368²; cannot warm-start a 2-block checkpoint | −17 % (10.1 vs 12.1 img/s), 22.2 GB |
|
|
638
|
+
| `iou_aware_cls=True` | train a matched query's class probability towards `(1 + mask IoU) / 2` instead of 1 (IoU-aware targets, as in VarifocalNet / Stable-DINO), so the score ranks loose masks below tight ones; above 0.5, so matched queries keep their class for the attribute gate. Hard targets reduce exactly to the weighted CE. Windscreen damage, 12 epochs at 644 px, against the same recipe: segm AP 0.0296 -> 0.0758, AP small 0.0033 -> 0.0106, AP large 0.128 -> 0.210, bbox AP 0.177 -> 0.217, attribute accuracy unchanged | one upsample of the matched masks |
|
|
639
|
+
| `quality_weight=2` | mask-quality head: a small MLP on the query embedding predicts each matched query's mask IoU (BCE against the measured IoU, final prediction only; Mask Scoring R-CNN), and postprocess scores class probability x predicted IoU instead of the mean mask probability. Windscreen damage, with `iou_aware_cls`, two seeds against the same recipe without the head: segm AP +7 / +13 %, AP small +19 % / 3.6x, AP medium +21 / +22 %, AR100 +2 / +5 %, bbox AP and attribute accuracy unchanged | one MLP on the queries |
|
|
640
|
+
| `val_imgsz=N` | validate (and pick `best.pt`) at another input size than `imgsz`. Scale jitter shows the objects larger, on average, than a val image resized to `imgsz`: validate at `imgsz` × the mean training scale. Windscreens (training scale 0.4-1.0 of the photo, val 0.5), same checkpoint: segm AP 0.0113 at 644, 0.0173 at 896 | val time grows with the pixels |
|
|
641
|
+
| `stop_after_epochs=N` | stop after N epochs while the LR schedule and mask annealing still span `epochs` (truncated A/B runs) | |
|
|
642
|
+
|
|
643
|
+
Always on (training-side; inference of existing checkpoints is unchanged):
|
|
644
|
+
|
|
645
|
+
- **Masked attention keeps small objects in view.** A query may attend to every patch its predicted mask touches (max-pool of
|
|
646
|
+
the mask logits onto the patch grid), and a query whose mask touches no patch attends to all of them, as in Mask2Former. The
|
|
647
|
+
reference EoMT downsamples bilinearly (only the centre of each patch counts) and leaves an empty-mask query blind, which in
|
|
648
|
+
practice blinded the queries of most objects smaller than a patch during the masked phase. The mask is a boolean
|
|
649
|
+
`[B, 1, N, N]` and is not built in a block whose annealing probability has reached 0.
|
|
650
|
+
- **Mask terms that see small objects.** Matching scores every query on every cell of the mask-logit grid against the GT
|
|
651
|
+
area-averaged onto it, instead of on 12,544 random points (one per ~33 px² at 644 px, so an object of a few dozen pixels often
|
|
652
|
+
got no point and its assigned query changed from step to step). The mask loss keeps the reference's BCE on 12,544 points
|
|
653
|
+
concentrated where the prediction is uncertain (the boundaries), and computes dice on every GT pixel against the upsampled
|
|
654
|
+
logits, so every instance is supervised however small. A fully dense loss gave coarser masks (AP75 down) on thin objects;
|
|
655
|
+
dropping the points altogether lost the boundaries. The `train_num_points` / `oversample_ratio` / `importance_sample_ratio`
|
|
656
|
+
settings are gone (fixed at the reference values; checkpoints that carry them still load).
|
|
657
|
+
- **Attribute IoU gate at full resolution.** The gate that decides which matched queries train the attribute heads compares the
|
|
658
|
+
upsampled prediction with the GT, instead of shrinking the GT to the logit grid with "nearest" (which erased many objects of a
|
|
659
|
+
few pixels, so they never trained the attribute heads).
|
|
660
|
+
|
|
661
|
+
Inference: `predict_image(..., amp=True)` (and `predict(..., amp=True)`) runs the network under bf16 autocast, ~2.5x faster at 644 px
|
|
662
|
+
on an RTX 5090 with the same accuracy on a trained ViT-L model (AP within 0.002 in every size bucket). Off by default.
|
|
663
|
+
|
|
664
|
+
**What is not established:** the accuracy effect of `box_head`, `deep_supervision=False`,
|
|
665
|
+
`iou_aware_cls` and `quality_weight` (beyond the windscreen data), dropping the FPN and the always-on changes above depends on the data. Measure it with truncated runs (`stop_after_epochs`), which follow the schedule of a full run, against a reference run at equal
|
|
666
|
+
epochs.
|
|
667
|
+
|
|
422
668
|
---
|
|
423
669
|
|
|
424
670
|
## Augmentation
|
|
@@ -435,7 +681,7 @@ mass-preserving), so thin objects are not shredded into dots by nearest-neighbou
|
|
|
435
681
|
| geometry | `flip_prob` 0.5 · `vflip_prob` 0 · `rot90_prob` 0 · LSJ `min_scale`/`max_scale` 0.1–2.0 · `rotate_prob` 0.3 (±10°) · `shear_prob` 0.2 (±5°) · `perspective_prob` 0.15 |
|
|
436
682
|
| optics / sensor | colour jitter 1.0 · `gamma_prob` 0.3 · `grayscale_prob` 0.05 · `blur_prob` 0.2 · `noise_prob` 0.2 · `jpeg_prob` 0.3 · `glare_prob` 0.15 · `erasing_prob` 0.2 |
|
|
437
683
|
| multi-image (instance family) | `mosaic_prob` 0 · `mixup_prob` 0 · `copy_paste_prob` 0 |
|
|
438
|
-
| crop | `instance_crop_prob` 0 (instance-aware crop) · `instance_crop_empty_prob` 0 |
|
|
684
|
+
| crop | `instance_crop_prob` 0 (instance-aware crop) · `instance_crop_empty_prob` 0 · `zoom_prob` 0 (close-ups, `zoom_scale` 4–5x) |
|
|
439
685
|
| masks | `mask_resize="area"` (or `"nearest"`) |
|
|
440
686
|
|
|
441
687
|
Override per run, per dataset, or from the CLI — precedence is explicit keywords (`flip_prob`, `min_scale`,
|
|
@@ -453,7 +699,7 @@ train_aug:
|
|
|
453
699
|
min_scale: 0.6 # thin / tiny objects: keep the effective scale r = imgsz/long_side * s in ~[0.4, 1.0]
|
|
454
700
|
max_scale: 1.6
|
|
455
701
|
instance_crop_prob: 0.8 # optional instance-aware crop: window placed around a (rarity-weighted) instance
|
|
456
|
-
instance_crop_rarity_attr:
|
|
702
|
+
instance_crop_rarity_attr: material
|
|
457
703
|
```
|
|
458
704
|
|
|
459
705
|
```bash
|
|
@@ -463,14 +709,24 @@ python scripts/train.py --data my/data.yaml --aug min_scale=0.6 --aug instance_c
|
|
|
463
709
|
Things worth knowing:
|
|
464
710
|
|
|
465
711
|
- **Orientation.** Flips, 90° turns and rotations change the pixels but not the label. If an attribute encodes
|
|
466
|
-
orientation (e.g. a `viewpoint` head) set `flip_prob=0` and `rotate_prob=0
|
|
467
|
-
|
|
712
|
+
orientation (e.g. a `viewpoint` head) set `flip_prob=0` and `rotate_prob=0`, or, for a left / right attribute,
|
|
713
|
+
mark it `flip_aware` (see *Flip-aware attributes*) and keep the horizontal flip; `vflip_prob` / `rot90_prob` stay
|
|
714
|
+
off unless your scenes have no canonical "up".
|
|
468
715
|
- **Thin or tiny objects.** Large-Scale Jitter at 0.1 shrinks a 5 px crack to under a pixel; keep the effective
|
|
469
716
|
scale within ~[0.4, 1.0] (`min_scale = 0.4·long_side/imgsz`, `max_scale = 1.0·long_side/imgsz`). The instance-aware
|
|
470
717
|
crop pays off when the crop is much smaller than the image (high zoom, small `imgsz`), and
|
|
471
718
|
`instance_crop_rarity_attr` makes rare attribute values the crop's focus more often.
|
|
719
|
+
- **Close-ups.** `zoom_prob` makes an occasional sample a close-up: scaled by `s ~ U(zoom_scale)` (default 4–5x
|
|
720
|
+
the whole-image view, past the photo's own resolution) instead of the LSJ range, with the window on one instance
|
|
721
|
+
(whole when it fits, else a window inside it). `zoom_attr` restricts it to instances with a value of that
|
|
722
|
+
attribute and to the samples that have one, e.g. `{"zoom_prob": 0.1, "zoom_attr": "area"}` so a crop-scaled area
|
|
723
|
+
head also learns that a part filling the frame keeps its real area, and that a cut part keeps the share in view.
|
|
724
|
+
It resizes the whole photo first (~1.5x the time of a normal sample at 5x), so keep it occasional.
|
|
472
725
|
- **Multi-image ops** (`mosaic`, `mixup`, `copy_paste`) add instances from other samples and are *off*: they splice or
|
|
473
726
|
blend long thin structures, so try them deliberately and look at the result.
|
|
727
|
+
- **Negatives.** Images with no annotations are dropped from the train set by default. `train(keep_empty=True)` keeps
|
|
728
|
+
them as zero-target samples (masks `(0, imgsz, imgsz)`, every query is trained as "no object"); validation already
|
|
729
|
+
includes them. Instance family only.
|
|
474
730
|
- `args.yaml` of a run records the resolved `aug` dict; the log prints the active ops at start.
|
|
475
731
|
- `build_train_transform(imgsz, aug=...)` returns the same pipeline for use in your own loaders
|
|
476
732
|
(`tf(image_uint8, masks) -> (image, masks)`).
|