Add support for training trackastra with SAM2 features

[anwai98]

Hi @C-Achard,

Here's are my minimal changes to make training work with SAM2 features.

Let me know how it looks!

PS. In case it helps, here's my yaml config file to train trackastra:

yaml config

# Trackastra finetuning config file for TOIAM dataset (using SAM2 features)
# Run: python /mnt/vast-nhr/home/archit/u12090/trackastra/scripts/train.py -c train_config.yaml

name: toiam_sam2_features
outdir: ./runs

# Data
ndim: 2
input_train:
  - /mnt/vast-nhr/projects/cidas/cca/data/toiam/data/00
  - /mnt/vast-nhr/projects/cidas/cca/data/toiam/data/01
input_val:
  - /mnt/vast-nhr/projects/cidas/cca/data/toiam/data/04
detection_folders:
  - TRA
  - SEG

# Feature backbone (aligned to pretrained model)
features: pretrained_feats_aug
pretrained_feats_model: facebook/sam2.1-hiera-base-plus
pretrained_feats_mode: mean_patches_exact
pretrained_feats_additional_props: regionprops_small
pretrained_n_augs: 15
reduced_pretrained_feat_dim: 128
rotate_features: true

# Finetuning from pretrained
model: /user/archit/u12090/.local/share/trackastra/models/general_2d_w_SAM2_features

# Model architecture (matching pretrained)
d_model: 256
num_encoder_layers: 4
num_decoder_layers: 4
dropout: 0.05
window: 4
attn_dist_mode: v1
causal_norm: none

# Training hyperparameters
epochs: 500
warmup_epochs: 5
train_samples: 32000
batch_size: 16
max_tokens: 2048
weight_decay: 0.01
weight_by_dataset: true

# Augmentation
crop_size:
  - 320
  - 320

# Caching
cachedir: ./runs/.cache

# Logging and other misc. stuff
logger: tensorboard
seed: 42

[anwai98]

Poof, some ruff linting structures are funny haha. All should be working now. Lemme know how it looks @C-Achard

[C-Achard]

Thanks @anwai98, had a quick look and the approach seems reasonable, if you end up requiring changes on the pretrained_feats repo happy to have a look as well.

One thing I noticed is that in train.py, if no model path is given (training from scratch), it will load the basic model from Trackastra, rather than the one from pretrained_feats, since in the inference-only version create() is called only from TrackingTransformer.from_folder and then it would likely crash due to the extra args.
Now you did mention you wanted to fine-tune only but maybe the best is to add some error handling if anyone tries to train a pretrained_feats model from scratch, since the resulting exception will likely look unclear if no guard is added.

Otherwise, I noticed some slightly misleading help strings in the CLI, perhaps have a look at the manuscript for better context on what these options do (I added comments on these with recommended defaults).

Finally, if your next step is to train a model, those previous configs may come in handy for that.

I hope this helps, I'm afraid I cannot test this extensively right now but happy to help further if anything is unclear in the review.

Best,
Cyril

[anwai98]

Hi @C-Achard,

Thank you so much for the detailed feedback. I'll check them out later in the evening and come back to you!

[anwai98]

Hi @C-Achard,

Sorry for the super late follow-up. I managed to come back to the PR this week only - had a couple of busy weeks in the past.

I took care of the comments you left. What do you think about the current state now?

[C-Achard]

Hi @C-Achard,

Sorry for the super late follow-up. I managed to come back to the PR this week only - had a couple of busy weeks in the past.

I took care of the comments you left. What do you think about the current state now?

Nice, thanks! This looks good as far as I can tell, definitely curious to see how this performs on your data. If it overfits too much I would look into the augmentation API (I can help) but this is likely not needed for a first check.

Let me know if I can help with anything else, thanks again

[anwai98]

Hi @C-Achard,

Super late update, sorry - the training pipeline works now, both finetuning from general_2d_w_SAM2_features and from scratch.

On the augmented copies: you were right that they're not generated, but FeatureExtractorAugWrapper and PretrainedAugmentations seem like are both already in trackastra_pretrained_feats. The only missing piece imo is WRAugPretrainedFeatures - around pretrained_features.py:1230, reaches for it via trackastra.data.wrfeat, but it only exists on your cy/aug-zarr-caching branch (around wrfeat.py:402). Moving that class into the package would make the wrapper self-contained, and I can wire up pretrained_n_augs on the trackastra side from there.

I also managed a first real run: finetuning on two (TOIAM) 800-frame sequences (ca. 570 detections/window), validating on a third. No sign of overfitting so far.

What do you think? @C-Achard

[anwai98]

Oh btw -- Claude noticed a bug in Trackastra-et-Ultra actually in a corner-case (I'll let you be the judge of it if it's good to safeguard it. Dropping Claude's opinion below)

In Trackastra-et-Ultra, trackastra_pretrained_feats/model/model.py:244:

pt_features = self.ptfeat_norm(pt_features).squeeze()

Bare .squeeze() drops every size-1 dim, so a batch of 1 loses its batch dim. The unsqueeze(0) repair at lines 251-255 only runs in the features_out is not None branch - when additional_features is unset, pretrained features are the only features, so it takes the else at line 266 and line 271 fails:

RuntimeError: Tensors must have same number of dimensions: got 3 and 2

Single-element batches occur in the val loop, so it's always broken for that config. Dropping the .squeeze() fixes it, and leaves the region-props path unchanged (checked before/after, losses match). Same code is in the 0.0.2 wheel.

[anwai98]

Ahha also spotted another issue in a round of review by Codex (also fixed it -- see f4bcd48): if rotate_features=True is enabled and the intention is (as far as I understand) to encode both Y and X, the previous implementation wasn't as expected because it encoded only Y.

[C-Achard]

Hi @anwai98,

Thanks for the further updates, glad to hear that the finetuning seems to be working.

Both issues you mentioned seem correctly reported to me, the feature rotation one is particularly important, thanks for finding that.
Just one thing for fine-tuning, since the current trained SAM2 model expects only the y-based rotation , fixing it might make the pretrained feature distribution different enough that it would worsen rather than improve performance, or it might disrupt the finetuning, so making it a parameter in that function might be useful beyond the immediate fix (i.e. specify whether to use both x/y or only one of x/y). Maybe propagate it to the config if possible.

I might have to rerun some of my offline analysis using the fixed version of the rotation to see if the initial conclusions regarding the disambiguation still hold; hopefully it only gets better from there.

Moving that class into the package would make the wrapper self-contained, and I can wire up pretrained_n_augs on the trackastra side from there.

If moving that is not too much work, it would definitely be great to have it. From what I recall we omitted it (and the training) to keep the results from the SAM2 models usable without creating too much maintenance overhead/API changes with training. But it would certainly make the PR more complete. I think it's the only important missing piece.

Happy to help with the review of the implementation when I have time, let me know if you go ahead!

Thanks again for trying this out and working on the API, really appreciate it.

Best,
Cyril

[anwai98]

Hi @C-Achard,

Both issues you mentioned seem correctly reported to me, the feature rotation one is particularly important, thanks for finding that.

Nice, thanks for the affirmation.

fixing it might make the pretrained feature distribution different enough that it would worsen rather than improve performance, or it might disrupt the finetuning, so making it a parameter in that function might be useful beyond the immediate fix (i.e. specify whether to use both x/y or only one of x/y). Maybe propagate it to the config if possible.

Oh yeah, good point. By default keeping it fixed at the previous setup, and can be overridden to the expected version. I'll take care of it!

I might have to rerun some of my offline analysis using the fixed version of the rotation to see if the initial conclusions regarding the disambiguation still hold; hopefully it only gets better from there.

That would be great, thanks!

From what I recall we omitted it (and the training) to keep the results from the SAM2 models usable without creating too much maintenance overhead/API changes with training. But it would certainly make the PR more complete. I think it's the only important missing piece.

Gotcha. Then I'll give it a shot and create a PR in the other repo ;)

Thanks again for trying this out and working on the API, really appreciate it.

Ofc. And we have our interests vested in this anyways hehe, so we thought bringing this in a correct working structure would be awesome for us, and others! 😁

I'd probably touch the updates I promised above on Wednesday and ping you then!

[anwai98]

Hi @C-Achard,

I opened a PR (C-Achard/Trackastra-et-Ultra#2) for the augmentation port. Also made the rotation axes configurable here, with defaults same as before.

I also wired in pretrained_n_augs from the other spot -- but wouldn't work to test it out without installation from source from my other open PR.

What do you think about this now?

[C-Achard]

l. 554 is missing:
if features_type in ["none", "wrfeat"]:

Right n_workers and feats extraction logics are mixed

[C-Achard]

Also "none" feature handling is missing it seems...

[C-Achard]

(Note that these do not seem to be introduced by you directly, they may be older oversights)

[C-Achard]

try:
    PRETRAINED_FEATS_INSTALLED = True
    if TYPE_CHECKING:
        from trackastra_pretrained_feats import FeatureExtractor
except ImportError:
    PRETRAINED_FEATS_INSTALLED = False
    if TYPE_CHECKING:
        FeatureExtractor = None  # type: ignore

The bool can be removed, with type checking this is safe and the bool is not used here

[C-Achard]

Hello again @anwai98,

Thanks for the update, and sorry for the late review.

Overall the PR in the other repo looks good, just a few comments, but it is mainly the remaining changes on this repo that may be more difficult.

A lot of the wiring was updated in the fork I used, https://github.com/C-Achard/trackastra/blob/cy/aug-zarr-caching/trackastra/data/wrfeat.py, e.g. when returning WRFeatures, self.__class__ should be used since now there are several subclasses of WRFeatures and it shouldn't just be cast to WRFeatures directly,
Similarly, the AugmentationFactory is missing here, and WRAugmentationPipeline can specify its return type in the fork, which we did not include in this version.
There are also a few other augmentation tweaks and fixes I did not port initially.

So I think this is only the first step of a larger migration, I purposefully omitted several not-so-trivial integration changes when porting this the first time, trying to make a minimal inference-only version.
The full pretrained_feats augmented training pipeline will require extra migration work; now the aug system is restored in the pretrained_feats repo, which is a great first step, but a lot of cross-repository wiring will still be needed.

None of this necessarily needs to be addressed as part of this PR, I just wanted to clarify the scope. Let me know if I can help :)

Best,
Cyril

[C-Achard]

@anwai98 Thinking a bit about this, would it perhaps be easier for you to use the training code in my fork directly first?

Fine-tuning should definitely be possible there, so you can at least try out the model with/without augmentations and this way determine whether refactoring here would be worth it?

Maybe you already have preliminary results to motivate add the augmentation API here though, but since there was no overfitting in your latest round you may want to check first whether it helps

[anwai98]

Hi @C-Achard,

Thanks for your response. I'll look at your comments and response closely later today.

Quick re: on something:

@anwai98 Thinking a bit about this, would it perhaps be easier for you to use the training code in my fork directly first?

Ofc, I'd be happy to. Then we can potentially have two parallels to draw conclusion from (hopefully not orthogonal haha)

Will come back to you later with more thoughts!

[C-Achard]

Let me know if you need anything that may be in my thesis archive as well (train scripts, config, etc), as I forget exactly what is and isn't in the fork!