code/docs/training.md

M4SEG training

PRIMA trains on M⁴Seg using one dataset adapter. Follow the data guide, then set the model assets, optional projector initialization weights, dataset root, and training annotation path in configs/train.yaml.

For a single GPU:

python train.py --config configs/train.yaml

For multiple GPUs:

deepspeed --num_gpus 4 train.py --config configs/train.yaml

To further fine-tune the complete released model, run:

python train.py --config configs/train.yaml --merged_ckpt_path PLAN-Lab/PRIMA

This initializes from the merged release and attaches fresh language-model LoRA branches; separate base, EVA, Q-Former, and projector initialization paths are ignored. The dataset paths are still required. The default training configuration uses separate initialization assets to train from the pretrained backbone.

The merged checkpoint includes Q-Former and its tokenizer, so further fine-tuning with merged_ckpt_path needs no separate Q-Former checkpoint. Only training from the separate backbone assets requires the Q-Former initialization paths in configs/train.yaml.

Training updates LLM query/value LoRA, SQuARE attention and query tokens, the text-to-mask projection, SAM mask decoder, and optionally the visual-to-language projector. The language backbone, EVA/Q-Former backbone, and SAM image/prompt encoders remain frozen. A pretrained visual-to-language projector can be supplied as initialization; this repository does not train a separate pretraining stage.

Effective batch size is per-device batch size multiplied by GPU count and accumulation. The example gives an effective batch of 160 on four GPUs. AdamW uses linear learning-rate warmup followed by cosine decay. The example is a starting configuration; asset-specific recipes and expected results are documented separately.

Use --max_steps for a short smoke run. Scalar losses are printed at optimizer steps. Non-finite losses or gradients stop training.

Recoverable state and portable trainable-state exports are saved under the configured output directory. Resume with --resume /path/to/recoverable_checkpoint or --auto_resume. Keep batch size and accumulation unchanged when resuming. Distributed runs resume their DeepSpeed state; single-GPU runs restore optimizer, scheduler, progress, and random states. Portable exports omit frozen backbones and require the same base assets at inference.