[Bug] KeyError: 'indices' in Sam3LossWrapper during Validation loop (Fine-tuning)

Hi there,
First of all, I would like to express my sincere gratitude for this great project and for making the code available. It is incredibly helpful for my research.

I am attempting to fine-tune SAM3 on a custom dataset (COCO format) using the provided training scripts. The training process runs perfectly without issues. However, as soon as the trainer enters the validation phase (run_val -> val_epoch), the process crashes with a KeyError: 'indices'.

It appears that during the validation step, the outputs dictionary passed to compute_loss inside Sam3LossWrapper does not contain the indices key, which implies that the Matcher might not be triggering or updating the dictionary correctly in the validation context, unlike in the training loop.

I am using a custom DebugSam3LossWrapper which inherits from Sam3LossWrapper solely to print debug info. The error originates from the parent class sam3_loss.py.

- Error Log & Traceback: The error occurs in sam3/train/loss/sam3_loss.py, line 104, inside compute_loss.
> terminate called without an active exception
[DEBUG][loss] KeyError in compute_loss: KeyError('indices')
[DEBUG][loss] outputs/targets detail:
outputs: dict keys=['encoder_hidden_states', 'prev_encoder_out', 'presence_feats', 'queries', 'presence_logit_dec', 'pred_logits', 'pred_boxes', 'pred_boxes_xyxy', 'pred_masks', 'semantic_seg', 'presence_logit']
encoder_hidden_states: tensor shape=(5184, 2, 256), dtype=torch.bfloat16, device=cuda:0
prev_encoder_out: dict keys=['encoder_out', 'backbone_out']
presence_feats: tensor shape=(1, 2, 256), dtype=torch.float32, device=cuda:0
queries: tensor shape=(2, 200, 256), dtype=torch.float32, device=cuda:0
presence_logit_dec: tensor shape=(2, 1), dtype=torch.bfloat16, device=cuda:0
targets: dict keys=['boxes', 'boxes_xyxy', 'boxes_padded', 'positive_map', 'num_boxes', 'masks', 'semantic_masks', 'is_valid_mask', 'is_exhaustive', 'object_ids_packed', 'object_ids_padded']
[DEBUG][loss] exception at step=2153: KeyError('indices')
[DEBUG][loss] step=2153 input summary:
find_stages: type= len=1
find_targets: type= len=1
targets(obj_count) first_items=[15] empty=0 unknown=0
[rank0]: Traceback (most recent call last):
[rank0]: File "/workspace/c244/sam3/sam3/train/train.py", line 339, in
[rank0]: main(args)
[rank0]: File "/workspace/c244/sam3/sam3/train/train.py", line 310, in main
[rank0]: single_node_runner(cfg, main_port)
[rank0]: File "/workspace/c244/sam3/sam3/train/train.py", line 71, in single_node_runner
[rank0]: single_proc_run(local_rank=0, main_port=main_port, cfg=cfg, world_size=num_proc)
[rank0]: File "/workspace/c244/sam3/sam3/train/train.py", line 58, in single_proc_run
[rank0]: trainer.run()
[rank0]: File "/workspace/c244/sam3/sam3/train/trainer.py", line 567, in run
[rank0]: self.run_train()
[rank0]: File "/workspace/c244/sam3/sam3/train/trainer.py", line 608, in run_train
[rank0]: self.run_val()
[rank0]: File "/workspace/c244/sam3/sam3/train/trainer.py", line 631, in run_val
[rank0]: outs = self.val_epoch(dataloader, phase=Phase.VAL)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/workspace/c244/sam3/sam3/train/trainer.py", line 695, in val_epoch
[rank0]: loss_dict, batch_size, extra_losses = self._step(
[rank0]: ^^^^^^^^^^^
[rank0]: File "/workspace/c244/sam3/sam3/train/trainer.py", line 506, in _step
[rank0]: loss = self._find_loss(key)(find_stages, find_targets)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/workspace/miniconda3/envs/c244/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
[rank0]: return self._call_impl(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/workspace/miniconda3/envs/c244/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl
[rank0]: return forward_call(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/workspace/c244/sam3/sam3/train/loss/debug_loss_wrapper.py", line 83, in forward
[rank0]: return super().forward(find_stages, find_targets)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/workspace/c244/sam3/sam3/train/loss/sam3_loss.py", line 173, in forward
[rank0]: cur_losses = self.compute_loss(outputs, targets)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/workspace/c244/sam3/sam3/train/loss/debug_loss_wrapper.py", line 92, in compute_loss
[rank0]: return super().compute_loss(outputs, targets)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/workspace/c244/sam3/sam3/train/loss/sam3_loss.py", line 104, in compute_loss
[rank0]: indices = out["indices"]
[rank0]: ~~~^^^^^^^^^^^
[rank0]: KeyError: 'indices'
[rank0]:[W204 08:28:41.901222244 ProcessGroupNCCL.cpp:1476] Warning: WARNING: destroy_process_group() was not called before program exit, which can leak resources. For more info, please see https://pytorch.org/docs/stable/distributed.html#shutdown (function operator())

- Configuration (.yaml): Here is the relevant configuration used.

```python
# @package _global_
defaults:
- _self_

# ============================================================================
# Paths Configuration (Change this to your own paths)
# ============================================================================
paths:
hsv_hel_mul_root: "/path/to/my_custom_dataset"
experiment_log_dir: "/path/to/log"
bpe_path: "/path/to/sam3/assets/bpe_simple_vocab_16e6.txt.gz" # This should be under sam3/assets/bpe_simple_vocab_16e6.txt.gz
checkpoint_path: "/path/to/sam3/sam3/assets/sam3.pt" # Use local checkpoint to avoid ProxyError

# Dataset configuration
hsv_hel_mul_train:
num_images: 1076 # Note: This is the number of images used for training. If null, all images are used.

# Training transforms pipeline
train_transforms:
- _target_: sam3.train.transforms.basic_for_api.ComposeAPI
transforms:
- _target_: sam3.train.transforms.filter_query_transforms.FlexibleFilterFindGetQueries
query_filter:
_target_: sam3.train.transforms.filter_query_transforms.FilterCrowds

- _target_: sam3.train.transforms.point_sampling.RandomizeInputBbox
box_noise_std: 0.1
box_noise_max: 20

- _target_: sam3.train.transforms.segmentation.DecodeRle

- _target_: sam3.train.transforms.basic_for_api.RandomResizeAPI
sizes:
_target_: sam3.train.transforms.basic.get_random_resize_scales
size: ${scratch.resolution}
min_size: 480
rounded: false
max_size:
_target_: sam3.train.transforms.basic.get_random_resize_max_size
size: ${scratch.resolution}
square: true
consistent_transform: ${scratch.consistent_transform}

- _target_: sam3.train.transforms.basic_for_api.PadToSizeAPI
size: ${scratch.resolution}
consistent_transform: ${scratch.consistent_transform}
- _target_: sam3.train.transforms.basic_for_api.ToTensorAPI

- _target_: sam3.train.transforms.filter_query_transforms.FlexibleFilterFindGetQueries
query_filter:
_target_: sam3.train.transforms.filter_query_transforms.FilterEmptyTargets

- _target_: sam3.train.transforms.basic_for_api.NormalizeAPI
mean: ${scratch.train_norm_mean}
std: ${scratch.train_norm_std}

- _target_: sam3.train.transforms.filter_query_transforms.FlexibleFilterFindGetQueries
query_filter:
_target_: sam3.train.transforms.filter_query_transforms.FilterEmptyTargets

- _target_: sam3.train.transforms.filter_query_transforms.FlexibleFilterFindGetQueries
query_filter:
_target_: sam3.train.transforms.filter_query_transforms.FilterFindQueriesWithTooManyOut
max_num_objects: ${scratch.max_ann_per_img}

# Validation transforms pipeline
val_transforms:
- _target_: sam3.train.transforms.basic_for_api.ComposeAPI
transforms:
- _target_: sam3.train.transforms.segmentation.DecodeRle
- _target_: sam3.train.transforms.basic_for_api.RandomResizeAPI
sizes: ${scratch.resolution}
max_size:
_target_: sam3.train.transforms.basic.get_random_resize_max_size
size: ${scratch.resolution}
square: true
consistent_transform: False
- _target_: sam3.train.transforms.basic_for_api.PadToSizeAPI
size: ${scratch.resolution}
consistent_transform: False
- _target_: sam3.train.transforms.basic_for_api.ToTensorAPI
- _target_: sam3.train.transforms.basic_for_api.NormalizeAPI
mean: ${scratch.train_norm_mean}
std: ${scratch.train_norm_std}

# loss config (no mask loss)
# loss:
# _target_: sam3.train.loss.sam3_loss.Sam3LossWrapper
# matcher: ${scratch.matcher}
# o2m_weight: 2.0
# o2m_matcher:
# _target_: sam3.train.matcher.BinaryOneToManyMatcher
# alpha: 0.3
# threshold: 0.4
# topk: 4
# use_o2m_matcher_on_o2m_aux: false # Another option is true
# loss_fns_find:
# - _target_: sam3.train.loss.loss_fns.Boxes
# weight_dict:
# loss_bbox: 5.0
# loss_giou: 2.0
# - _target_: sam3.train.loss.loss_fns.IABCEMdetr
# weak_loss: False
# weight_dict:
# loss_ce: 20.0 # Another option is 100.0
# presence_loss: 20.0
# pos_weight: 10.0 # Another option is 5.0
# alpha: 0.25
# gamma: 2
# use_presence: True # Change
# pos_focal: false
# pad_n_queries: 200
# pad_scale_pos: 1.0

# loss_fn_semantic_seg: null
# scale_by_find_batch_size: ${scratch.scale_by_find_batch_size}

# NOTE: Loss to be used for training in case of segmentation
loss:
# ----- Debug -----
_target_: sam3.train.loss.debug_loss_wrapper.DebugSam3LossWrapper # sam3.train.loss.sam3_loss.Sam3LossWrapper

debug_print_every: 50
debug_max_list_items: 5
debug_only_on_exception: false

# -----------------
matcher: ${scratch.matcher} # HungarianMatcherV2
o2m_weight: 2.0
o2m_matcher:
_target_: sam3.train.matcher.BinaryOneToManyMatcher
alpha: 0.3
threshold: 0.4
topk: 4
use_o2m_matcher_on_o2m_aux: false

loss_fns_find:
# Detection
- _target_: sam3.train.loss.loss_fns.Boxes
weight_dict:
loss_bbox: 5.0
loss_giou: 2.0
# Presence (Binary classification)
- _target_: sam3.train.loss.loss_fns.IABCEMdetr
weak_loss: False
weight_dict:
loss_ce: 20.0 # Another option is 100.0
presence_loss: 20.0
pos_weight: 10.0 # Another option is 5.0
alpha: 0.25
gamma: 2
use_presence: True
pos_focal: false
pad_n_queries: 200
pad_scale_pos: 1.0
# Instance segmentation
- _target_: sam3.train.loss.loss_fns.Masks
focal_alpha: 0.25
focal_gamma: 2.0
weight_dict:
loss_mask: 200.0
loss_dice: 10.0
compute_aux: false

loss_fn_semantic_seg: # Semantic segmentation
_target_: sam3.train.loss.loss_fns.SemanticSegCriterion
presence_head: True
presence_loss: False
focal: True
focal_alpha: 0.6
focal_gamma: 2.0
downsample: False
weight_dict:
loss_semantic_seg: 20.0
loss_semantic_presence: 1.0
loss_semantic_dice: 30.0
scale_by_find_batch_size: ${scratch.scale_by_find_batch_size}

# ============================================================================
# Different helper parameters and functions
# ============================================================================
scratch:
enable_segmentation: True # NOTE: This is the number of queries used for segmentation
# Model parameters
d_model: 256
pos_embed:
_target_: sam3.model.position_encoding.PositionEmbeddingSine
num_pos_feats: ${scratch.d_model}
normalize: true
scale: null
temperature: 10000

# Box processing
use_presence_eval: True
original_box_postprocessor:
_target_: sam3.eval.postprocessors.PostProcessImage
max_dets_per_img: -1 # infinite detections
use_original_ids: true
use_original_sizes_box: true
use_presence: ${scratch.use_presence_eval}
iou_type: "bbox"

# Mask processing
original_segm_postprocessor:
_target_: sam3.eval.postprocessors.PostProcessImage
max_dets_per_img: -1
use_original_ids: true
use_original_sizes_box: true
use_original_sizes_mask: true
use_presence: ${scratch.use_presence_eval}
# always_interpolate_masks_on_gpu: false
convert_mask_to_rle: false
iou_type: "segm"

# Matcher configuration (Hungarian + One-to-Many)
matcher:
_target_: sam3.train.matcher.BinaryHungarianMatcherV2
focal: true # with `focal: true` it is equivalent to BinaryFocalHungarianMatcher
cost_class: 2.0
cost_bbox: 5.0
cost_giou: 2.0
alpha: 0.25
gamma: 2
stable: False
scale_by_find_batch_size: True

# Image processing parameters
resolution: 1008
consistent_transform: False
max_ann_per_img: 200

# Normalization parameters
train_norm_mean: [0.5, 0.5, 0.5]
train_norm_std: [0.5, 0.5, 0.5]
val_norm_mean: [0.5, 0.5, 0.5]
val_norm_std: [0.5, 0.5, 0.5]

# Training parameters
num_train_workers: 24
num_val_workers: 0 # 8
max_data_epochs: 20
target_epoch_size: 1500
hybrid_repeats: 1
context_length: 2
gather_pred_via_filesys: false

# Learning rate and scheduler parameters
lr_scale: 0.1
lr_transformer: ${times:8e-4,${scratch.lr_scale}}
lr_vision_backbone: ${times:2.5e-4,${scratch.lr_scale}}
lr_language_backbone: ${times:5e-5,${scratch.lr_scale}}
lrd_vision_backbone: 0.9
wd: 0.1
scheduler_timescale: 20
scheduler_warmup: 20
scheduler_cooldown: 20

val_batch_size: 1
collate_fn_val:
_target_: sam3.train.data.collator.collate_fn_api
_partial_: true
repeats: ${scratch.hybrid_repeats}
dict_key: all
with_seg_masks: ${scratch.enable_segmentation} # Note: Set this to true if using segmentation masks!

gradient_accumulation_steps: 1
train_batch_size: 1
collate_fn:
_target_: sam3.train.data.collator.collate_fn_api
_partial_: true
repeats: ${scratch.hybrid_repeats}
dict_key: all
with_seg_masks: ${scratch.enable_segmentation} # Note: Set this to true if using segmentation masks!

# ============================================================================
# Trainer Configuration
# ============================================================================

trainer:
_target_: sam3.train.trainer.Trainer
skip_saving_ckpts: False
empty_gpu_mem_cache_after_eval: True
skip_first_val: True
max_epochs: 20
accelerator: cuda
seed_value: 123
val_epoch_freq: 1
mode: train
gradient_accumulation_steps: ${scratch.gradient_accumulation_steps}

distributed:
backend: nccl
find_unused_parameters: True
gradient_as_bucket_view: True

loss:
all: ${hsv_hel_mul_train.loss}
default:
_target_: sam3.train.loss.sam3_loss.DummyLoss

data:
train:
_target_: sam3.train.data.torch_dataset.TorchDataset
dataset:
_target_: sam3.train.data.sam3_image_dataset.Sam3ImageDataset
limit_ids: ${hsv_hel_mul_train.num_images}
transforms: ${hsv_hel_mul_train.train_transforms}
load_segmentation: ${scratch.enable_segmentation}
max_ann_per_img: 500000
multiplier: 1
max_train_queries: 50000
max_val_queries: 50000
training: true
use_caching: False
img_folder: ${paths.hsv_hel_mul_root}/train/
ann_file: ${paths.hsv_hel_mul_root}/train/_annotations.coco.json

shuffle: True
batch_size: ${scratch.train_batch_size}
num_workers: ${scratch.num_train_workers}
pin_memory: True
drop_last: True
collate_fn: ${scratch.collate_fn}

val:
_target_: sam3.train.data.torch_dataset.TorchDataset
dataset:
_target_: sam3.train.data.sam3_image_dataset.Sam3ImageDataset
load_segmentation: ${scratch.enable_segmentation}
coco_json_loader:
_target_: sam3.train.data.coco_json_loaders.COCO_FROM_JSON
include_negatives: true
category_chunk_size: 2
_partial_: true
img_folder: ${paths.hsv_hel_mul_root}/valid/
ann_file: ${paths.hsv_hel_mul_root}/valid/_annotations.coco.json
transforms: ${hsv_hel_mul_train.val_transforms}
max_ann_per_img: 100000
multiplier: 1
training: false

shuffle: False
batch_size: ${scratch.val_batch_size}
num_workers: ${scratch.num_val_workers}
pin_memory: True
drop_last: False
collate_fn: ${scratch.collate_fn_val}

model:
_target_: sam3.model_builder.build_sam3_image_model
bpe_path: ${paths.bpe_path}
checkpoint_path: ${paths.checkpoint_path}
device: cpus
load_from_HF: false
eval_mode: false
enable_segmentation: ${scratch.enable_segmentation} # Warning: Enable this if using segmentation.

meters:
val:
all: # this key matches the "dict_key" in the dataloader's collate function
detection:
_target_: sam3.eval.coco_writer.PredictionDumper
iou_type: "bbox"
dump_dir: ${launcher.experiment_log_dir}/dumps/hsv_hel_mul_val_det/
merge_predictions: True
postprocessor: ${scratch.original_box_postprocessor}
gather_pred_via_filesys: ${scratch.gather_pred_via_filesys}
maxdets: 100
pred_file_evaluators:
- _target_: sam3.eval.coco_eval_offline.CocoEvaluatorOfflineWithPredFileEvaluators
gt_path: ${paths.hsv_hel_mul_root}/valid/_annotations.coco.json
tide: False
iou_type: "bbox"

# New: segmentation meter
segmentation:
_target_: sam3.eval.coco_writer.PredictionDumper
iou_type: "segm"
dump_dir: ${launcher.experiment_log_dir}/dumps/hsv_hel_mul_val_segm
merge_predictions: True
postprocessor: ${scratch.original_segm_postprocessor}
gather_pred_via_filesys: ${scratch.gather_pred_via_filesys}
maxdets: 100
pred_file_evaluators:
- _target_: sam3.eval.coco_eval_offline.CocoEvaluatorOfflineWithPredFileEvaluators
gt_path: ${paths.hsv_hel_mul_root}/valid/_annotations.coco.json
tide: False
iou_type: "segm"

optim:
amp:
enabled: True
amp_dtype: bfloat16

optimizer:
_target_: torch.optim.AdamW

gradient_clip:
_target_: sam3.train.optim.optimizer.GradientClipper
max_norm: 0.1
norm_type: 2

param_group_modifiers:
- _target_: sam3.train.optim.optimizer.layer_decay_param_modifier
_partial_: True
layer_decay_value: ${scratch.lrd_vision_backbone}
apply_to: 'backbone.vision_backbone.trunk'
overrides:
- pattern: '*pos_embed*'
value: 1.0

options:
lr:
- scheduler: # transformer and class_embed
_target_: sam3.train.optim.schedulers.InverseSquareRootParamScheduler
base_lr: ${scratch.lr_transformer}
timescale: ${scratch.scheduler_timescale}
warmup_steps: ${scratch.scheduler_warmup}
cooldown_steps: ${scratch.scheduler_cooldown}
- scheduler:
_target_: sam3.train.optim.schedulers.InverseSquareRootParamScheduler
base_lr: ${scratch.lr_vision_backbone}
timescale: ${scratch.scheduler_timescale}
warmup_steps: ${scratch.scheduler_warmup}
cooldown_steps: ${scratch.scheduler_cooldown}
param_names:
- 'backbone.vision_backbone.*'
- scheduler:
_target_: sam3.train.optim.schedulers.InverseSquareRootParamScheduler
base_lr: ${scratch.lr_language_backbone}
timescale: ${scratch.scheduler_timescale}
warmup_steps: ${scratch.scheduler_warmup}
cooldown_steps: ${scratch.scheduler_cooldown}
param_names:
- 'backbone.language_backbone.*'

weight_decay:
- scheduler:
_target_: fvcore.common.param_scheduler.ConstantParamScheduler
value: ${scratch.wd}
- scheduler:
_target_: fvcore.common.param_scheduler.ConstantParamScheduler
value: 0.0
param_names:
- '*bias*'
module_cls_names: ['torch.nn.LayerNorm']

checkpoint:
save_dir: ${launcher.experiment_log_dir}/checkpoints
save_freq: 5 # 0 only last checkpoint is saved.

logging:
tensorboard_writer:
_target_: sam3.train.utils.logger.make_tensorboard_logger
log_dir: ${launcher.experiment_log_dir}/tensorboard
flush_secs: 120
should_log: True
wandb_writer: null
log_dir: ${launcher.experiment_log_dir}/logs
log_freq: 10

# ============================================================================
# Launcher and Submitit Configuration
# ============================================================================

launcher:
num_nodes: 1
gpus_per_node: 1
experiment_log_dir: ${paths.experiment_log_dir}
multiprocessing_context: forkserver

submitit:
account: null
partition: null
qos: null
timeout_hour: 72
use_cluster: False
cpus_per_task: 10
port_range: [10000, 65000]
constraint: null
# # Uncomment for job array configuration
# job_array:
# num_tasks: 100
# task_index: 0

# ============================================================================
# Available HSV_HEL_MUL Supercategories (for reference)
# ============================================================================

all_hsv_hel_mul_supercategories:
- Herpes Simplex Virus Type-1 infected cell
- Herpes Simplex Virus Type-2 infected cell
```

It seems that Sam3LossWrapper.forward (or the underlying logic triggering the matcher) treats the validation phase differently, or the validation loop in Trainer invokes the loss function without pre-calculating the matching indices. Since compute_loss strictly requires out["indices"], it fails.

Is there a specific configuration required to enable the Matcher during the validation phase, or is this a bug in how val_epoch handles loss calculation?

Any guidance would be appreciated. Thanks!

Start at sam3/train/loss/sam3_loss.py line 104 and trace the validation path through sam3/train/trainer.py, especially run_val, val_epoch, and _step. Compare how outputs reach Sam3LossWrapper in training versus validation using the supplied configuration and traceback. Done means the validation loop completes without KeyError: 'indices' and reports its losses normally.

Written by the indexing model from the issue text.