PRIMA / code /docs /training.md
kanguyen-vn's picture
Remove debug instrumentation and rename the M4Seg dataset adapter
5e66a57 verified
|
Raw History Blame Contribute Delete
2.54 kB
# M4SEG training
PRIMA trains on [M⁴Seg](https://huggingface.co/datasets/PLAN-Lab/M4Seg) using one dataset adapter. Follow the [data guide](datasets.md), then set the model assets, optional projector initialization weights, dataset root, and training annotation path in `configs/train.yaml`.
For a single GPU:
```bash
python train.py --config configs/train.yaml
```
For multiple GPUs:
```bash
deepspeed --num_gpus 4 train.py --config configs/train.yaml
```
To further fine-tune the complete [released model](https://huggingface.co/PLAN-Lab/PRIMA), run:
```bash
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.