Coverage for collectra/tasks/image_segmentation/sam3_training.py: 91.67%
84 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-16 09:56 +0000
1"""Adapt Meta's installed image-training template for mask-supervised COCO data."""
3from pathlib import Path
4from types import MethodType, SimpleNamespace
6import yaml
9def training_config(dataset, output, checkpoint, *, epochs, batch, num_gpus,
10 wandb=False, run_name=None, project=None, confidence_threshold=0.5):
11 from .sam3 import _sam3
13 package = Path(_sam3().__file__).parent
14 template = (
15 package / "train/configs/roboflow_v100/roboflow_v100_full_ft_100_images.yaml"
16 )
17 if not template.is_file():
18 raise FileNotFoundError(
19 f"Install SAM3's training sources; template missing: {template}"
20 )
21 config = yaml.safe_load(template.read_text())
22 config["paths"].update(
23 roboflow_vl_100_root=str(dataset),
24 experiment_log_dir=str(output / "training"),
25 bpe_path=str(package / "assets/bpe_simple_vocab_16e6.txt.gz"),
26 )
27 settings = config["roboflow_train"]
28 settings.update(supercategory="collectra", num_images=None)
29 losses = settings["loss"]["loss_fns_find"]
30 losses.append(
31 {
32 "_target_": "sam3.train.loss.loss_fns.Masks",
33 "focal_alpha": 0.25,
34 "focal_gamma": 2.0,
35 "weight_dict": {"loss_mask": 200.0, "loss_dice": 10.0},
36 "compute_aux": False,
37 }
38 )
39 # Validation masks must be decoded too; Meta's detection template only
40 # decodes them in the training transform sequence.
41 settings["val_transforms"][0]["transforms"].insert(
42 0,
43 {
44 "_target_": "sam3.train.transforms.segmentation.DecodeRle",
45 },
46 )
47 config["scratch"].update(
48 enable_segmentation=True,
49 train_batch_size=batch,
50 num_train_workers=0,
51 num_val_workers=0,
52 )
53 trainer = config["trainer"]
54 trainer.update(
55 _target_="collectra.tasks.image_segmentation.sam3_metrics.create_trainer",
56 max_epochs=epochs, skip_saving_ckpts=True, val_epoch_freq=1, skip_first_val=False,
57 collectra_output=str(output), collectra_wandb=wandb,
58 collectra_run_name=run_name,
59 collectra_project=project,
60 )
61 trainer["model"] = {
62 "_target_": "collectra.tasks.image_segmentation.sam3_training.build_training_model",
63 "checkpoint": str(checkpoint),
64 "bpe_path": config["paths"]["bpe_path"],
65 }
66 for split, folder in (("train", "train"), ("val", "valid")):
67 data = trainer["data"][split]
68 data["dataset"].update(
69 img_folder=str(dataset / folder),
70 ann_file=str(dataset / folder / "_annotations.coco.json"),
71 load_segmentation=True,
72 )
73 data["drop_last"] = False
74 import json
76 valid = json.loads((dataset / "valid/_annotations.coco.json").read_text())
77 if valid["images"]:
78 evaluator = trainer["meters"]["val"]["roboflow100"]["detection"][
79 "pred_file_evaluators"
80 ][0]
81 evaluator["gt_path"] = str(dataset / "valid/_annotations.coco.json")
82 trainer["meters"]["val"]["roboflow100"]["mask_iou"] = {
83 "_target_": "collectra.tasks.image_segmentation.sam3_metrics.MaskIoU",
84 "confidence_threshold": confidence_threshold,
85 }
86 else:
87 # Training without held-out images must not silently evaluate on training data.
88 trainer["data"].pop("val", None)
89 trainer["meters"] = {}
90 trainer["mode"] = "train_only"
91 config["launcher"].update(gpus_per_node=num_gpus, num_nodes=1)
92 config["submitit"].update(use_cluster=False, job_array=None)
93 return config
96def load_image_state(model, state):
97 """Load image weights, adapting SAM3.1's three-level detection neck."""
98 if any(key.startswith("detector.") for key in state):
99 state = {
100 key.removeprefix("detector."): value
101 for key, value in state.items()
102 if key.startswith("detector.")
103 }
104 prefix = "backbone.vision_backbone."
105 tracking = (prefix + "interactive_convs.", prefix + "propagation_convs.")
106 if any(key.startswith(tracking) for key in state):
107 state = {key: value for key, value in state.items() if not key.startswith(tracking)}
108 # SAM3.1 stores three detection scales. The image builder creates four,
109 # then SAM3VLBackbone(scalp=1) discards the fourth. Initialize only that
110 # unused scale from the equivalent stride-1 convolutions; never tolerate
111 # missing weights in any feature level actually consumed by the model.
112 neck = model.backbone.vision_backbone
113 if model.backbone.scalp != 1 or tuple(neck.scale_factors) != (4.0, 2.0, 1.0, 0.5):
114 raise ValueError("Unsupported SAM3 image backbone for SAM3.1 checkpoint")
115 for suffix in ("conv_1x1.weight", "conv_1x1.bias", "conv_3x3.weight", "conv_3x3.bias"):
116 target = prefix + "convs.3." + suffix
117 if target not in state:
118 state[target] = state[prefix + "convs.2." + suffix].clone()
119 model.load_state_dict(state, strict=True)
122def _training_mlp_forward(self, x):
123 """SAM3 ViT MLP with ordinary autograd-capable PyTorch operations."""
124 x = self.act(self.fc1(x))
125 x = self.drop1(x)
126 x = self.norm(x)
127 x = self.fc2(x)
128 return self.drop2(x)
131def enable_training_mlps(model):
132 from sam3.model.vitdet import Mlp
134 # Upstream Mlp.forward uses addmm_act, an inference-only fused kernel.
135 # Bind only this model's MLPs; keep parameter names/checkpoints intact and
136 # leave SAM3's inference implementation untouched.
137 for module in model.modules():
138 if isinstance(module, Mlp):
139 module.forward = MethodType(_training_mlp_forward, module)
142def freeze_unused_semantic_head(model):
143 """Instance-mask losses do not supervise the separate semantic output."""
144 segmentation = getattr(model, "segmentation_head", None)
145 semantic = getattr(segmentation, "semantic_seg_head", None)
146 if semantic is not None:
147 # DDP sees this head in forward's outputs, so find_unused_parameters
148 # cannot detect that the loss never consumes it. Freeze its parameters
149 # before DDP/optimizer construction, preserving checkpoint compatibility.
150 semantic.requires_grad_(False)
153def build_training_model(checkpoint, bpe_path):
154 import torch
155 from sam3.model_builder import build_sam3_image_model
157 state = torch.load(checkpoint, map_location="cpu", weights_only=True)
158 state = state.get("model", state)
159 model = build_sam3_image_model(
160 bpe_path=bpe_path,
161 load_from_HF=False,
162 enable_segmentation=True,
163 device="cpu",
164 eval_mode=False,
165 )
166 load_image_state(model, state)
167 enable_training_mlps(model)
168 freeze_unused_semantic_head(model)
169 return model
172def launch(config_path):
173 from hydra import initialize_config_dir
174 from sam3.train.train import main, register_omegaconf_resolvers
176 config_path = Path(config_path).resolve()
177 register_omegaconf_resolvers()
178 with initialize_config_dir(str(config_path.parent), version_base="1.2"):
179 main(
180 SimpleNamespace(
181 config=config_path.name, use_cluster=False, num_gpus=None, num_nodes=1
182 )
183 )