From 5e198d7399ec0245f23a32b87df4c93a1c28d6a6 Mon Sep 17 00:00:00 2001 From: AI Lab <129358391+1038lab@users.noreply.github.com> Date: Mon, 24 Nov 2025 20:58:18 -0800 Subject: [PATCH] Delete train directory --- train/__init__.py | 1 - train/configs/eval_base.yaml | 279 ---- .../sam3_gold_image_attributes.yaml | 66 - .../sam3_gold_image_crowded.yaml | 66 - .../sam3_gold_image_fg_food.yaml | 66 - .../sam3_gold_image_fg_sports.yaml | 66 - .../sam3_gold_image_metaclip_nps.yaml | 66 - .../sam3_gold_image_sa1b_nps.yaml | 66 - .../sam3_gold_image_wiki_common.yaml | 66 - .../odinw13/odinw_text_and_visual.yaml | 255 --- train/configs/odinw13/odinw_text_only.yaml | 253 --- .../odinw13/odinw_text_only_positive.yaml | 253 --- .../odinw13/odinw_text_only_train.yaml | 591 ------- train/configs/odinw13/odinw_visual_only.yaml | 256 --- .../roboflow_v100/roboflow_v100_eval.yaml | 539 ------- .../roboflow_v100_full_ft_100_images.yaml | 539 ------- .../saco_video_evals/saco_veval_sav_test.yaml | 174 -- .../saco_veval_sav_test_noheur.yaml | 174 -- .../saco_video_evals/saco_veval_sav_val.yaml | 174 -- .../saco_veval_sav_val_noheur.yaml | 174 -- .../saco_veval_smartglasses_test.yaml | 174 -- .../saco_veval_smartglasses_test_noheur.yaml | 174 -- .../saco_veval_smartglasses_val.yaml | 174 -- .../saco_veval_smartglasses_val_noheur.yaml | 174 -- .../saco_veval_yt1b_test.yaml | 174 -- .../saco_veval_yt1b_test_noheur.yaml | 174 -- .../saco_video_evals/saco_veval_yt1b_val.yaml | 174 -- .../saco_veval_yt1b_val_noheur.yaml | 174 -- .../sam3_silver_image_bdd100k.yaml | 64 - .../sam3_silver_image_droid.yaml | 64 - .../sam3_silver_image_ego4d.yaml | 64 - .../sam3_silver_image_fathomnet.yaml | 64 - .../sam3_silver_image_food_rec.yaml | 64 - .../sam3_silver_image_geode.yaml | 64 - .../sam3_silver_image_inaturalist.yaml | 64 - .../sam3_silver_image_nga.yaml | 64 - .../sam3_silver_image_sav.yaml | 64 - .../sam3_silver_image_yt1b.yaml | 64 - train/data/__init__.py | 1 - train/data/coco_json_loaders.py | 465 ------ train/data/collator.py | 360 ----- train/data/sam3_image_dataset.py | 528 ------- train/data/sam3_video_dataset.py | 327 ---- train/data/torch_dataset.py | 52 - train/loss/__init__.py | 1 - train/loss/loss_fns.py | 1319 ---------------- train/loss/mask_sampling.py | 113 -- train/loss/sam3_loss.py | 203 --- train/loss/sigmoid_focal_loss.py | 321 ---- train/masks_ops.py | 272 ---- train/matcher.py | 806 ---------- train/nms_helper.py | 306 ---- train/optim/__init__.py | 1 - train/optim/optimizer.py | 498 ------ train/optim/schedulers.py | 41 - train/train.py | 339 ---- train/trainer.py | 1193 -------------- train/transforms/__init__.py | 1 - train/transforms/basic.py | 455 ------ train/transforms/basic_for_api.py | 1396 ----------------- train/transforms/filter_query_transforms.py | 607 ------- train/transforms/point_sampling.py | 345 ---- train/transforms/segmentation.py | 157 -- train/utils/__init__.py | 1 - train/utils/checkpoint_utils.py | 358 ----- train/utils/distributed.py | 585 ------- train/utils/logger.py | 241 --- train/utils/train_utils.py | 285 ---- 68 files changed, 17733 deletions(-) delete mode 100644 train/__init__.py delete mode 100644 train/configs/eval_base.yaml delete mode 100644 train/configs/gold_image_evals/sam3_gold_image_attributes.yaml delete mode 100644 train/configs/gold_image_evals/sam3_gold_image_crowded.yaml delete mode 100644 train/configs/gold_image_evals/sam3_gold_image_fg_food.yaml delete mode 100644 train/configs/gold_image_evals/sam3_gold_image_fg_sports.yaml delete mode 100644 train/configs/gold_image_evals/sam3_gold_image_metaclip_nps.yaml delete mode 100644 train/configs/gold_image_evals/sam3_gold_image_sa1b_nps.yaml delete mode 100644 train/configs/gold_image_evals/sam3_gold_image_wiki_common.yaml delete mode 100644 train/configs/odinw13/odinw_text_and_visual.yaml delete mode 100644 train/configs/odinw13/odinw_text_only.yaml delete mode 100644 train/configs/odinw13/odinw_text_only_positive.yaml delete mode 100644 train/configs/odinw13/odinw_text_only_train.yaml delete mode 100644 train/configs/odinw13/odinw_visual_only.yaml delete mode 100644 train/configs/roboflow_v100/roboflow_v100_eval.yaml delete mode 100644 train/configs/roboflow_v100/roboflow_v100_full_ft_100_images.yaml delete mode 100644 train/configs/saco_video_evals/saco_veval_sav_test.yaml delete mode 100644 train/configs/saco_video_evals/saco_veval_sav_test_noheur.yaml delete mode 100644 train/configs/saco_video_evals/saco_veval_sav_val.yaml delete mode 100644 train/configs/saco_video_evals/saco_veval_sav_val_noheur.yaml delete mode 100644 train/configs/saco_video_evals/saco_veval_smartglasses_test.yaml delete mode 100644 train/configs/saco_video_evals/saco_veval_smartglasses_test_noheur.yaml delete mode 100644 train/configs/saco_video_evals/saco_veval_smartglasses_val.yaml delete mode 100644 train/configs/saco_video_evals/saco_veval_smartglasses_val_noheur.yaml delete mode 100644 train/configs/saco_video_evals/saco_veval_yt1b_test.yaml delete mode 100644 train/configs/saco_video_evals/saco_veval_yt1b_test_noheur.yaml delete mode 100644 train/configs/saco_video_evals/saco_veval_yt1b_val.yaml delete mode 100644 train/configs/saco_video_evals/saco_veval_yt1b_val_noheur.yaml delete mode 100644 train/configs/silver_image_evals/sam3_silver_image_bdd100k.yaml delete mode 100644 train/configs/silver_image_evals/sam3_silver_image_droid.yaml delete mode 100644 train/configs/silver_image_evals/sam3_silver_image_ego4d.yaml delete mode 100644 train/configs/silver_image_evals/sam3_silver_image_fathomnet.yaml delete mode 100644 train/configs/silver_image_evals/sam3_silver_image_food_rec.yaml delete mode 100644 train/configs/silver_image_evals/sam3_silver_image_geode.yaml delete mode 100644 train/configs/silver_image_evals/sam3_silver_image_inaturalist.yaml delete mode 100644 train/configs/silver_image_evals/sam3_silver_image_nga.yaml delete mode 100644 train/configs/silver_image_evals/sam3_silver_image_sav.yaml delete mode 100644 train/configs/silver_image_evals/sam3_silver_image_yt1b.yaml delete mode 100644 train/data/__init__.py delete mode 100644 train/data/coco_json_loaders.py delete mode 100644 train/data/collator.py delete mode 100644 train/data/sam3_image_dataset.py delete mode 100644 train/data/sam3_video_dataset.py delete mode 100644 train/data/torch_dataset.py delete mode 100644 train/loss/__init__.py delete mode 100644 train/loss/loss_fns.py delete mode 100644 train/loss/mask_sampling.py delete mode 100644 train/loss/sam3_loss.py delete mode 100644 train/loss/sigmoid_focal_loss.py delete mode 100644 train/masks_ops.py delete mode 100644 train/matcher.py delete mode 100644 train/nms_helper.py delete mode 100644 train/optim/__init__.py delete mode 100644 train/optim/optimizer.py delete mode 100644 train/optim/schedulers.py delete mode 100644 train/train.py delete mode 100644 train/trainer.py delete mode 100644 train/transforms/__init__.py delete mode 100644 train/transforms/basic.py delete mode 100644 train/transforms/basic_for_api.py delete mode 100644 train/transforms/filter_query_transforms.py delete mode 100644 train/transforms/point_sampling.py delete mode 100644 train/transforms/segmentation.py delete mode 100644 train/utils/__init__.py delete mode 100644 train/utils/checkpoint_utils.py delete mode 100644 train/utils/distributed.py delete mode 100644 train/utils/logger.py delete mode 100644 train/utils/train_utils.py diff --git a/train/__init__.py b/train/__init__.py deleted file mode 100644 index 46d37d2..0000000 --- a/train/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved diff --git a/train/configs/eval_base.yaml b/train/configs/eval_base.yaml deleted file mode 100644 index 20890d8..0000000 --- a/train/configs/eval_base.yaml +++ /dev/null @@ -1,279 +0,0 @@ -# @package _global_ -defaults: - - _self_ - -# This config is the base configuration for all evaluations. Amongst other things, it defines: -# - the model -# - the image transforms -# - the post processors -# - cluster configuration (only relevant for slurm-based evals, ignored otherwise) -# -# Most of the parameters should be kept as-is. The main modifications you may want to make are: -# - the cluster configuration, to adjust partitions/qos to your system -# - the flag gather_pred_via_filesys if you ram is tight -# - num_val_workers if your number of cores is small (should be roughly number of cores / number of gpus) -# - the paths below - - -# ============================================================================ -# Paths Configuration (Chage this to your own paths) -# ============================================================================ -paths: - # If you leave the checkpoint path to null, the model will be downloaded from hugging-face. Otherwise provide a path - checkpoint_path: null - # the experiments will be subfolders of this - base_experiment_log_dir: - - # base path to the annotation folder for gold (refer to the readmes on how to download) - base_annotation_path: - - # base path to the annotation folder for silver (refer to the readmes on how to download) - base_annotation_path_silver: - - # path to the metaclip images, used for SA-Co gold (refer to the readme for instructions). Can be null if you don't intend on evaluating on this dataset. - metaclip_img_path: - - # path to the sa1b images, used for SA-Co gold (refer to the readme for instructions). Can be null if you don't intend on evaluating on this dataset. - sa1b_img_path: - - # path to the SA-Co/silver images - silver_img_path: - - bpe_path: # This should be under assets/bpe_simple_vocab_16e6.txt.gz - - -# ============================================================================ -# Different helper parameters and functions -# ============================================================================ -scratch: - - use_presence_eval: True - - base_val_transform: - - _target_: sam3.train.transforms.basic_for_api.ComposeAPI - transforms: - ######## transforms for validation (begin) ######## - - _target_: sam3.train.transforms.basic_for_api.RandomResizeAPI - sizes: ${scratch.resolution} # originally `resolution: 1024` - max_size: - _target_: sam3.train.transforms.basic.get_random_resize_max_size - size: ${scratch.resolution} # originally `resolution: 1024` - square: true - consistent_transform: False - ######## transforms for validation (end) ######## - - _target_: sam3.train.transforms.basic_for_api.ToTensorAPI - - _target_: sam3.train.transforms.basic_for_api.NormalizeAPI - mean: ${scratch.val_norm_mean} - std: ${scratch.val_norm_std} - - loss: null - - # Model parameters - d_model: 256 - input_box_embedding_dim: ${add:${scratch.d_model},2} - - # Box processing - 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} - - box_postprocessor: - _target_: sam3.eval.postprocessors.PostProcessImage - max_dets_per_img: -1 #infinite detections - use_original_ids: false - use_original_sizes_box: false - use_presence: ${scratch.use_presence_eval} - - box_postprocessor_thresholded: - _target_: sam3.eval.postprocessors.PostProcessImage - max_dets_per_img: -1 #infinite detections - use_original_ids: false - use_original_sizes_box: false - detection_threshold: 0.3 - use_presence: ${scratch.use_presence_eval} - - mask_postprocessor_thresholded: - _target_: sam3.eval.postprocessors.PostProcessImage - max_dets_per_img: -1 #infinite detections - iou_type: "segm" - use_original_ids: false - use_original_sizes_box: false - use_original_sizes_mask: true - convert_mask_to_rle: True - detection_threshold: 0.3 - use_presence: ${scratch.use_presence_eval} - - # Image processing parameters - resolution: 1008 - 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 - train_batch_size: 1 - val_batch_size: 1 - num_train_workers: 0 - num_val_workers: 10 # change this depending on the number of cpu cores available - max_data_epochs: 20 - target_epoch_size: 1500 - hybrid_repeats: 1 - context_length: 2 - - # All reduce - this controls how the predictions are sent back to node 0. - # If you have a lot of ram, CPU gather is faster. Otherwise, we provide a fallback through filesystem (eg NFS) - # Switch to true if you get cpu ooms during gather. - gather_pred_via_filesys: false - - # Learning rate and scheduler parameters (unused for eval) - 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 # (lower for in-domain adn higher for ood) - wd: 0.1 - scheduler_timescale: 20 - scheduler_warmup: 20 - scheduler_cooldown: 20 - - -# ============================================================================ -# Trainer Configuration -# ============================================================================ - -trainer: - _target_: sam3.train.trainer.Trainer - skip_saving_ckpts: true - empty_gpu_mem_cache_after_eval: True - skip_first_val: True - max_epochs: ${scratch.max_data_epochs} - accelerator: cuda - seed_value: 123 - val_epoch_freq: 10 - mode: val - - distributed: - backend: nccl - find_unused_parameters: True - gradient_as_bucket_view: True - - loss: - all: - _target_: sam3.train.loss.sam3_loss.DummyLoss - default: - _target_: sam3.train.loss.sam3_loss.DummyLoss - - data: - train: null - val: null - - model: - _target_: sam3.model_builder.build_sam3_image_model - bpe_path: ${paths.bpe_path} - device: cpus - eval_mode: true - enable_segmentation: true # Warning: Enable this if using segmentation. - checkpoint_path: ${paths.checkpoint_path} - - meters: - val: null - - 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: 0 # 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: 4 - gpus_per_node: 8 - experiment_log_dir: ${paths.experiment_log_dir} - multiprocessing_context: forkserver - - -submitit: - account: null # Add your SLURM account if use_cluster == 1 - partition: null - qos: null # Add your QoS if use_cluster == 1 - timeout_hour: 72 - use_cluster: True - cpus_per_task: 10 - port_range: [10000, 65000] - constraint: null diff --git a/train/configs/gold_image_evals/sam3_gold_image_attributes.yaml b/train/configs/gold_image_evals/sam3_gold_image_attributes.yaml deleted file mode 100644 index 8646b69..0000000 --- a/train/configs/gold_image_evals/sam3_gold_image_attributes.yaml +++ /dev/null @@ -1,66 +0,0 @@ -# @package _global_ -defaults: - - /configs/eval_base.yaml - - _self_ - -# ============================================================================ -# Paths Configuration (you can override here, but it shouldn't require further changes if eval_base.yaml is correct -# ============================================================================ -paths: - experiment_log_dir: ${paths.base_experiment_log_dir}/gold_attributes/ - coco_gt: ${paths.base_annotation_path}/gold_attributes_merged_a_release_test.json - coco_gts: - - ${paths.base_annotation_path}/gold_attributes_merged_a_release_test.json - - ${paths.base_annotation_path}/gold_attributes_merged_b_release_test.json - - ${paths.base_annotation_path}/gold_attributes_merged_c_release_test.json - - -# ============================================================================ -# Trainer Configuration -# ============================================================================ - -trainer: - data: - val: - _target_: sam3.train.data.torch_dataset.TorchDataset - dataset: - _target_: sam3.train.data.sam3_image_dataset.Sam3ImageDataset - coco_json_loader: - _target_: sam3.train.data.coco_json_loaders.SAM3_EVAL_API_FROM_JSON_NP - _partial_: true - img_folder: ${paths.metaclip_img_path} - ann_file: ${paths.coco_gt} - transforms: ${scratch.base_val_transform} - 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: False - drop_last: False - collate_fn: - _target_: sam3.train.data.collator.collate_fn_api - _partial_: true - repeats: ${scratch.hybrid_repeats} - dict_key: gold_attributes - - meters: - val: - gold_attributes: # this key matches the "dict_key" in the dataloader's collate function - cgf1: - _target_: sam3.eval.coco_writer.PredictionDumper - iou_type: "segm" - dump_dir: ${launcher.experiment_log_dir}/dumps/gold_attributes - merge_predictions: True - postprocessor: ${scratch.mask_postprocessor_thresholded} - gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} - maxdets: 1000000 # no limit - pred_file_evaluators: - - _target_: sam3.eval.cgf1_eval.CGF1Evaluator - gt_path: ${paths.coco_gts} - iou_type: "bbox" - - _target_: sam3.eval.cgf1_eval.CGF1Evaluator - gt_path: ${paths.coco_gts} - iou_type: "segm" diff --git a/train/configs/gold_image_evals/sam3_gold_image_crowded.yaml b/train/configs/gold_image_evals/sam3_gold_image_crowded.yaml deleted file mode 100644 index fef74a6..0000000 --- a/train/configs/gold_image_evals/sam3_gold_image_crowded.yaml +++ /dev/null @@ -1,66 +0,0 @@ -# @package _global_ -defaults: - - /configs/eval_base.yaml - - _self_ - -# ============================================================================ -# Paths Configuration (you can override here, but it shouldn't require further changes if eval_base.yaml is correct -# ============================================================================ -paths: - experiment_log_dir: ${paths.base_experiment_log_dir}/gold_crowded/ - coco_gt: ${paths.base_annotation_path}/gold_crowded_merged_a_release_test.json - coco_gts: - - ${paths.base_annotation_path}/gold_crowded_merged_a_release_test.json - - ${paths.base_annotation_path}/gold_crowded_merged_b_release_test.json - - ${paths.base_annotation_path}/gold_crowded_merged_c_release_test.json - - -# ============================================================================ -# Trainer Configuration -# ============================================================================ - -trainer: - data: - val: - _target_: sam3.train.data.torch_dataset.TorchDataset - dataset: - _target_: sam3.train.data.sam3_image_dataset.Sam3ImageDataset - coco_json_loader: - _target_: sam3.train.data.coco_json_loaders.SAM3_EVAL_API_FROM_JSON_NP - _partial_: true - img_folder: ${paths.metaclip_img_path} - ann_file: ${paths.coco_gt} - transforms: ${scratch.base_val_transform} - 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: False - drop_last: False - collate_fn: - _target_: sam3.train.data.collator.collate_fn_api - _partial_: true - repeats: ${scratch.hybrid_repeats} - dict_key: gold_crowded - - meters: - val: - gold_crowded: # this key matches the "dict_key" in the dataloader's collate function - cgf1: - _target_: sam3.eval.coco_writer.PredictionDumper - iou_type: "segm" - dump_dir: ${launcher.experiment_log_dir}/dumps/gold_crowded - merge_predictions: True - postprocessor: ${scratch.mask_postprocessor_thresholded} - gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} - maxdets: 1000000 # no limit - pred_file_evaluators: - - _target_: sam3.eval.cgf1_eval.CGF1Evaluator - gt_path: ${paths.coco_gts} - iou_type: "bbox" - - _target_: sam3.eval.cgf1_eval.CGF1Evaluator - gt_path: ${paths.coco_gts} - iou_type: "segm" diff --git a/train/configs/gold_image_evals/sam3_gold_image_fg_food.yaml b/train/configs/gold_image_evals/sam3_gold_image_fg_food.yaml deleted file mode 100644 index b08c4a4..0000000 --- a/train/configs/gold_image_evals/sam3_gold_image_fg_food.yaml +++ /dev/null @@ -1,66 +0,0 @@ -# @package _global_ -defaults: - - /configs/eval_base.yaml - - _self_ - -# ============================================================================ -# Paths Configuration (you can override here, but it shouldn't require further changes if eval_base.yaml is correct -# ============================================================================ -paths: - experiment_log_dir: ${paths.base_experiment_log_dir}/gold_fg_food/ - coco_gt: ${paths.base_annotation_path}/gold_fg_food_merged_a_release_test.json - coco_gts: - - ${paths.base_annotation_path}/gold_fg_food_merged_a_release_test.json - - ${paths.base_annotation_path}/gold_fg_food_merged_b_release_test.json - - ${paths.base_annotation_path}/gold_fg_food_merged_c_release_test.json - - -# ============================================================================ -# Trainer Configuration -# ============================================================================ - -trainer: - data: - val: - _target_: sam3.train.data.torch_dataset.TorchDataset - dataset: - _target_: sam3.train.data.sam3_image_dataset.Sam3ImageDataset - coco_json_loader: - _target_: sam3.train.data.coco_json_loaders.SAM3_EVAL_API_FROM_JSON_NP - _partial_: true - img_folder: ${paths.metaclip_img_path} - ann_file: ${paths.coco_gt} - transforms: ${scratch.base_val_transform} - 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: False - drop_last: False - collate_fn: - _target_: sam3.train.data.collator.collate_fn_api - _partial_: true - repeats: ${scratch.hybrid_repeats} - dict_key: gold_fg_food - - meters: - val: - gold_fg_food: # this key matches the "dict_key" in the dataloader's collate function - cgf1: - _target_: sam3.eval.coco_writer.PredictionDumper - iou_type: "segm" - dump_dir: ${launcher.experiment_log_dir}/dumps/gold_fg_food - merge_predictions: True - postprocessor: ${scratch.mask_postprocessor_thresholded} - gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} - maxdets: 1000000 # no limit - pred_file_evaluators: - - _target_: sam3.eval.cgf1_eval.CGF1Evaluator - gt_path: ${paths.coco_gts} - iou_type: "bbox" - - _target_: sam3.eval.cgf1_eval.CGF1Evaluator - gt_path: ${paths.coco_gts} - iou_type: "segm" diff --git a/train/configs/gold_image_evals/sam3_gold_image_fg_sports.yaml b/train/configs/gold_image_evals/sam3_gold_image_fg_sports.yaml deleted file mode 100644 index 89a93be..0000000 --- a/train/configs/gold_image_evals/sam3_gold_image_fg_sports.yaml +++ /dev/null @@ -1,66 +0,0 @@ -# @package _global_ -defaults: - - /configs/eval_base.yaml - - _self_ - -# ============================================================================ -# Paths Configuration (you can override here, but it shouldn't require further changes if eval_base.yaml is correct -# ============================================================================ -paths: - experiment_log_dir: ${paths.base_experiment_log_dir}/gold_fg_sports_equipment/ - coco_gt: ${paths.base_annotation_path}/gold_fg_sports_equipment_merged_a_release_test.json - coco_gts: - - ${paths.base_annotation_path}/gold_fg_sports_equipment_merged_a_release_test.json - - ${paths.base_annotation_path}/gold_fg_sports_equipment_merged_b_release_test.json - - ${paths.base_annotation_path}/gold_fg_sports_equipment_merged_c_release_test.json - - -# ============================================================================ -# Trainer Configuration -# ============================================================================ - -trainer: - data: - val: - _target_: sam3.train.data.torch_dataset.TorchDataset - dataset: - _target_: sam3.train.data.sam3_image_dataset.Sam3ImageDataset - coco_json_loader: - _target_: sam3.train.data.coco_json_loaders.SAM3_EVAL_API_FROM_JSON_NP - _partial_: true - img_folder: ${paths.metaclip_img_path} - ann_file: ${paths.coco_gt} - transforms: ${scratch.base_val_transform} - 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: False - drop_last: False - collate_fn: - _target_: sam3.train.data.collator.collate_fn_api - _partial_: true - repeats: ${scratch.hybrid_repeats} - dict_key: gold_fg_sports_equipment - - meters: - val: - gold_fg_sports_equipment: # this key matches the "dict_key" in the dataloader's collate function - cgf1: - _target_: sam3.eval.coco_writer.PredictionDumper - iou_type: "segm" - dump_dir: ${launcher.experiment_log_dir}/dumps/gold_fg_sports_equipment - merge_predictions: True - postprocessor: ${scratch.mask_postprocessor_thresholded} - gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} - maxdets: 1000000 # no limit - pred_file_evaluators: - - _target_: sam3.eval.cgf1_eval.CGF1Evaluator - gt_path: ${paths.coco_gts} - iou_type: "bbox" - - _target_: sam3.eval.cgf1_eval.CGF1Evaluator - gt_path: ${paths.coco_gts} - iou_type: "segm" diff --git a/train/configs/gold_image_evals/sam3_gold_image_metaclip_nps.yaml b/train/configs/gold_image_evals/sam3_gold_image_metaclip_nps.yaml deleted file mode 100644 index e9c276f..0000000 --- a/train/configs/gold_image_evals/sam3_gold_image_metaclip_nps.yaml +++ /dev/null @@ -1,66 +0,0 @@ -# @package _global_ -defaults: - - /configs/eval_base.yaml - - _self_ - -# ============================================================================ -# Paths Configuration (you can override here, but it shouldn't require further changes if eval_base.yaml is correct -# ============================================================================ -paths: - experiment_log_dir: ${paths.base_experiment_log_dir}/gold_metaclip_nps/ - coco_gt: ${paths.base_annotation_path}/gold_metaclip_merged_a_release_test.json - coco_gts: - - ${paths.base_annotation_path}/gold_metaclip_merged_a_release_test.json - - ${paths.base_annotation_path}/gold_metaclip_merged_b_release_test.json - - ${paths.base_annotation_path}/gold_metaclip_merged_c_release_test.json - - -# ============================================================================ -# Trainer Configuration -# ============================================================================ - -trainer: - data: - val: - _target_: sam3.train.data.torch_dataset.TorchDataset - dataset: - _target_: sam3.train.data.sam3_image_dataset.Sam3ImageDataset - coco_json_loader: - _target_: sam3.train.data.coco_json_loaders.SAM3_EVAL_API_FROM_JSON_NP - _partial_: true - img_folder: ${paths.metaclip_img_path} - ann_file: ${paths.coco_gt} - transforms: ${scratch.base_val_transform} - 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: False - drop_last: False - collate_fn: - _target_: sam3.train.data.collator.collate_fn_api - _partial_: true - repeats: ${scratch.hybrid_repeats} - dict_key: gold_metaclip_nps - - meters: - val: - gold_metaclip_nps: # this key matches the "dict_key" in the dataloader's collate function - cgf1: - _target_: sam3.eval.coco_writer.PredictionDumper - iou_type: "segm" - dump_dir: ${launcher.experiment_log_dir}/dumps/gold_metaclip_nps - merge_predictions: True - postprocessor: ${scratch.mask_postprocessor_thresholded} - gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} - maxdets: 1000000 # no limit - pred_file_evaluators: - - _target_: sam3.eval.cgf1_eval.CGF1Evaluator - gt_path: ${paths.coco_gts} - iou_type: "bbox" - - _target_: sam3.eval.cgf1_eval.CGF1Evaluator - gt_path: ${paths.coco_gts} - iou_type: "segm" diff --git a/train/configs/gold_image_evals/sam3_gold_image_sa1b_nps.yaml b/train/configs/gold_image_evals/sam3_gold_image_sa1b_nps.yaml deleted file mode 100644 index 52c87ee..0000000 --- a/train/configs/gold_image_evals/sam3_gold_image_sa1b_nps.yaml +++ /dev/null @@ -1,66 +0,0 @@ -# @package _global_ -defaults: - - /configs/eval_base.yaml - - _self_ - -# ============================================================================ -# Paths Configuration (you can override here, but it shouldn't require further changes if eval_base.yaml is correct -# ============================================================================ -paths: - experiment_log_dir: ${paths.base_experiment_log_dir}/gold_sa1b_nps/ - coco_gt: ${paths.base_annotation_path}/gold_sa1b_merged_a_release_test.json - coco_gts: - - ${paths.base_annotation_path}/gold_sa1b_merged_a_release_test.json - - ${paths.base_annotation_path}/gold_sa1b_merged_b_release_test.json - - ${paths.base_annotation_path}/gold_sa1b_merged_c_release_test.json - - -# ============================================================================ -# Trainer Configuration -# ============================================================================ - -trainer: - data: - val: - _target_: sam3.train.data.torch_dataset.TorchDataset - dataset: - _target_: sam3.train.data.sam3_image_dataset.Sam3ImageDataset - coco_json_loader: - _target_: sam3.train.data.coco_json_loaders.SAM3_EVAL_API_FROM_JSON_NP - _partial_: true - img_folder: ${paths.sa1b_img_path} - ann_file: ${paths.coco_gt} - transforms: ${scratch.base_val_transform} - 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: False - drop_last: False - collate_fn: - _target_: sam3.train.data.collator.collate_fn_api - _partial_: true - repeats: ${scratch.hybrid_repeats} - dict_key: gold_sa1b_nps - - meters: - val: - gold_sa1b_nps: # this key matches the "dict_key" in the dataloader's collate function - cgf1: - _target_: sam3.eval.coco_writer.PredictionDumper - iou_type: "segm" - dump_dir: ${launcher.experiment_log_dir}/dumps/gold_sa1b_nps - merge_predictions: True - postprocessor: ${scratch.mask_postprocessor_thresholded} - gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} - maxdets: 1000000 # no limit - pred_file_evaluators: - - _target_: sam3.eval.cgf1_eval.CGF1Evaluator - gt_path: ${paths.coco_gts} - iou_type: "bbox" - - _target_: sam3.eval.cgf1_eval.CGF1Evaluator - gt_path: ${paths.coco_gts} - iou_type: "segm" diff --git a/train/configs/gold_image_evals/sam3_gold_image_wiki_common.yaml b/train/configs/gold_image_evals/sam3_gold_image_wiki_common.yaml deleted file mode 100644 index 6304954..0000000 --- a/train/configs/gold_image_evals/sam3_gold_image_wiki_common.yaml +++ /dev/null @@ -1,66 +0,0 @@ -# @package _global_ -defaults: - - /configs/eval_base.yaml - - _self_ - -# ============================================================================ -# Paths Configuration (you can override here, but it shouldn't require further changes if eval_base.yaml is correct -# ============================================================================ -paths: - experiment_log_dir: ${paths.base_experiment_log_dir}/gold_wiki_common/ - coco_gt: ${paths.base_annotation_path}/gold_wiki_common_merged_a_release_test.json - coco_gts: - - ${paths.base_annotation_path}/gold_wiki_common_merged_a_release_test.json - - ${paths.base_annotation_path}/gold_wiki_common_merged_b_release_test.json - - ${paths.base_annotation_path}/gold_wiki_common_merged_c_release_test.json - - -# ============================================================================ -# Trainer Configuration -# ============================================================================ - -trainer: - data: - val: - _target_: sam3.train.data.torch_dataset.TorchDataset - dataset: - _target_: sam3.train.data.sam3_image_dataset.Sam3ImageDataset - coco_json_loader: - _target_: sam3.train.data.coco_json_loaders.SAM3_EVAL_API_FROM_JSON_NP - _partial_: true - img_folder: ${paths.metaclip_img_path} - ann_file: ${paths.coco_gt} - transforms: ${scratch.base_val_transform} - 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: False - drop_last: False - collate_fn: - _target_: sam3.train.data.collator.collate_fn_api - _partial_: true - repeats: ${scratch.hybrid_repeats} - dict_key: gold_wiki_common - - meters: - val: - gold_wiki_common: # this key matches the "dict_key" in the dataloader's collate function - cgf1: - _target_: sam3.eval.coco_writer.PredictionDumper - iou_type: "segm" - dump_dir: ${launcher.experiment_log_dir}/dumps/gold_wiki_common - merge_predictions: True - postprocessor: ${scratch.mask_postprocessor_thresholded} - gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} - maxdets: 1000000 # no limit - pred_file_evaluators: - - _target_: sam3.eval.cgf1_eval.CGF1Evaluator - gt_path: ${paths.coco_gts} - iou_type: "bbox" - - _target_: sam3.eval.cgf1_eval.CGF1Evaluator - gt_path: ${paths.coco_gts} - iou_type: "segm" diff --git a/train/configs/odinw13/odinw_text_and_visual.yaml b/train/configs/odinw13/odinw_text_and_visual.yaml deleted file mode 100644 index 51e93b4..0000000 --- a/train/configs/odinw13/odinw_text_and_visual.yaml +++ /dev/null @@ -1,255 +0,0 @@ -# @package _global_ -defaults: - - _self_ - -# ============================================================================ -# Paths Configuration (Chage this to your own paths) -# ============================================================================ -# python sam3/train/train.py -c configs/odinw_text_only.yaml --use-cluster 1 --partition ${PARTITION} --account ${ACCOUNT} --qos ${QoS} - -paths: - odinw_data_root: - experiment_log_dir: - bpe_path: # This should be under assets/bpe_simple_vocab_16e6.txt.gz - -supercategory_tuple: ${all_odinw_supercategories.${string:${submitit.job_array.task_index}}} -# Validation transforms pipeline -val_transforms: - - _target_: sam3.train.transforms.basic_for_api.ComposeAPI - transforms: - - _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.ToTensorAPI - - _target_: sam3.train.transforms.basic_for_api.NormalizeAPI - mean: ${scratch.val_norm_mean} - std: ${scratch.val_norm_std} - - _target_: sam3.train.transforms.filter_query_transforms.TextQueryToVisual - keep_text_queries: true # Note: set this to false if you only want visual - probability: 1.0 # always - -# ============================================================================ -# Different helper parameters and functions -# ============================================================================ -scratch: - enable_segmentation: True - # 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} - - # Image processing parameters - resolution: 1008 - # Normalization parameters - val_norm_mean: [0.5, 0.5, 0.5] - val_norm_std: [0.5, 0.5, 0.5] - - # Training parameters - val_batch_size: 2 - num_val_workers: 0 - gather_pred_via_filesys: false - -# ============================================================================ -# Trainer Configuration -# ============================================================================ - -trainer: - _target_: sam3.train.trainer.Trainer - skip_saving_ckpts: true - empty_gpu_mem_cache_after_eval: True - max_epochs: 1 - accelerator: cuda - seed_value: 123 - mode: val - - distributed: - backend: nccl - find_unused_parameters: True - gradient_as_bucket_view: True - - loss: - default: - _target_: sam3.train.loss.sam3_loss.DummyLoss - - data: - val: - _target_: sam3.train.data.torch_dataset.TorchDataset - dataset: - _target_: sam3.train.data.sam3_image_dataset.Sam3ImageDataset - coco_json_loader: - _target_: sam3.train.data.coco_json_loaders.COCO_FROM_JSON - prompts: ${odinw35_prompts.${supercategory_tuple.name}} - include_negatives: true - category_chunk_size: 20 # Note: Since we are doing AP +ve we need to include all categories! - _partial_: true - img_folder: ${paths.odinw_data_root}/${supercategory_tuple.val.img_folder} - ann_file: - _target_: sam3.eval.coco_reindex.reindex_coco_to_temp - input_json_path: ${paths.odinw_data_root}/${supercategory_tuple.val.json} - transforms: ${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: False - drop_last: False - collate_fn: - _target_: sam3.train.data.collator.collate_fn_api - _partial_: true - repeats: 1 - dict_key: odinw35 - - model: - _target_: sam3.model_builder.build_sam3_image_model - bpe_path: ${paths.bpe_path} - device: cpus - eval_mode: true # Set to false if training - enable_segmentation: ${scratch.enable_segmentation} # Warning: Enable this if using segmentation. - - meters: - val: - odinw35: - detection: - _target_: sam3.eval.coco_writer.PredictionDumper - iou_type: "bbox" - dump_dir: ${launcher.experiment_log_dir}/dumps/roboflow/${supercategory_tuple.name} - 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: - _target_: sam3.eval.coco_reindex.reindex_coco_to_temp - input_json_path: ${paths.odinw_data_root}/${supercategory_tuple.val.json} - tide: False - iou_type: "bbox" - positive_split: true - - checkpoint: - save_dir: ${launcher.experiment_log_dir}/checkpoints - save_freq: 0 # 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/${supercategory_tuple.name} - log_freq: 10 - -# ============================================================================ -# Launcher and Submitit Configuration -# ============================================================================ - -launcher: - num_nodes: 1 - gpus_per_node: 2 - experiment_log_dir: ${paths.experiment_log_dir} - multiprocessing_context: forkserver - -submitit: - account: null - partition: null - qos: null - timeout_hour: 72 - use_cluster: True - cpus_per_task: 10 - port_range: [10000, 65000] - constraint: null - - job_array: - num_tasks: 13 - task_index: 0 - -# ============================================================================ -# ODinW13 Supercategories -# ============================================================================ - -all_odinw_supercategories: - - name: AerialMaritimeDrone_large - val: - img_folder: AerialMaritimeDrone/large/test/ - json: AerialMaritimeDrone/large/test/annotations_without_background.json - - name: Aquarium - val: - img_folder: Aquarium/Aquarium Combined.v2-raw-1024.coco/test/ - json: Aquarium/Aquarium Combined.v2-raw-1024.coco/test/annotations_without_background.json - - name: CottontailRabbits - val: - img_folder: CottontailRabbits/test/ - json: CottontailRabbits/test/annotations_without_background.json - - name: EgoHands_generic - val: - img_folder: EgoHands/generic/test/ - json: EgoHands/generic/test/annotations_without_background.json - - name: NorthAmericaMushrooms - val: - img_folder: NorthAmericaMushrooms/North American Mushrooms.v1-416x416.coco/test/ - json: NorthAmericaMushrooms/North American Mushrooms.v1-416x416.coco/test/annotations_without_background.json - - name: Packages - val: - img_folder: Packages/Raw/test/ - json: Packages/Raw/test/annotations_without_background.json - - name: PascalVOC - val: - img_folder: PascalVOC/valid/ - json: PascalVOC/valid/annotations_without_background.json - - name: Raccoon - val: - img_folder: Raccoon/Raccoon.v2-raw.coco/test/ - json: Raccoon/Raccoon.v2-raw.coco/test/annotations_without_background.json - - name: ShellfishOpenImages - val: - img_folder: ShellfishOpenImages/raw/test/ - json: ShellfishOpenImages/raw/test/annotations_without_background.json - - name: VehiclesOpenImages - val: - img_folder: VehiclesOpenImages/416x416/test/ - json: VehiclesOpenImages/416x416/test/annotations_without_background.json - - name: pistols - val: - img_folder: pistols/export/ - json: pistols/export/test_annotations_without_background.json - - name: pothole - val: - img_folder: pothole/test/ - json: pothole/test/annotations_without_background.json - - name: thermalDogsAndPeople - val: - img_folder: thermalDogsAndPeople/test/ - json: thermalDogsAndPeople/test/annotations_without_background.json - - -odinw35_prompts: - AerialMaritimeDrone_large: '[{"id": 1, "name": "boat", "supercategory": "movable-objects"}, - {"id": 2, "name": "car", "supercategory": "movable-objects"}, {"id": 3, "name": "dock", - "supercategory": "movable-objects"}, {"id": 4, "name": "jet ski", "supercategory": "movable-objects"}, - {"id": 5, "name": "boat lift", "supercategory": "movable-objects"}]' - Aquarium: null - CottontailRabbits: null - EgoHands_generic: null - NorthAmericaMushrooms: '[{''id'': 1, ''name'': - ''chicken of the woods'', ''supercategory'': ''mushroom''}, {''id'': 2, ''name'': ''chanterelle'', ''supercategory'': ''mushroom''}]' - Packages: null - PascalVOC: null - Raccoon: null - ShellfishOpenImages: null - VehiclesOpenImages: null - pistols: null - pothole: null - thermalDogsAndPeople: null diff --git a/train/configs/odinw13/odinw_text_only.yaml b/train/configs/odinw13/odinw_text_only.yaml deleted file mode 100644 index e28fa5d..0000000 --- a/train/configs/odinw13/odinw_text_only.yaml +++ /dev/null @@ -1,253 +0,0 @@ -# @package _global_ -defaults: - - _self_ - -# ============================================================================ -# Paths Configuration (Chage this to your own paths) -# ============================================================================ -# python sam3/train/train.py -c configs/odinw_text_only.yaml --use-cluster 1 --partition ${PARTITION} --account ${ACCOUNT} --qos ${QoS} - -paths: - odinw_data_root: - experiment_log_dir: - bpe_path: # This should be under assets/bpe_simple_vocab_16e6.txt.gz - - -supercategory_tuple: ${all_odinw_supercategories.${string:${submitit.job_array.task_index}}} -# Validation transforms pipeline -val_transforms: - - _target_: sam3.train.transforms.basic_for_api.ComposeAPI - transforms: - - _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.ToTensorAPI - - _target_: sam3.train.transforms.basic_for_api.NormalizeAPI - mean: ${scratch.val_norm_mean} - std: ${scratch.val_norm_std} - -# ============================================================================ -# Different helper parameters and functions -# ============================================================================ -scratch: - enable_segmentation: True - # 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} - - # Image processing parameters - resolution: 1008 - # Normalization parameters - val_norm_mean: [0.5, 0.5, 0.5] - val_norm_std: [0.5, 0.5, 0.5] - - # Training parameters - val_batch_size: 2 - num_val_workers: 0 - gather_pred_via_filesys: false - -# ============================================================================ -# Trainer Configuration -# ============================================================================ - -trainer: - _target_: sam3.train.trainer.Trainer - skip_saving_ckpts: true - empty_gpu_mem_cache_after_eval: True - max_epochs: 1 - accelerator: cuda - seed_value: 123 - mode: val - - distributed: - backend: nccl - find_unused_parameters: True - gradient_as_bucket_view: True - - loss: - default: - _target_: sam3.train.loss.sam3_loss.DummyLoss - - data: - val: - _target_: sam3.train.data.torch_dataset.TorchDataset - dataset: - _target_: sam3.train.data.sam3_image_dataset.Sam3ImageDataset - coco_json_loader: - _target_: sam3.train.data.coco_json_loaders.COCO_FROM_JSON - prompts: ${odinw35_prompts.${supercategory_tuple.name}} - include_negatives: true - category_chunk_size: 20 # Note: Since we are doing AP +ve we need to include all categories! - _partial_: true - img_folder: ${paths.odinw_data_root}/${supercategory_tuple.val.img_folder} - ann_file: - _target_: sam3.eval.coco_reindex.reindex_coco_to_temp - input_json_path: ${paths.odinw_data_root}/${supercategory_tuple.val.json} - transforms: ${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: False - drop_last: False - collate_fn: - _target_: sam3.train.data.collator.collate_fn_api - _partial_: true - repeats: 1 - dict_key: odinw35 - - model: - _target_: sam3.model_builder.build_sam3_image_model - bpe_path: ${paths.bpe_path} - device: cpus - eval_mode: true # Set to false if training - enable_segmentation: ${scratch.enable_segmentation} # Warning: Enable this if using segmentation. - - meters: - val: - odinw35: - detection: - _target_: sam3.eval.coco_writer.PredictionDumper - iou_type: "bbox" - dump_dir: ${launcher.experiment_log_dir}/dumps/odinw/${supercategory_tuple.name} - 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: - _target_: sam3.eval.coco_reindex.reindex_coco_to_temp - input_json_path: ${paths.odinw_data_root}/${supercategory_tuple.val.json} - tide: False - iou_type: "bbox" - positive_split: False - - checkpoint: - save_dir: ${launcher.experiment_log_dir}/checkpoints - save_freq: 0 # 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/${supercategory_tuple.name} - log_freq: 10 - -# ============================================================================ -# Launcher and Submitit Configuration -# ============================================================================ - -launcher: - num_nodes: 1 - gpus_per_node: 2 - experiment_log_dir: ${paths.experiment_log_dir} - multiprocessing_context: forkserver - -submitit: - account: null - partition: null - qos: null - timeout_hour: 72 - use_cluster: True - cpus_per_task: 10 - port_range: [10000, 65000] - constraint: null - - job_array: - num_tasks: 13 - task_index: 0 - -# ============================================================================ -# ODinW13 Supercategories -# ============================================================================ - -all_odinw_supercategories: - - name: AerialMaritimeDrone_large - val: - img_folder: AerialMaritimeDrone/large/test/ - json: AerialMaritimeDrone/large/test/annotations_without_background.json - - name: Aquarium - val: - img_folder: Aquarium/Aquarium Combined.v2-raw-1024.coco/test/ - json: Aquarium/Aquarium Combined.v2-raw-1024.coco/test/annotations_without_background.json - - name: CottontailRabbits - val: - img_folder: CottontailRabbits/test/ - json: CottontailRabbits/test/annotations_without_background.json - - name: EgoHands_generic - val: - img_folder: EgoHands/generic/test/ - json: EgoHands/generic/test/annotations_without_background.json - - name: NorthAmericaMushrooms - val: - img_folder: NorthAmericaMushrooms/North American Mushrooms.v1-416x416.coco/test/ - json: NorthAmericaMushrooms/North American Mushrooms.v1-416x416.coco/test/annotations_without_background.json - - name: Packages - val: - img_folder: Packages/Raw/test/ - json: Packages/Raw/test/annotations_without_background.json - - name: PascalVOC - val: - img_folder: PascalVOC/valid/ - json: PascalVOC/valid/annotations_without_background.json - - name: Raccoon - val: - img_folder: Raccoon/Raccoon.v2-raw.coco/test/ - json: Raccoon/Raccoon.v2-raw.coco/test/annotations_without_background.json - - name: ShellfishOpenImages - val: - img_folder: ShellfishOpenImages/raw/test/ - json: ShellfishOpenImages/raw/test/annotations_without_background.json - - name: VehiclesOpenImages - val: - img_folder: VehiclesOpenImages/416x416/test/ - json: VehiclesOpenImages/416x416/test/annotations_without_background.json - - name: pistols - val: - img_folder: pistols/export/ - json: pistols/export/test_annotations_without_background.json - - name: pothole - val: - img_folder: pothole/test/ - json: pothole/test/annotations_without_background.json - - name: thermalDogsAndPeople - val: - img_folder: thermalDogsAndPeople/test/ - json: thermalDogsAndPeople/test/annotations_without_background.json - - -odinw35_prompts: - AerialMaritimeDrone_large: '[{"id": 1, "name": "boat", "supercategory": "movable-objects"}, - {"id": 2, "name": "car", "supercategory": "movable-objects"}, {"id": 3, "name": "dock", - "supercategory": "movable-objects"}, {"id": 4, "name": "jet ski", "supercategory": "movable-objects"}, - {"id": 5, "name": "boat lift", "supercategory": "movable-objects"}]' - Aquarium: null - CottontailRabbits: null - EgoHands_generic: null - NorthAmericaMushrooms: '[{''id'': 1, ''name'': - ''chicken of the woods'', ''supercategory'': ''mushroom''}, {''id'': 2, ''name'': ''chanterelle'', ''supercategory'': ''mushroom''}]' - Packages: null - PascalVOC: null - Raccoon: null - ShellfishOpenImages: null - VehiclesOpenImages: null - pistols: null - pothole: null - thermalDogsAndPeople: null diff --git a/train/configs/odinw13/odinw_text_only_positive.yaml b/train/configs/odinw13/odinw_text_only_positive.yaml deleted file mode 100644 index 9a86a52..0000000 --- a/train/configs/odinw13/odinw_text_only_positive.yaml +++ /dev/null @@ -1,253 +0,0 @@ -# @package _global_ -defaults: - - _self_ - -# ============================================================================ -# Paths Configuration (Chage this to your own paths) -# ============================================================================ -# python sam3/train/train.py -c configs/odinw_text_only.yaml --use-cluster 1 --partition ${PARTITION} --account ${ACCOUNT} --qos ${QoS} - -paths: - odinw_data_root: - experiment_log_dir: - bpe_path: # This should be under assets/bpe_simple_vocab_16e6.txt.gz - - -supercategory_tuple: ${all_odinw_supercategories.${string:${submitit.job_array.task_index}}} -# Validation transforms pipeline -val_transforms: - - _target_: sam3.train.transforms.basic_for_api.ComposeAPI - transforms: - - _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.ToTensorAPI - - _target_: sam3.train.transforms.basic_for_api.NormalizeAPI - mean: ${scratch.val_norm_mean} - std: ${scratch.val_norm_std} - -# ============================================================================ -# Different helper parameters and functions -# ============================================================================ -scratch: - enable_segmentation: True - # 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} - - # Image processing parameters - resolution: 1008 - # Normalization parameters - val_norm_mean: [0.5, 0.5, 0.5] - val_norm_std: [0.5, 0.5, 0.5] - - # Training parameters - val_batch_size: 2 - num_val_workers: 0 - gather_pred_via_filesys: false - -# ============================================================================ -# Trainer Configuration -# ============================================================================ - -trainer: - _target_: sam3.train.trainer.Trainer - skip_saving_ckpts: true - empty_gpu_mem_cache_after_eval: True - max_epochs: 1 - accelerator: cuda - seed_value: 123 - mode: val - - distributed: - backend: nccl - find_unused_parameters: True - gradient_as_bucket_view: True - - loss: - default: - _target_: sam3.train.loss.sam3_loss.DummyLoss - - data: - val: - _target_: sam3.train.data.torch_dataset.TorchDataset - dataset: - _target_: sam3.train.data.sam3_image_dataset.Sam3ImageDataset - coco_json_loader: - _target_: sam3.train.data.coco_json_loaders.COCO_FROM_JSON - prompts: ${odinw35_prompts.${supercategory_tuple.name}} - include_negatives: true - category_chunk_size: 20 # Note: Since we are doing AP +ve we need to include all categories! - _partial_: true - img_folder: ${paths.odinw_data_root}/${supercategory_tuple.val.img_folder} - ann_file: - _target_: sam3.eval.coco_reindex.reindex_coco_to_temp - input_json_path: ${paths.odinw_data_root}/${supercategory_tuple.val.json} - transforms: ${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: False - drop_last: False - collate_fn: - _target_: sam3.train.data.collator.collate_fn_api - _partial_: true - repeats: 1 - dict_key: odinw35 - - model: - _target_: sam3.model_builder.build_sam3_image_model - bpe_path: ${paths.bpe_path} - device: cpus - eval_mode: true # Set to false if training - enable_segmentation: ${scratch.enable_segmentation} # Warning: Enable this if using segmentation. - - meters: - val: - odinw35: - detection: - _target_: sam3.eval.coco_writer.PredictionDumper - iou_type: "bbox" - dump_dir: ${launcher.experiment_log_dir}/dumps/roboflow/${supercategory_tuple.name} - 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: - _target_: sam3.eval.coco_reindex.reindex_coco_to_temp - input_json_path: ${paths.odinw_data_root}/${supercategory_tuple.val.json} - tide: False - iou_type: "bbox" - positive_split: true - - checkpoint: - save_dir: ${launcher.experiment_log_dir}/checkpoints - save_freq: 0 # 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/${supercategory_tuple.name} - log_freq: 10 - -# ============================================================================ -# Launcher and Submitit Configuration -# ============================================================================ - -launcher: - num_nodes: 1 - gpus_per_node: 2 - experiment_log_dir: ${paths.experiment_log_dir} - multiprocessing_context: forkserver - -submitit: - account: null - partition: null - qos: null - timeout_hour: 72 - use_cluster: True - cpus_per_task: 10 - port_range: [10000, 65000] - constraint: null - - job_array: - num_tasks: 13 - task_index: 0 - -# ============================================================================ -# ODinW13 Supercategories -# ============================================================================ - -all_odinw_supercategories: - - name: AerialMaritimeDrone_large - val: - img_folder: AerialMaritimeDrone/large/test/ - json: AerialMaritimeDrone/large/test/annotations_without_background.json - - name: Aquarium - val: - img_folder: Aquarium/Aquarium Combined.v2-raw-1024.coco/test/ - json: Aquarium/Aquarium Combined.v2-raw-1024.coco/test/annotations_without_background.json - - name: CottontailRabbits - val: - img_folder: CottontailRabbits/test/ - json: CottontailRabbits/test/annotations_without_background.json - - name: EgoHands_generic - val: - img_folder: EgoHands/generic/test/ - json: EgoHands/generic/test/annotations_without_background.json - - name: NorthAmericaMushrooms - val: - img_folder: NorthAmericaMushrooms/North American Mushrooms.v1-416x416.coco/test/ - json: NorthAmericaMushrooms/North American Mushrooms.v1-416x416.coco/test/annotations_without_background.json - - name: Packages - val: - img_folder: Packages/Raw/test/ - json: Packages/Raw/test/annotations_without_background.json - - name: PascalVOC - val: - img_folder: PascalVOC/valid/ - json: PascalVOC/valid/annotations_without_background.json - - name: Raccoon - val: - img_folder: Raccoon/Raccoon.v2-raw.coco/test/ - json: Raccoon/Raccoon.v2-raw.coco/test/annotations_without_background.json - - name: ShellfishOpenImages - val: - img_folder: ShellfishOpenImages/raw/test/ - json: ShellfishOpenImages/raw/test/annotations_without_background.json - - name: VehiclesOpenImages - val: - img_folder: VehiclesOpenImages/416x416/test/ - json: VehiclesOpenImages/416x416/test/annotations_without_background.json - - name: pistols - val: - img_folder: pistols/export/ - json: pistols/export/test_annotations_without_background.json - - name: pothole - val: - img_folder: pothole/test/ - json: pothole/test/annotations_without_background.json - - name: thermalDogsAndPeople - val: - img_folder: thermalDogsAndPeople/test/ - json: thermalDogsAndPeople/test/annotations_without_background.json - - -odinw35_prompts: - AerialMaritimeDrone_large: '[{"id": 1, "name": "boat", "supercategory": "movable-objects"}, - {"id": 2, "name": "car", "supercategory": "movable-objects"}, {"id": 3, "name": "dock", - "supercategory": "movable-objects"}, {"id": 4, "name": "jet ski", "supercategory": "movable-objects"}, - {"id": 5, "name": "boat lift", "supercategory": "movable-objects"}]' - Aquarium: null - CottontailRabbits: null - EgoHands_generic: null - NorthAmericaMushrooms: '[{''id'': 1, ''name'': - ''chicken of the woods'', ''supercategory'': ''mushroom''}, {''id'': 2, ''name'': ''chanterelle'', ''supercategory'': ''mushroom''}]' - Packages: null - PascalVOC: null - Raccoon: null - ShellfishOpenImages: null - VehiclesOpenImages: null - pistols: null - pothole: null - thermalDogsAndPeople: null diff --git a/train/configs/odinw13/odinw_text_only_train.yaml b/train/configs/odinw13/odinw_text_only_train.yaml deleted file mode 100644 index eb03cdf..0000000 --- a/train/configs/odinw13/odinw_text_only_train.yaml +++ /dev/null @@ -1,591 +0,0 @@ -# @package _global_ -defaults: - - _self_ - -# ============================================================================ -# Paths Configuration (Chage this to your own paths) -# ============================================================================ -# python sam3/train/train.py -c configs/odinw_text_only.yaml --use-cluster 1 --partition ${PARTITION} --account ${ACCOUNT} --qos ${QoS} - -paths: - odinw_data_root: - experiment_log_dir: - bpe_path: # This should be under assets/bpe_simple_vocab_16e6.txt.gz - - -odinw_train: - train_file: fewshot_train_shot10_seed300 - num_images: null - supercategory_tuple: ${all_odinw_supercategories.${string:${submitit.job_array.task_index}}} - # 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.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.ToTensorAPI - - _target_: sam3.train.transforms.basic_for_api.NormalizeAPI - mean: ${scratch.val_norm_mean} - std: ${scratch.val_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: ${scratch.use_o2m_matcher_on_o2m_aux} - 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: ${scratch.loss_ce_weight} # Change - presence_loss: ${scratch.presence_weight} # Change - pos_weight: ${scratch.iabce_pos_weight} - alpha: ${scratch.iabce_alpha} - gamma: 2 - use_presence: True # Change - pos_focal: ${scratch.iabce_pos_focal} - pad_n_queries: ${scratch.num_queries} - pad_scale_pos: ${scratch.instance_query_loss_pad_scale_pos} - - loss_fn_semantic_seg: null - scale_by_find_batch_size: ${scratch.scale_by_find_batch_size} - - -# ============================================================================ -# Different helper parameters and functions -# ============================================================================ -scratch: - enable_segmentation: False - use_act_checkpoint_geo_encoder: True - input_geometry_encoder: - _target_: sam3.model.geometry_encoders.SequenceGeometryEncoder - pos_enc: ${scratch.pos_embed} - encode_boxes_as_points: False - points_direct_project: True - points_pool: True - points_pos_enc: True - boxes_direct_project: True - boxes_pool: True - boxes_pos_enc: True - d_model: ${scratch.d_model} - num_layers: 3 - use_act_ckpt: ${scratch.use_act_checkpoint_geo_encoder} - layer: - _target_: sam3.model.encoder.TransformerEncoderLayer - activation: "relu" - d_model: ${scratch.d_model} - dim_feedforward: 2048 - dropout: ${scratch.encoder_dropout} - pos_enc_at_attn: false - pre_norm: True - pos_enc_at_cross_attn_queries: false - pos_enc_at_cross_attn_keys: true - self_attention: - _target_: sam3.model.attention.MultiheadAttention - attn_type: Vanilla - num_heads: 8 - dropout: ${scratch.encoder_dropout} - embed_dim: ${scratch.d_model} - batch_first: False - cross_attention: - _target_: sam3.model.attention.MultiheadAttention - attn_type: Vanilla - num_heads: 8 - dropout: ${scratch.encoder_dropout} - embed_dim: ${scratch.d_model} - batch_first: False - add_cls: true - add_post_encode_proj: True - - boxRPB: "log" - dac: True - use_early_fusion: true - o2m_mask: false - num_feature_levels: 1 # > 1 not implemented - encoder_dropout: 0.1 - decoder_dropout: 0.1 - - tokenizer_ve: - _target_: sam3.model.tokenizer_ve.SimpleTokenizer - bpe_path: ${paths.bpe_path} - - - freeze_text_tower: False - freeze_image_tower: NoFreeze - vis_backbone_dp: 0.0 - # Activation checkpointing (Save memory) - use_act_checkpoint_vision_backbone: True - use_act_checkpoint_text_backbone: True - use_act_checkpoint_encoder: True - use_act_checkpoint_decoder: True - - loss: null - # Loss parameters - num_queries: 200 - presence_weight: 20.0 - loss_ce_weight: 20.0 - iabce_pos_weight: 5.0 - iabce_pos_focal: false - iabce_alpha: 0.25 - instance_query_loss_pad_scale_pos: 1.0 - use_o2m_matcher_on_o2m_aux: false - - # Model parameters - use_instance_query: true - 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} - - - # Matcher configuration - matcher: - _target_: sam3.train.matcher.BinaryHungarianMatcherV2 - focal: true - 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 - train_batch_size: 1 - val_batch_size: 1 - num_train_workers: 0 - num_val_workers: 0 - max_data_epochs: 40 - 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 - - - - -# ============================================================================ -# Trainer Configuration -# ============================================================================ - -trainer: - _target_: sam3.train.trainer.Trainer - skip_saving_ckpts: true - # _target_: sam3.train.trainer.Trainer - # skip_saving_ckpts: true - empty_gpu_mem_cache_after_eval: True - skip_first_val: True - max_epochs: ${scratch.max_data_epochs} - accelerator: cuda - seed_value: 123 - val_epoch_freq: 10 - mode: train - - distributed: - backend: nccl - find_unused_parameters: True - gradient_as_bucket_view: True - - loss: - all: ${odinw_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: ${odinw_train.num_images} - transforms: ${odinw_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.odinw_data_root}/${odinw_train.supercategory_tuple.train.img_folder} - ann_file: - _target_: sam3.eval.coco_reindex.reindex_coco_to_temp - input_json_path: ${paths.odinw_data_root}/${odinw_train.supercategory_tuple.train.json} - coco_json_loader: - _target_: sam3.train.data.coco_json_loaders.COCO_FROM_JSON - prompts: ${odinw35_prompts.${odinw_train.supercategory_tuple.name}} #${odinw_train.supercategory_tuple.name) - _partial_: true - shuffle: True - batch_size: ${scratch.train_batch_size} - num_workers: ${scratch.num_train_workers} - pin_memory: False - drop_last: True - 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} - - 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 - prompts: ${odinw35_prompts.${odinw_train.supercategory_tuple.name}} - include_negatives: true - category_chunk_size: 20 # Note: Since we are doing AP +ve we need to include all categories! - _partial_: true - img_folder: ${paths.odinw_data_root}/${odinw_train.supercategory_tuple.val.img_folder} - ann_file: - _target_: sam3.eval.coco_reindex.reindex_coco_to_temp - input_json_path: ${paths.odinw_data_root}/${odinw_train.supercategory_tuple.val.json} - transforms: ${odinw_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: False - drop_last: False - collate_fn: - _target_: sam3.train.data.collator.collate_fn_api - _partial_: true - repeats: 1 - dict_key: odinw35 - with_seg_masks: ${scratch.enable_segmentation} - - model: - _target_: sam3.model_builder.build_sam3_image_model - bpe_path: ${paths.bpe_path} - device: cpus - eval_mode: false # Set to false if training - enable_segmentation: ${scratch.enable_segmentation} # Warning: Enable this if using segmentation. - - meters: - val: - odinw35: - detection: - _target_: sam3.eval.coco_writer.PredictionDumper - iou_type: "bbox" - dump_dir: ${launcher.experiment_log_dir}/dumps/odinw/${odinw_train.supercategory_tuple.name} - 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: - _target_: sam3.eval.coco_reindex.reindex_coco_to_temp - input_json_path: ${paths.odinw_data_root}/${odinw_train.supercategory_tuple.val.json} - tide: False - iou_type: "bbox" - positive_split: False - - 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: 0 # 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/${odinw_train.supercategory_tuple.name} - log_freq: 10 - -# ============================================================================ -# Launcher and Submitit Configuration -# ============================================================================ - -launcher: - num_nodes: 1 - gpus_per_node: 2 - experiment_log_dir: null #${paths.experiment_log_dir} - multiprocessing_context: forkserver - -submitit: - account: null - partition: null - qos: null - timeout_hour: 72 - use_cluster: True - cpus_per_task: 10 - port_range: [10000, 65000] - constraint: null - - # task_index: 2 - # Uncomment for job array configuration - job_array: - num_tasks: 13 - task_index: 0 - - -# ============================================================================ -# ODinW13 Supercategories -# ============================================================================ - -all_odinw_supercategories: - - name: AerialMaritimeDrone_large - val: - img_folder: AerialMaritimeDrone/large/test/ - json: AerialMaritimeDrone/large/test/annotations_without_background.json - train: - img_folder: AerialMaritimeDrone/large/train/ - json: AerialMaritimeDrone/large/train/${odinw_train.train_file}.json - - name: Aquarium - val: - img_folder: Aquarium/Aquarium Combined.v2-raw-1024.coco/test/ - json: Aquarium/Aquarium Combined.v2-raw-1024.coco/test/annotations_without_background.json - train: - img_folder: Aquarium/Aquarium Combined.v2-raw-1024.coco/train/ - json: Aquarium/Aquarium Combined.v2-raw-1024.coco/train/${odinw_train.train_file}.json - - name: CottontailRabbits - val: - img_folder: CottontailRabbits/test/ - json: CottontailRabbits/test/annotations_without_background.json - train: - img_folder: CottontailRabbits/train/ - json: CottontailRabbits/train/${odinw_train.train_file}.json - - name: EgoHands_generic - val: - img_folder: EgoHands/generic/test/ - json: EgoHands/generic/test/annotations_without_background.json - train: - img_folder: EgoHands/generic/train/ - json: EgoHands/generic/train/${odinw_train.train_file}.json - - name: NorthAmericaMushrooms - val: - img_folder: NorthAmericaMushrooms/North American Mushrooms.v1-416x416.coco/test/ - json: NorthAmericaMushrooms/North American Mushrooms.v1-416x416.coco/test/annotations_without_background.json - train: - img_folder: NorthAmericaMushrooms/North American Mushrooms.v1-416x416.coco/train/ - json: NorthAmericaMushrooms/North American Mushrooms.v1-416x416.coco/train/${odinw_train.train_file}.json - - name: Packages - val: - img_folder: Packages/Raw/test/ - json: Packages/Raw/test/annotations_without_background.json - train: - img_folder: Packages/Raw/train/ - json: Packages/Raw/train/${odinw_train.train_file}.json - - name: PascalVOC - val: - img_folder: PascalVOC/valid/ - json: PascalVOC/valid/annotations_without_background.json - train: - img_folder: PascalVOC/train/ - json: PascalVOC/train/${odinw_train.train_file}.json - - name: Raccoon - val: - img_folder: Raccoon/Raccoon.v2-raw.coco/test/ - json: Raccoon/Raccoon.v2-raw.coco/test/annotations_without_background.json - train: - img_folder: Raccoon/Raccoon.v2-raw.coco/train/ - json: Raccoon/Raccoon.v2-raw.coco/train/${odinw_train.train_file}.json - - name: ShellfishOpenImages - val: - img_folder: ShellfishOpenImages/raw/test/ - json: ShellfishOpenImages/raw/test/annotations_without_background.json - train: - img_folder: ShellfishOpenImages/raw/train/ - json: ShellfishOpenImages/raw/train/${odinw_train.train_file}.json - - name: VehiclesOpenImages - val: - img_folder: VehiclesOpenImages/416x416/test/ - json: VehiclesOpenImages/416x416/test/annotations_without_background.json - train: - img_folder: VehiclesOpenImages/416x416/train/ - json: VehiclesOpenImages/416x416/train/${odinw_train.train_file}.json - - name: pistols - val: - img_folder: pistols/export/ - json: pistols/export/test_annotations_without_background.json - train: - img_folder: pistols/export/ - json: pistols/export/${odinw_train.train_file}.json - - name: pothole - val: - img_folder: pothole/test/ - json: pothole/test/annotations_without_background.json - train: - img_folder: pothole/train/ - json: pothole/train/${odinw_train.train_file}.json - - name: thermalDogsAndPeople - val: - img_folder: thermalDogsAndPeople/test/ - json: thermalDogsAndPeople/test/annotations_without_background.json - train: - img_folder: thermalDogsAndPeople/train/ - json: thermalDogsAndPeople/train/${odinw_train.train_file}.json - - -odinw35_prompts: - AerialMaritimeDrone_large: '[{"id": 1, "name": "boat", "supercategory": "movable-objects"}, - {"id": 2, "name": "car", "supercategory": "movable-objects"}, {"id": 3, "name": "dock", - "supercategory": "movable-objects"}, {"id": 4, "name": "jet ski", "supercategory": "movable-objects"}, - {"id": 5, "name": "boat lift", "supercategory": "movable-objects"}]' - Aquarium: null - CottontailRabbits: null - EgoHands_generic: null - NorthAmericaMushrooms: '[{''id'': 1, ''name'': - ''chicken of the woods'', ''supercategory'': ''mushroom''}, {''id'': 2, ''name'': ''chanterelle'', ''supercategory'': ''mushroom''}]' - Packages: null - PascalVOC: null - Raccoon: null - ShellfishOpenImages: null - VehiclesOpenImages: null - pistols: null - pothole: null - thermalDogsAndPeople: null diff --git a/train/configs/odinw13/odinw_visual_only.yaml b/train/configs/odinw13/odinw_visual_only.yaml deleted file mode 100644 index e724f2d..0000000 --- a/train/configs/odinw13/odinw_visual_only.yaml +++ /dev/null @@ -1,256 +0,0 @@ -# @package _global_ -defaults: - - _self_ - -# ============================================================================ -# Paths Configuration (Chage this to your own paths) -# ============================================================================ -# python sam3/train/train.py -c configs/odinw_text_only.yaml --use-cluster 1 --partition ${PARTITION} --account ${ACCOUNT} --qos ${QoS} - -paths: - odinw_data_root: - experiment_log_dir: - bpe_path: # This should be under assets/bpe_simple_vocab_16e6.txt.gz - - -supercategory_tuple: ${all_odinw_supercategories.${string:${submitit.job_array.task_index}}} -# Validation transforms pipeline -val_transforms: - - _target_: sam3.train.transforms.basic_for_api.ComposeAPI - transforms: - - _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.ToTensorAPI - - _target_: sam3.train.transforms.basic_for_api.NormalizeAPI - mean: ${scratch.val_norm_mean} - std: ${scratch.val_norm_std} - - _target_: sam3.train.transforms.filter_query_transforms.TextQueryToVisual - keep_text_queries: false # Note: set this to false if you only want visual - probability: 1.0 # always - -# ============================================================================ -# Different helper parameters and functions -# ============================================================================ -scratch: - enable_segmentation: True - # 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} - - # Image processing parameters - resolution: 1008 - # Normalization parameters - val_norm_mean: [0.5, 0.5, 0.5] - val_norm_std: [0.5, 0.5, 0.5] - - # Training parameters - val_batch_size: 2 - num_val_workers: 0 - gather_pred_via_filesys: false - -# ============================================================================ -# Trainer Configuration -# ============================================================================ - -trainer: - _target_: sam3.train.trainer.Trainer - skip_saving_ckpts: true - empty_gpu_mem_cache_after_eval: True - max_epochs: 1 - accelerator: cuda - seed_value: 123 - mode: val - - distributed: - backend: nccl - find_unused_parameters: True - gradient_as_bucket_view: True - - loss: - default: - _target_: sam3.train.loss.sam3_loss.DummyLoss - - data: - val: - _target_: sam3.train.data.torch_dataset.TorchDataset - dataset: - _target_: sam3.train.data.sam3_image_dataset.Sam3ImageDataset - coco_json_loader: - _target_: sam3.train.data.coco_json_loaders.COCO_FROM_JSON - prompts: ${odinw35_prompts.${supercategory_tuple.name}} - include_negatives: true - category_chunk_size: 20 # Note: Since we are doing AP +ve we need to include all categories! - _partial_: true - img_folder: ${paths.odinw_data_root}/${supercategory_tuple.val.img_folder} - ann_file: - _target_: sam3.eval.coco_reindex.reindex_coco_to_temp - input_json_path: ${paths.odinw_data_root}/${supercategory_tuple.val.json} - transforms: ${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: False - drop_last: False - collate_fn: - _target_: sam3.train.data.collator.collate_fn_api - _partial_: true - repeats: 1 - dict_key: odinw35 - - model: - _target_: sam3.model_builder.build_sam3_image_model - bpe_path: ${paths.bpe_path} - device: cpus - eval_mode: true # Set to false if training - enable_segmentation: ${scratch.enable_segmentation} # Warning: Enable this if using segmentation. - - meters: - val: - odinw35: - detection: - _target_: sam3.eval.coco_writer.PredictionDumper - iou_type: "bbox" - dump_dir: ${launcher.experiment_log_dir}/dumps/roboflow/${supercategory_tuple.name} - 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: - _target_: sam3.eval.coco_reindex.reindex_coco_to_temp - input_json_path: ${paths.odinw_data_root}/${supercategory_tuple.val.json} - tide: False - iou_type: "bbox" - positive_split: true - - checkpoint: - save_dir: ${launcher.experiment_log_dir}/checkpoints - save_freq: 0 # 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/${supercategory_tuple.name} - log_freq: 10 - -# ============================================================================ -# Launcher and Submitit Configuration -# ============================================================================ - -launcher: - num_nodes: 1 - gpus_per_node: 2 - experiment_log_dir: ${paths.experiment_log_dir} - multiprocessing_context: forkserver - -submitit: - account: null - partition: null - qos: null - timeout_hour: 72 - use_cluster: True - cpus_per_task: 10 - port_range: [10000, 65000] - constraint: null - - job_array: - num_tasks: 13 - task_index: 0 - -# ============================================================================ -# ODinW13 Supercategories -# ============================================================================ - -all_odinw_supercategories: - - name: AerialMaritimeDrone_large - val: - img_folder: AerialMaritimeDrone/large/test/ - json: AerialMaritimeDrone/large/test/annotations_without_background.json - - name: Aquarium - val: - img_folder: Aquarium/Aquarium Combined.v2-raw-1024.coco/test/ - json: Aquarium/Aquarium Combined.v2-raw-1024.coco/test/annotations_without_background.json - - name: CottontailRabbits - val: - img_folder: CottontailRabbits/test/ - json: CottontailRabbits/test/annotations_without_background.json - - name: EgoHands_generic - val: - img_folder: EgoHands/generic/test/ - json: EgoHands/generic/test/annotations_without_background.json - - name: NorthAmericaMushrooms - val: - img_folder: NorthAmericaMushrooms/North American Mushrooms.v1-416x416.coco/test/ - json: NorthAmericaMushrooms/North American Mushrooms.v1-416x416.coco/test/annotations_without_background.json - - name: Packages - val: - img_folder: Packages/Raw/test/ - json: Packages/Raw/test/annotations_without_background.json - - name: PascalVOC - val: - img_folder: PascalVOC/valid/ - json: PascalVOC/valid/annotations_without_background.json - - name: Raccoon - val: - img_folder: Raccoon/Raccoon.v2-raw.coco/test/ - json: Raccoon/Raccoon.v2-raw.coco/test/annotations_without_background.json - - name: ShellfishOpenImages - val: - img_folder: ShellfishOpenImages/raw/test/ - json: ShellfishOpenImages/raw/test/annotations_without_background.json - - name: VehiclesOpenImages - val: - img_folder: VehiclesOpenImages/416x416/test/ - json: VehiclesOpenImages/416x416/test/annotations_without_background.json - - name: pistols - val: - img_folder: pistols/export/ - json: pistols/export/test_annotations_without_background.json - - name: pothole - val: - img_folder: pothole/test/ - json: pothole/test/annotations_without_background.json - - name: thermalDogsAndPeople - val: - img_folder: thermalDogsAndPeople/test/ - json: thermalDogsAndPeople/test/annotations_without_background.json - - -odinw35_prompts: - AerialMaritimeDrone_large: '[{"id": 1, "name": "boat", "supercategory": "movable-objects"}, - {"id": 2, "name": "car", "supercategory": "movable-objects"}, {"id": 3, "name": "dock", - "supercategory": "movable-objects"}, {"id": 4, "name": "jet ski", "supercategory": "movable-objects"}, - {"id": 5, "name": "boat lift", "supercategory": "movable-objects"}]' - Aquarium: null - CottontailRabbits: null - EgoHands_generic: null - NorthAmericaMushrooms: '[{''id'': 1, ''name'': - ''chicken of the woods'', ''supercategory'': ''mushroom''}, {''id'': 2, ''name'': ''chanterelle'', ''supercategory'': ''mushroom''}]' - Packages: null - PascalVOC: null - Raccoon: null - ShellfishOpenImages: null - VehiclesOpenImages: null - pistols: null - pothole: null - thermalDogsAndPeople: null diff --git a/train/configs/roboflow_v100/roboflow_v100_eval.yaml b/train/configs/roboflow_v100/roboflow_v100_eval.yaml deleted file mode 100644 index 361e622..0000000 --- a/train/configs/roboflow_v100/roboflow_v100_eval.yaml +++ /dev/null @@ -1,539 +0,0 @@ -# @package _global_ -defaults: - - _self_ - -# ============================================================================ -# Paths Configuration (Chage this to your own paths) -# ============================================================================ -paths: - roboflow_vl_100_root: - experiment_log_dir: - bpe_path: # This should be under assets/bpe_simple_vocab_16e6.txt.gz - -# Roboflow dataset configuration -roboflow_train: - num_images: 100 # Note: This is the number of images used for training. If null, all images are used. - supercategory: ${all_roboflow_supercategories.${string:${submitit.job_array.task_index}}} - - # 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.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.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: - # _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 - # 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 - # - _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: - # _target_: sam3.losses.loss_fns.SemanticSegCriterion - # presence_head: True - # presence_loss: False # Change - # 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: False # 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} - - # Matcher configuration - 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: 10 - num_val_workers: 0 - 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: roboflow100 - 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: true - empty_gpu_mem_cache_after_eval: True - skip_first_val: True - max_epochs: 20 - accelerator: cuda - seed_value: 123 - val_epoch_freq: 10 - mode: val - gradient_accumulation_steps: ${scratch.gradient_accumulation_steps} - - distributed: - backend: nccl - find_unused_parameters: True - gradient_as_bucket_view: True - - loss: - all: ${roboflow_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: ${roboflow_train.num_images} - transforms: ${roboflow_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.roboflow_vl_100_root}/${roboflow_train.supercategory}/train/ - ann_file: ${paths.roboflow_vl_100_root}/${roboflow_train.supercategory}/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 # Note: You can increase this based on the memory of your GPU. - _partial_: true - img_folder: ${paths.roboflow_vl_100_root}/${roboflow_train.supercategory}/test/ - ann_file: ${paths.roboflow_vl_100_root}/${roboflow_train.supercategory}/test/_annotations.coco.json - transforms: ${roboflow_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} - device: cpus - eval_mode: true - enable_segmentation: ${scratch.enable_segmentation} # Warning: Enable this if using segmentation. - - meters: - val: - roboflow100: - detection: - _target_: sam3.eval.coco_writer.PredictionDumper - iou_type: "bbox" - dump_dir: ${launcher.experiment_log_dir}/dumps/roboflow/${roboflow_train.supercategory} - 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.roboflow_vl_100_root}/${roboflow_train.supercategory}/test/_annotations.coco.json - tide: False - iou_type: "bbox" - - 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: 0 # 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/${roboflow_train.supercategory} - log_freq: 10 - -# ============================================================================ -# Launcher and Submitit Configuration -# ============================================================================ - -launcher: - num_nodes: 1 - gpus_per_node: 2 - experiment_log_dir: ${paths.experiment_log_dir} - multiprocessing_context: forkserver - -submitit: - account: null - partition: null - qos: null - timeout_hour: 72 - use_cluster: True - cpus_per_task: 10 - port_range: [10000, 65000] - constraint: null - # Uncomment for job array configuration - job_array: - num_tasks: 100 - task_index: 0 - -# ============================================================================ -# Available Roboflow Supercategories (for reference) -# ============================================================================ - -all_roboflow_supercategories: - - -grccs - - zebrasatasturias - - cod-mw-warzone - - canalstenosis - - label-printing-defect-version-2 - - new-defects-in-wood - - orionproducts - - aquarium-combined - - varroa-mites-detection--test-set - - clashroyalechardetector - - stomata-cells - - halo-infinite-angel-videogame - - pig-detection - - urine-analysis1 - - aerial-sheep - - orgharvest - - actions - - mahjong - - liver-disease - - needle-base-tip-min-max - - wheel-defect-detection - - aircraft-turnaround-dataset - - xray - - wildfire-smoke - - spinefrxnormalvindr - - ufba-425 - - speech-bubbles-detection - - train - - pill - - truck-movement - - car-logo-detection - - inbreast - - sea-cucumbers-new-tiles - - uavdet-small - - penguin-finder-seg - - aerial-airport - - bibdetection - - taco-trash-annotations-in-context - - bees - - recode-waste - - screwdetectclassification - - wine-labels - - aerial-cows - - into-the-vale - - gwhd2021 - - lacrosse-object-detection - - defect-detection - - dataconvert - - x-ray-id - - ball - - tube - - 2024-frc - - crystal-clean-brain-tumors-mri-dataset - - grapes-5 - - human-detection-in-floods - - buoy-onboarding - - apoce-aerial-photographs-for-object-detection-of-construction-equipment - - l10ul502 - - floating-waste - - deeppcb - - ism-band-packet-detection - - weeds4 - - invoice-processing - - thermal-cheetah - - tomatoes-2 - - marine-sharks - - peixos-fish - - sssod - - aerial-pool - - countingpills - - asphaltdistressdetection - - roboflow-trained-dataset - - everdaynew - - underwater-objects - - soda-bottles - - dentalai - - jellyfish - - deepfruits - - activity-diagrams - - circuit-voltages - - all-elements - - macro-segmentation - - exploratorium-daphnia - - signatures - - conveyor-t-shirts - - fruitjes - - grass-weeds - - infraredimageofpowerequipment - - 13-lkc01 - - wb-prova - - flir-camera-objects - - paper-parts - - football-player-detection - - trail-camera - - smd-components - - water-meter - - nih-xray - - the-dreidel-project - - electric-pylon-detection-in-rsi - - cable-damage diff --git a/train/configs/roboflow_v100/roboflow_v100_full_ft_100_images.yaml b/train/configs/roboflow_v100/roboflow_v100_full_ft_100_images.yaml deleted file mode 100644 index 6b95f62..0000000 --- a/train/configs/roboflow_v100/roboflow_v100_full_ft_100_images.yaml +++ /dev/null @@ -1,539 +0,0 @@ -# @package _global_ -defaults: - - _self_ - -# ============================================================================ -# Paths Configuration (Chage this to your own paths) -# ============================================================================ -paths: - roboflow_vl_100_root: - experiment_log_dir: - bpe_path: # This should be under assets/bpe_simple_vocab_16e6.txt.gz - -# Roboflow dataset configuration -roboflow_train: - num_images: 100 # Note: This is the number of images used for training. If null, all images are used. - supercategory: ${all_roboflow_supercategories.${string:${submitit.job_array.task_index}}} - - # 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.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.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: - # _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 - # 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 - # - _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: - # _target_: sam3.losses.loss_fns.SemanticSegCriterion - # presence_head: True - # presence_loss: False # Change - # 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: False # 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} - - # Matcher configuration - 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: 10 - num_val_workers: 0 - 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: roboflow100 - 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: true - empty_gpu_mem_cache_after_eval: True - skip_first_val: True - max_epochs: 20 - accelerator: cuda - seed_value: 123 - val_epoch_freq: 10 - mode: train - gradient_accumulation_steps: ${scratch.gradient_accumulation_steps} - - distributed: - backend: nccl - find_unused_parameters: True - gradient_as_bucket_view: True - - loss: - all: ${roboflow_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: ${roboflow_train.num_images} - transforms: ${roboflow_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.roboflow_vl_100_root}/${roboflow_train.supercategory}/train/ - ann_file: ${paths.roboflow_vl_100_root}/${roboflow_train.supercategory}/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 # Note: You can increase this based on the memory of your GPU. - _partial_: true - img_folder: ${paths.roboflow_vl_100_root}/${roboflow_train.supercategory}/test/ - ann_file: ${paths.roboflow_vl_100_root}/${roboflow_train.supercategory}/test/_annotations.coco.json - transforms: ${roboflow_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} - device: cpus - eval_mode: false - enable_segmentation: ${scratch.enable_segmentation} # Warning: Enable this if using segmentation. - - meters: - val: - roboflow100: - detection: - _target_: sam3.eval.coco_writer.PredictionDumper - iou_type: "bbox" - dump_dir: ${launcher.experiment_log_dir}/dumps/roboflow/${roboflow_train.supercategory} - 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.roboflow_vl_100_root}/${roboflow_train.supercategory}/test/_annotations.coco.json - tide: False - iou_type: "bbox" - - 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: 0 # 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/${roboflow_train.supercategory} - log_freq: 10 - -# ============================================================================ -# Launcher and Submitit Configuration -# ============================================================================ - -launcher: - num_nodes: 1 - gpus_per_node: 2 - experiment_log_dir: ${paths.experiment_log_dir} - multiprocessing_context: forkserver - -submitit: - account: null - partition: null - qos: null - timeout_hour: 72 - use_cluster: True - cpus_per_task: 10 - port_range: [10000, 65000] - constraint: null - # Uncomment for job array configuration - job_array: - num_tasks: 100 - task_index: 0 - -# ============================================================================ -# Available Roboflow Supercategories (for reference) -# ============================================================================ - -all_roboflow_supercategories: - - -grccs - - zebrasatasturias - - cod-mw-warzone - - canalstenosis - - label-printing-defect-version-2 - - new-defects-in-wood - - orionproducts - - aquarium-combined - - varroa-mites-detection--test-set - - clashroyalechardetector - - stomata-cells - - halo-infinite-angel-videogame - - pig-detection - - urine-analysis1 - - aerial-sheep - - orgharvest - - actions - - mahjong - - liver-disease - - needle-base-tip-min-max - - wheel-defect-detection - - aircraft-turnaround-dataset - - xray - - wildfire-smoke - - spinefrxnormalvindr - - ufba-425 - - speech-bubbles-detection - - train - - pill - - truck-movement - - car-logo-detection - - inbreast - - sea-cucumbers-new-tiles - - uavdet-small - - penguin-finder-seg - - aerial-airport - - bibdetection - - taco-trash-annotations-in-context - - bees - - recode-waste - - screwdetectclassification - - wine-labels - - aerial-cows - - into-the-vale - - gwhd2021 - - lacrosse-object-detection - - defect-detection - - dataconvert - - x-ray-id - - ball - - tube - - 2024-frc - - crystal-clean-brain-tumors-mri-dataset - - grapes-5 - - human-detection-in-floods - - buoy-onboarding - - apoce-aerial-photographs-for-object-detection-of-construction-equipment - - l10ul502 - - floating-waste - - deeppcb - - ism-band-packet-detection - - weeds4 - - invoice-processing - - thermal-cheetah - - tomatoes-2 - - marine-sharks - - peixos-fish - - sssod - - aerial-pool - - countingpills - - asphaltdistressdetection - - roboflow-trained-dataset - - everdaynew - - underwater-objects - - soda-bottles - - dentalai - - jellyfish - - deepfruits - - activity-diagrams - - circuit-voltages - - all-elements - - macro-segmentation - - exploratorium-daphnia - - signatures - - conveyor-t-shirts - - fruitjes - - grass-weeds - - infraredimageofpowerequipment - - 13-lkc01 - - wb-prova - - flir-camera-objects - - paper-parts - - football-player-detection - - trail-camera - - smd-components - - water-meter - - nih-xray - - the-dreidel-project - - electric-pylon-detection-in-rsi - - cable-damage diff --git a/train/configs/saco_video_evals/saco_veval_sav_test.yaml b/train/configs/saco_video_evals/saco_veval_sav_test.yaml deleted file mode 100644 index b5bed47..0000000 --- a/train/configs/saco_video_evals/saco_veval_sav_test.yaml +++ /dev/null @@ -1,174 +0,0 @@ -# @package _global_ -defaults: - - _self_ - -# ============================================================================ -# Paths Configuration (Chage this to your own paths) -# ============================================================================ -paths: - - dump_file_name: saco_veval_sav_test - experiment_log_dir: - ytvis_json: /saco_veval_sav_test.json - ytvis_dir : - bpe_path: # This should be under assets/bpe_simple_vocab_16e6.txt.gz - num_videos: null - -# ============================================================================ -# Different helper parameters and functions -# ============================================================================ -scratch: - vid_mask_postprocessor: - _target_: sam3.eval.postprocessors.PostProcessNullOp - - use_presence_eval: True - - video_transforms_val: - - _target_: sam3.train.transforms.basic_for_api.ComposeAPI - transforms: - - _target_: sam3.train.transforms.segmentation.DecodeRle - # resize the image to 1024x1024 resolution - - _target_: sam3.train.transforms.basic_for_api.RandomResizeAPI - sizes: ${scratch.resolution} # originally `resolution: 1024` - square: true - consistent_transform: true - - _target_: sam3.train.transforms.basic_for_api.ToTensorAPI - - _target_: sam3.train.transforms.basic_for_api.NormalizeAPI - mean: ${scratch.val_norm_mean} - std: ${scratch.val_norm_std} - - # Model parameters - d_model: 256 - - # Image processing parameters - resolution: 1008 - - # 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] - - val_batch_size: 1 - num_val_workers: 0 - max_data_epochs: 20 - hybrid_repeats: 1 - gather_pred_via_filesys: false - - -# ============================================================================ -# Trainer Configuration -# ============================================================================ - -trainer: - _target_: sam3.train.trainer.Trainer - skip_saving_ckpts: true - empty_gpu_mem_cache_after_eval: True - skip_first_val: True - max_epochs: ${scratch.max_data_epochs} - accelerator: cuda - seed_value: 123 - val_epoch_freq: 10 - mode: val - - distributed: - backend: nccl - find_unused_parameters: True - gradient_as_bucket_view: True - - loss: - all: - _target_: sam3.train.loss.sam3_loss.DummyLoss - default: - _target_: sam3.train.loss.sam3_loss.DummyLoss - - data: - train: null - val: - _target_: sam3.train.data.torch_dataset.TorchDataset - dataset: - _target_: sam3.train.data.sam3_video_dataset.VideoGroundingDataset - limit_ids: ${paths.num_videos} - img_folder: ${paths.ytvis_dir} - ann_file: ${paths.ytvis_json} - coco_json_loader: - _target_: sam3.train.data.coco_json_loaders.SAM3_VEVAL_API_FROM_JSON_NP - _partial_: true - - transforms: ${scratch.video_transforms_val} - max_ann_per_img: 100000 # filtered in transforms - max_val_queries: 100000 - multiplier: 1 - load_segmentation: true - training: false - - - shuffle: False - batch_size: ${scratch.val_batch_size} - num_workers: ${scratch.num_val_workers} - pin_memory: True - drop_last: False - collate_fn: - _target_: sam3.train.data.collator.collate_fn_api - _partial_: true - repeats: ${scratch.hybrid_repeats} - dict_key: ytvis_val - with_seg_masks: true - - - model: - _target_: sam3.model_builder.build_sam3_video_model - bpe_path: ${paths.bpe_path} - has_presence_token: True - geo_encoder_use_img_cross_attn: True - apply_temporal_disambiguation: True - - meters: - val: - ytvis_val: - pred_file: # key - _target_: sam3.eval.ytvis_eval.YTVISResultsWriter - dump_file: ${launcher.experiment_log_dir}/preds/${paths.dump_file_name}.json - postprocessor: ${scratch.vid_mask_postprocessor} - gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} - - optim: - amp: - enabled: True - amp_dtype: bfloat16 - - - checkpoint: - save_dir: ${launcher.experiment_log_dir}/checkpoints - save_freq: 0 # 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: 8 - gpus_per_node: 8 - experiment_log_dir: ${paths.experiment_log_dir} - multiprocessing_context: forkserver - -submitit: - account: null - partition: null - qos: null - timeout_hour: 72 - use_cluster: True - cpus_per_task: 10 - port_range: [10000, 65000] - constraint: null diff --git a/train/configs/saco_video_evals/saco_veval_sav_test_noheur.yaml b/train/configs/saco_video_evals/saco_veval_sav_test_noheur.yaml deleted file mode 100644 index abc3289..0000000 --- a/train/configs/saco_video_evals/saco_veval_sav_test_noheur.yaml +++ /dev/null @@ -1,174 +0,0 @@ -# @package _global_ -defaults: - - _self_ - -# ============================================================================ -# Paths Configuration (Chage this to your own paths) -# ============================================================================ -paths: - - dump_file_name: saco_veval_sav_test - experiment_log_dir: - ytvis_json: /saco_veval_sav_test.json - ytvis_dir : - bpe_path: # This should be under assets/bpe_simple_vocab_16e6.txt.gz - num_videos: null - -# ============================================================================ -# Different helper parameters and functions -# ============================================================================ -scratch: - vid_mask_postprocessor: - _target_: sam3.eval.postprocessors.PostProcessNullOp - - use_presence_eval: True - - video_transforms_val: - - _target_: sam3.train.transforms.basic_for_api.ComposeAPI - transforms: - - _target_: sam3.train.transforms.segmentation.DecodeRle - # resize the image to 1024x1024 resolution - - _target_: sam3.train.transforms.basic_for_api.RandomResizeAPI - sizes: ${scratch.resolution} # originally `resolution: 1024` - square: true - consistent_transform: true - - _target_: sam3.train.transforms.basic_for_api.ToTensorAPI - - _target_: sam3.train.transforms.basic_for_api.NormalizeAPI - mean: ${scratch.val_norm_mean} - std: ${scratch.val_norm_std} - - # Model parameters - d_model: 256 - - # Image processing parameters - resolution: 1008 - - # 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] - - val_batch_size: 1 - num_val_workers: 0 - max_data_epochs: 20 - hybrid_repeats: 1 - gather_pred_via_filesys: false - - -# ============================================================================ -# Trainer Configuration -# ============================================================================ - -trainer: - _target_: sam3.train.trainer.Trainer - skip_saving_ckpts: true - empty_gpu_mem_cache_after_eval: True - skip_first_val: True - max_epochs: ${scratch.max_data_epochs} - accelerator: cuda - seed_value: 123 - val_epoch_freq: 10 - mode: val - - distributed: - backend: nccl - find_unused_parameters: True - gradient_as_bucket_view: True - - loss: - all: - _target_: sam3.train.loss.sam3_loss.DummyLoss - default: - _target_: sam3.train.loss.sam3_loss.DummyLoss - - data: - train: null - val: - _target_: sam3.train.data.torch_dataset.TorchDataset - dataset: - _target_: sam3.train.data.sam3_video_dataset.VideoGroundingDataset - limit_ids: ${paths.num_videos} - img_folder: ${paths.ytvis_dir} - ann_file: ${paths.ytvis_json} - coco_json_loader: - _target_: sam3.train.data.coco_json_loaders.SAM3_VEVAL_API_FROM_JSON_NP - _partial_: true - - transforms: ${scratch.video_transforms_val} - max_ann_per_img: 100000 # filtered in transforms - max_val_queries: 100000 - multiplier: 1 - load_segmentation: true - training: false - - - shuffle: False - batch_size: ${scratch.val_batch_size} - num_workers: ${scratch.num_val_workers} - pin_memory: True - drop_last: False - collate_fn: - _target_: sam3.train.data.collator.collate_fn_api - _partial_: true - repeats: ${scratch.hybrid_repeats} - dict_key: ytvis_val - with_seg_masks: true - - - model: - _target_: sam3.model_builder.build_sam3_video_model - bpe_path: ${paths.bpe_path} - has_presence_token: True - geo_encoder_use_img_cross_attn: True - apply_temporal_disambiguation: False - - meters: - val: - ytvis_val: - pred_file: # key - _target_: sam3.eval.ytvis_eval.YTVISResultsWriter - dump_file: ${launcher.experiment_log_dir}/preds/${paths.dump_file_name}.json - postprocessor: ${scratch.vid_mask_postprocessor} - gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} - - optim: - amp: - enabled: True - amp_dtype: bfloat16 - - - checkpoint: - save_dir: ${launcher.experiment_log_dir}/checkpoints - save_freq: 0 # 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: 8 - gpus_per_node: 8 - experiment_log_dir: ${paths.experiment_log_dir} - multiprocessing_context: forkserver - -submitit: - account: null - partition: null - qos: null - timeout_hour: 72 - use_cluster: True - cpus_per_task: 10 - port_range: [10000, 65000] - constraint: null diff --git a/train/configs/saco_video_evals/saco_veval_sav_val.yaml b/train/configs/saco_video_evals/saco_veval_sav_val.yaml deleted file mode 100644 index 25c6e60..0000000 --- a/train/configs/saco_video_evals/saco_veval_sav_val.yaml +++ /dev/null @@ -1,174 +0,0 @@ -# @package _global_ -defaults: - - _self_ - -# ============================================================================ -# Paths Configuration (Chage this to your own paths) -# ============================================================================ -paths: - - dump_file_name: saco_veval_sav_val - experiment_log_dir: - ytvis_json: /saco_veval_sav_val.json - ytvis_dir : - bpe_path: # This should be under assets/bpe_simple_vocab_16e6.txt.gz - num_videos: null - -# ============================================================================ -# Different helper parameters and functions -# ============================================================================ -scratch: - vid_mask_postprocessor: - _target_: sam3.eval.postprocessors.PostProcessNullOp - - use_presence_eval: True - - video_transforms_val: - - _target_: sam3.train.transforms.basic_for_api.ComposeAPI - transforms: - - _target_: sam3.train.transforms.segmentation.DecodeRle - # resize the image to 1024x1024 resolution - - _target_: sam3.train.transforms.basic_for_api.RandomResizeAPI - sizes: ${scratch.resolution} # originally `resolution: 1024` - square: true - consistent_transform: true - - _target_: sam3.train.transforms.basic_for_api.ToTensorAPI - - _target_: sam3.train.transforms.basic_for_api.NormalizeAPI - mean: ${scratch.val_norm_mean} - std: ${scratch.val_norm_std} - - # Model parameters - d_model: 256 - - # Image processing parameters - resolution: 1008 - - # 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] - - val_batch_size: 1 - num_val_workers: 0 - max_data_epochs: 20 - hybrid_repeats: 1 - gather_pred_via_filesys: false - - -# ============================================================================ -# Trainer Configuration -# ============================================================================ - -trainer: - _target_: sam3.train.trainer.Trainer - skip_saving_ckpts: true - empty_gpu_mem_cache_after_eval: True - skip_first_val: True - max_epochs: ${scratch.max_data_epochs} - accelerator: cuda - seed_value: 123 - val_epoch_freq: 10 - mode: val - - distributed: - backend: nccl - find_unused_parameters: True - gradient_as_bucket_view: True - - loss: - all: - _target_: sam3.train.loss.sam3_loss.DummyLoss - default: - _target_: sam3.train.loss.sam3_loss.DummyLoss - - data: - train: null - val: - _target_: sam3.train.data.torch_dataset.TorchDataset - dataset: - _target_: sam3.train.data.sam3_video_dataset.VideoGroundingDataset - limit_ids: ${paths.num_videos} - img_folder: ${paths.ytvis_dir} - ann_file: ${paths.ytvis_json} - coco_json_loader: - _target_: sam3.train.data.coco_json_loaders.SAM3_VEVAL_API_FROM_JSON_NP - _partial_: true - - transforms: ${scratch.video_transforms_val} - max_ann_per_img: 100000 # filtered in transforms - max_val_queries: 100000 - multiplier: 1 - load_segmentation: true - training: false - - - shuffle: False - batch_size: ${scratch.val_batch_size} - num_workers: ${scratch.num_val_workers} - pin_memory: True - drop_last: False - collate_fn: - _target_: sam3.train.data.collator.collate_fn_api - _partial_: true - repeats: ${scratch.hybrid_repeats} - dict_key: ytvis_val - with_seg_masks: true - - - model: - _target_: sam3.model_builder.build_sam3_video_model - bpe_path: ${paths.bpe_path} - has_presence_token: True - geo_encoder_use_img_cross_attn: True - apply_temporal_disambiguation: True - - meters: - val: - ytvis_val: - pred_file: # key - _target_: sam3.eval.ytvis_eval.YTVISResultsWriter - dump_file: ${launcher.experiment_log_dir}/preds/${paths.dump_file_name}.json - postprocessor: ${scratch.vid_mask_postprocessor} - gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} - - optim: - amp: - enabled: True - amp_dtype: bfloat16 - - - checkpoint: - save_dir: ${launcher.experiment_log_dir}/checkpoints - save_freq: 0 # 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: 8 - gpus_per_node: 8 - experiment_log_dir: ${paths.experiment_log_dir} - multiprocessing_context: forkserver - -submitit: - account: null - partition: null - qos: null - timeout_hour: 72 - use_cluster: True - cpus_per_task: 10 - port_range: [10000, 65000] - constraint: null diff --git a/train/configs/saco_video_evals/saco_veval_sav_val_noheur.yaml b/train/configs/saco_video_evals/saco_veval_sav_val_noheur.yaml deleted file mode 100644 index 9a89eba..0000000 --- a/train/configs/saco_video_evals/saco_veval_sav_val_noheur.yaml +++ /dev/null @@ -1,174 +0,0 @@ -# @package _global_ -defaults: - - _self_ - -# ============================================================================ -# Paths Configuration (Chage this to your own paths) -# ============================================================================ -paths: - - dump_file_name: saco_veval_sav_val - experiment_log_dir: - ytvis_json: /saco_veval_sav_val.json - ytvis_dir : - bpe_path: # This should be under assets/bpe_simple_vocab_16e6.txt.gz - num_videos: null - -# ============================================================================ -# Different helper parameters and functions -# ============================================================================ -scratch: - vid_mask_postprocessor: - _target_: sam3.eval.postprocessors.PostProcessNullOp - - use_presence_eval: True - - video_transforms_val: - - _target_: sam3.train.transforms.basic_for_api.ComposeAPI - transforms: - - _target_: sam3.train.transforms.segmentation.DecodeRle - # resize the image to 1024x1024 resolution - - _target_: sam3.train.transforms.basic_for_api.RandomResizeAPI - sizes: ${scratch.resolution} # originally `resolution: 1024` - square: true - consistent_transform: true - - _target_: sam3.train.transforms.basic_for_api.ToTensorAPI - - _target_: sam3.train.transforms.basic_for_api.NormalizeAPI - mean: ${scratch.val_norm_mean} - std: ${scratch.val_norm_std} - - # Model parameters - d_model: 256 - - # Image processing parameters - resolution: 1008 - - # 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] - - val_batch_size: 1 - num_val_workers: 0 - max_data_epochs: 20 - hybrid_repeats: 1 - gather_pred_via_filesys: false - - -# ============================================================================ -# Trainer Configuration -# ============================================================================ - -trainer: - _target_: sam3.train.trainer.Trainer - skip_saving_ckpts: true - empty_gpu_mem_cache_after_eval: True - skip_first_val: True - max_epochs: ${scratch.max_data_epochs} - accelerator: cuda - seed_value: 123 - val_epoch_freq: 10 - mode: val - - distributed: - backend: nccl - find_unused_parameters: True - gradient_as_bucket_view: True - - loss: - all: - _target_: sam3.train.loss.sam3_loss.DummyLoss - default: - _target_: sam3.train.loss.sam3_loss.DummyLoss - - data: - train: null - val: - _target_: sam3.train.data.torch_dataset.TorchDataset - dataset: - _target_: sam3.train.data.sam3_video_dataset.VideoGroundingDataset - limit_ids: ${paths.num_videos} - img_folder: ${paths.ytvis_dir} - ann_file: ${paths.ytvis_json} - coco_json_loader: - _target_: sam3.train.data.coco_json_loaders.SAM3_VEVAL_API_FROM_JSON_NP - _partial_: true - - transforms: ${scratch.video_transforms_val} - max_ann_per_img: 100000 # filtered in transforms - max_val_queries: 100000 - multiplier: 1 - load_segmentation: true - training: false - - - shuffle: False - batch_size: ${scratch.val_batch_size} - num_workers: ${scratch.num_val_workers} - pin_memory: True - drop_last: False - collate_fn: - _target_: sam3.train.data.collator.collate_fn_api - _partial_: true - repeats: ${scratch.hybrid_repeats} - dict_key: ytvis_val - with_seg_masks: true - - - model: - _target_: sam3.model_builder.build_sam3_video_model - bpe_path: ${paths.bpe_path} - has_presence_token: True - geo_encoder_use_img_cross_attn: True - apply_temporal_disambiguation: False - - meters: - val: - ytvis_val: - pred_file: # key - _target_: sam3.eval.ytvis_eval.YTVISResultsWriter - dump_file: ${launcher.experiment_log_dir}/preds/${paths.dump_file_name}.json - postprocessor: ${scratch.vid_mask_postprocessor} - gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} - - optim: - amp: - enabled: True - amp_dtype: bfloat16 - - - checkpoint: - save_dir: ${launcher.experiment_log_dir}/checkpoints - save_freq: 0 # 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: 8 - gpus_per_node: 8 - experiment_log_dir: ${paths.experiment_log_dir} - multiprocessing_context: forkserver - -submitit: - account: null - partition: null - qos: null - timeout_hour: 72 - use_cluster: True - cpus_per_task: 10 - port_range: [10000, 65000] - constraint: null diff --git a/train/configs/saco_video_evals/saco_veval_smartglasses_test.yaml b/train/configs/saco_video_evals/saco_veval_smartglasses_test.yaml deleted file mode 100644 index fdc8185..0000000 --- a/train/configs/saco_video_evals/saco_veval_smartglasses_test.yaml +++ /dev/null @@ -1,174 +0,0 @@ -# @package _global_ -defaults: - - _self_ - -# ============================================================================ -# Paths Configuration (Chage this to your own paths) -# ============================================================================ -paths: - - dump_file_name: saco_veval_smartglasses_test - experiment_log_dir: - ytvis_json: /saco_veval_smartglasses_test.json - ytvis_dir : - bpe_path: # This should be under assets/bpe_simple_vocab_16e6.txt.gz - num_videos: null - -# ============================================================================ -# Different helper parameters and functions -# ============================================================================ -scratch: - vid_mask_postprocessor: - _target_: sam3.eval.postprocessors.PostProcessNullOp - - use_presence_eval: True - - video_transforms_val: - - _target_: sam3.train.transforms.basic_for_api.ComposeAPI - transforms: - - _target_: sam3.train.transforms.segmentation.DecodeRle - # resize the image to 1024x1024 resolution - - _target_: sam3.train.transforms.basic_for_api.RandomResizeAPI - sizes: ${scratch.resolution} # originally `resolution: 1024` - square: true - consistent_transform: true - - _target_: sam3.train.transforms.basic_for_api.ToTensorAPI - - _target_: sam3.train.transforms.basic_for_api.NormalizeAPI - mean: ${scratch.val_norm_mean} - std: ${scratch.val_norm_std} - - # Model parameters - d_model: 256 - - # Image processing parameters - resolution: 1008 - - # 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] - - val_batch_size: 1 - num_val_workers: 0 - max_data_epochs: 20 - hybrid_repeats: 1 - gather_pred_via_filesys: false - - -# ============================================================================ -# Trainer Configuration -# ============================================================================ - -trainer: - _target_: sam3.train.trainer.Trainer - skip_saving_ckpts: true - empty_gpu_mem_cache_after_eval: True - skip_first_val: True - max_epochs: ${scratch.max_data_epochs} - accelerator: cuda - seed_value: 123 - val_epoch_freq: 10 - mode: val - - distributed: - backend: nccl - find_unused_parameters: True - gradient_as_bucket_view: True - - loss: - all: - _target_: sam3.train.loss.sam3_loss.DummyLoss - default: - _target_: sam3.train.loss.sam3_loss.DummyLoss - - data: - train: null - val: - _target_: sam3.train.data.torch_dataset.TorchDataset - dataset: - _target_: sam3.train.data.sam3_video_dataset.VideoGroundingDataset - limit_ids: ${paths.num_videos} - img_folder: ${paths.ytvis_dir} - ann_file: ${paths.ytvis_json} - coco_json_loader: - _target_: sam3.train.data.coco_json_loaders.SAM3_VEVAL_API_FROM_JSON_NP - _partial_: true - - transforms: ${scratch.video_transforms_val} - max_ann_per_img: 100000 # filtered in transforms - max_val_queries: 100000 - multiplier: 1 - load_segmentation: true - training: false - - - shuffle: False - batch_size: ${scratch.val_batch_size} - num_workers: ${scratch.num_val_workers} - pin_memory: True - drop_last: False - collate_fn: - _target_: sam3.train.data.collator.collate_fn_api - _partial_: true - repeats: ${scratch.hybrid_repeats} - dict_key: ytvis_val - with_seg_masks: true - - - model: - _target_: sam3.model_builder.build_sam3_video_model - bpe_path: ${paths.bpe_path} - has_presence_token: True - geo_encoder_use_img_cross_attn: True - apply_temporal_disambiguation: True - - meters: - val: - ytvis_val: - pred_file: # key - _target_: sam3.eval.ytvis_eval.YTVISResultsWriter - dump_file: ${launcher.experiment_log_dir}/preds/${paths.dump_file_name}.json - postprocessor: ${scratch.vid_mask_postprocessor} - gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} - - optim: - amp: - enabled: True - amp_dtype: bfloat16 - - - checkpoint: - save_dir: ${launcher.experiment_log_dir}/checkpoints - save_freq: 0 # 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: 8 - gpus_per_node: 8 - experiment_log_dir: ${paths.experiment_log_dir} - multiprocessing_context: forkserver - -submitit: - account: null - partition: null - qos: null - timeout_hour: 72 - use_cluster: True - cpus_per_task: 10 - port_range: [10000, 65000] - constraint: null diff --git a/train/configs/saco_video_evals/saco_veval_smartglasses_test_noheur.yaml b/train/configs/saco_video_evals/saco_veval_smartglasses_test_noheur.yaml deleted file mode 100644 index 2d6150e..0000000 --- a/train/configs/saco_video_evals/saco_veval_smartglasses_test_noheur.yaml +++ /dev/null @@ -1,174 +0,0 @@ -# @package _global_ -defaults: - - _self_ - -# ============================================================================ -# Paths Configuration (Chage this to your own paths) -# ============================================================================ -paths: - - dump_file_name: saco_veval_smartglasses_test - experiment_log_dir: - ytvis_json: /saco_veval_smartglasses_test.json - ytvis_dir : - bpe_path: # This should be under assets/bpe_simple_vocab_16e6.txt.gz - num_videos: null - -# ============================================================================ -# Different helper parameters and functions -# ============================================================================ -scratch: - vid_mask_postprocessor: - _target_: sam3.eval.postprocessors.PostProcessNullOp - - use_presence_eval: True - - video_transforms_val: - - _target_: sam3.train.transforms.basic_for_api.ComposeAPI - transforms: - - _target_: sam3.train.transforms.segmentation.DecodeRle - # resize the image to 1024x1024 resolution - - _target_: sam3.train.transforms.basic_for_api.RandomResizeAPI - sizes: ${scratch.resolution} # originally `resolution: 1024` - square: true - consistent_transform: true - - _target_: sam3.train.transforms.basic_for_api.ToTensorAPI - - _target_: sam3.train.transforms.basic_for_api.NormalizeAPI - mean: ${scratch.val_norm_mean} - std: ${scratch.val_norm_std} - - # Model parameters - d_model: 256 - - # Image processing parameters - resolution: 1008 - - # 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] - - val_batch_size: 1 - num_val_workers: 0 - max_data_epochs: 20 - hybrid_repeats: 1 - gather_pred_via_filesys: false - - -# ============================================================================ -# Trainer Configuration -# ============================================================================ - -trainer: - _target_: sam3.train.trainer.Trainer - skip_saving_ckpts: true - empty_gpu_mem_cache_after_eval: True - skip_first_val: True - max_epochs: ${scratch.max_data_epochs} - accelerator: cuda - seed_value: 123 - val_epoch_freq: 10 - mode: val - - distributed: - backend: nccl - find_unused_parameters: True - gradient_as_bucket_view: True - - loss: - all: - _target_: sam3.train.loss.sam3_loss.DummyLoss - default: - _target_: sam3.train.loss.sam3_loss.DummyLoss - - data: - train: null - val: - _target_: sam3.train.data.torch_dataset.TorchDataset - dataset: - _target_: sam3.train.data.sam3_video_dataset.VideoGroundingDataset - limit_ids: ${paths.num_videos} - img_folder: ${paths.ytvis_dir} - ann_file: ${paths.ytvis_json} - coco_json_loader: - _target_: sam3.train.data.coco_json_loaders.SAM3_VEVAL_API_FROM_JSON_NP - _partial_: true - - transforms: ${scratch.video_transforms_val} - max_ann_per_img: 100000 # filtered in transforms - max_val_queries: 100000 - multiplier: 1 - load_segmentation: true - training: false - - - shuffle: False - batch_size: ${scratch.val_batch_size} - num_workers: ${scratch.num_val_workers} - pin_memory: True - drop_last: False - collate_fn: - _target_: sam3.train.data.collator.collate_fn_api - _partial_: true - repeats: ${scratch.hybrid_repeats} - dict_key: ytvis_val - with_seg_masks: true - - - model: - _target_: sam3.model_builder.build_sam3_video_model - bpe_path: ${paths.bpe_path} - has_presence_token: True - geo_encoder_use_img_cross_attn: True - apply_temporal_disambiguation: False - - meters: - val: - ytvis_val: - pred_file: # key - _target_: sam3.eval.ytvis_eval.YTVISResultsWriter - dump_file: ${launcher.experiment_log_dir}/preds/${paths.dump_file_name}.json - postprocessor: ${scratch.vid_mask_postprocessor} - gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} - - optim: - amp: - enabled: True - amp_dtype: bfloat16 - - - checkpoint: - save_dir: ${launcher.experiment_log_dir}/checkpoints - save_freq: 0 # 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: 8 - gpus_per_node: 8 - experiment_log_dir: ${paths.experiment_log_dir} - multiprocessing_context: forkserver - -submitit: - account: null - partition: null - qos: null - timeout_hour: 72 - use_cluster: True - cpus_per_task: 10 - port_range: [10000, 65000] - constraint: null diff --git a/train/configs/saco_video_evals/saco_veval_smartglasses_val.yaml b/train/configs/saco_video_evals/saco_veval_smartglasses_val.yaml deleted file mode 100644 index a9bffda..0000000 --- a/train/configs/saco_video_evals/saco_veval_smartglasses_val.yaml +++ /dev/null @@ -1,174 +0,0 @@ -# @package _global_ -defaults: - - _self_ - -# ============================================================================ -# Paths Configuration (Chage this to your own paths) -# ============================================================================ -paths: - - dump_file_name: saco_veval_smartglasses_val - experiment_log_dir: - ytvis_json: /saco_veval_smartglasses_val.json - ytvis_dir : - bpe_path: # This should be under assets/bpe_simple_vocab_16e6.txt.gz - num_videos: null - -# ============================================================================ -# Different helper parameters and functions -# ============================================================================ -scratch: - vid_mask_postprocessor: - _target_: sam3.eval.postprocessors.PostProcessNullOp - - use_presence_eval: True - - video_transforms_val: - - _target_: sam3.train.transforms.basic_for_api.ComposeAPI - transforms: - - _target_: sam3.train.transforms.segmentation.DecodeRle - # resize the image to 1024x1024 resolution - - _target_: sam3.train.transforms.basic_for_api.RandomResizeAPI - sizes: ${scratch.resolution} # originally `resolution: 1024` - square: true - consistent_transform: true - - _target_: sam3.train.transforms.basic_for_api.ToTensorAPI - - _target_: sam3.train.transforms.basic_for_api.NormalizeAPI - mean: ${scratch.val_norm_mean} - std: ${scratch.val_norm_std} - - # Model parameters - d_model: 256 - - # Image processing parameters - resolution: 1008 - - # 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] - - val_batch_size: 1 - num_val_workers: 0 - max_data_epochs: 20 - hybrid_repeats: 1 - gather_pred_via_filesys: false - - -# ============================================================================ -# Trainer Configuration -# ============================================================================ - -trainer: - _target_: sam3.train.trainer.Trainer - skip_saving_ckpts: true - empty_gpu_mem_cache_after_eval: True - skip_first_val: True - max_epochs: ${scratch.max_data_epochs} - accelerator: cuda - seed_value: 123 - val_epoch_freq: 10 - mode: val - - distributed: - backend: nccl - find_unused_parameters: True - gradient_as_bucket_view: True - - loss: - all: - _target_: sam3.train.loss.sam3_loss.DummyLoss - default: - _target_: sam3.train.loss.sam3_loss.DummyLoss - - data: - train: null - val: - _target_: sam3.train.data.torch_dataset.TorchDataset - dataset: - _target_: sam3.train.data.sam3_video_dataset.VideoGroundingDataset - limit_ids: ${paths.num_videos} - img_folder: ${paths.ytvis_dir} - ann_file: ${paths.ytvis_json} - coco_json_loader: - _target_: sam3.train.data.coco_json_loaders.SAM3_VEVAL_API_FROM_JSON_NP - _partial_: true - - transforms: ${scratch.video_transforms_val} - max_ann_per_img: 100000 # filtered in transforms - max_val_queries: 100000 - multiplier: 1 - load_segmentation: true - training: false - - - shuffle: False - batch_size: ${scratch.val_batch_size} - num_workers: ${scratch.num_val_workers} - pin_memory: True - drop_last: False - collate_fn: - _target_: sam3.train.data.collator.collate_fn_api - _partial_: true - repeats: ${scratch.hybrid_repeats} - dict_key: ytvis_val - with_seg_masks: true - - - model: - _target_: sam3.model_builder.build_sam3_video_model - bpe_path: ${paths.bpe_path} - has_presence_token: True - geo_encoder_use_img_cross_attn: True - apply_temporal_disambiguation: True - - meters: - val: - ytvis_val: - pred_file: # key - _target_: sam3.eval.ytvis_eval.YTVISResultsWriter - dump_file: ${launcher.experiment_log_dir}/preds/${paths.dump_file_name}.json - postprocessor: ${scratch.vid_mask_postprocessor} - gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} - - optim: - amp: - enabled: True - amp_dtype: bfloat16 - - - checkpoint: - save_dir: ${launcher.experiment_log_dir}/checkpoints - save_freq: 0 # 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: 8 - gpus_per_node: 8 - experiment_log_dir: ${paths.experiment_log_dir} - multiprocessing_context: forkserver - -submitit: - account: null - partition: null - qos: null - timeout_hour: 72 - use_cluster: True - cpus_per_task: 10 - port_range: [10000, 65000] - constraint: null diff --git a/train/configs/saco_video_evals/saco_veval_smartglasses_val_noheur.yaml b/train/configs/saco_video_evals/saco_veval_smartglasses_val_noheur.yaml deleted file mode 100644 index e1f6443..0000000 --- a/train/configs/saco_video_evals/saco_veval_smartglasses_val_noheur.yaml +++ /dev/null @@ -1,174 +0,0 @@ -# @package _global_ -defaults: - - _self_ - -# ============================================================================ -# Paths Configuration (Chage this to your own paths) -# ============================================================================ -paths: - - dump_file_name: saco_veval_smartglasses_val - experiment_log_dir: - ytvis_json: /saco_veval_smartglasses_val.json - ytvis_dir : - bpe_path: # This should be under assets/bpe_simple_vocab_16e6.txt.gz - num_videos: null - -# ============================================================================ -# Different helper parameters and functions -# ============================================================================ -scratch: - vid_mask_postprocessor: - _target_: sam3.eval.postprocessors.PostProcessNullOp - - use_presence_eval: True - - video_transforms_val: - - _target_: sam3.train.transforms.basic_for_api.ComposeAPI - transforms: - - _target_: sam3.train.transforms.segmentation.DecodeRle - # resize the image to 1024x1024 resolution - - _target_: sam3.train.transforms.basic_for_api.RandomResizeAPI - sizes: ${scratch.resolution} # originally `resolution: 1024` - square: true - consistent_transform: true - - _target_: sam3.train.transforms.basic_for_api.ToTensorAPI - - _target_: sam3.train.transforms.basic_for_api.NormalizeAPI - mean: ${scratch.val_norm_mean} - std: ${scratch.val_norm_std} - - # Model parameters - d_model: 256 - - # Image processing parameters - resolution: 1008 - - # 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] - - val_batch_size: 1 - num_val_workers: 0 - max_data_epochs: 20 - hybrid_repeats: 1 - gather_pred_via_filesys: false - - -# ============================================================================ -# Trainer Configuration -# ============================================================================ - -trainer: - _target_: sam3.train.trainer.Trainer - skip_saving_ckpts: true - empty_gpu_mem_cache_after_eval: True - skip_first_val: True - max_epochs: ${scratch.max_data_epochs} - accelerator: cuda - seed_value: 123 - val_epoch_freq: 10 - mode: val - - distributed: - backend: nccl - find_unused_parameters: True - gradient_as_bucket_view: True - - loss: - all: - _target_: sam3.train.loss.sam3_loss.DummyLoss - default: - _target_: sam3.train.loss.sam3_loss.DummyLoss - - data: - train: null - val: - _target_: sam3.train.data.torch_dataset.TorchDataset - dataset: - _target_: sam3.train.data.sam3_video_dataset.VideoGroundingDataset - limit_ids: ${paths.num_videos} - img_folder: ${paths.ytvis_dir} - ann_file: ${paths.ytvis_json} - coco_json_loader: - _target_: sam3.train.data.coco_json_loaders.SAM3_VEVAL_API_FROM_JSON_NP - _partial_: true - - transforms: ${scratch.video_transforms_val} - max_ann_per_img: 100000 # filtered in transforms - max_val_queries: 100000 - multiplier: 1 - load_segmentation: true - training: false - - - shuffle: False - batch_size: ${scratch.val_batch_size} - num_workers: ${scratch.num_val_workers} - pin_memory: True - drop_last: False - collate_fn: - _target_: sam3.train.data.collator.collate_fn_api - _partial_: true - repeats: ${scratch.hybrid_repeats} - dict_key: ytvis_val - with_seg_masks: true - - - model: - _target_: sam3.model_builder.build_sam3_video_model - bpe_path: ${paths.bpe_path} - has_presence_token: True - geo_encoder_use_img_cross_attn: True - apply_temporal_disambiguation: False - - meters: - val: - ytvis_val: - pred_file: # key - _target_: sam3.eval.ytvis_eval.YTVISResultsWriter - dump_file: ${launcher.experiment_log_dir}/preds/${paths.dump_file_name}.json - postprocessor: ${scratch.vid_mask_postprocessor} - gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} - - optim: - amp: - enabled: True - amp_dtype: bfloat16 - - - checkpoint: - save_dir: ${launcher.experiment_log_dir}/checkpoints - save_freq: 0 # 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: 8 - gpus_per_node: 8 - experiment_log_dir: ${paths.experiment_log_dir} - multiprocessing_context: forkserver - -submitit: - account: null - partition: null - qos: null - timeout_hour: 72 - use_cluster: True - cpus_per_task: 10 - port_range: [10000, 65000] - constraint: null diff --git a/train/configs/saco_video_evals/saco_veval_yt1b_test.yaml b/train/configs/saco_video_evals/saco_veval_yt1b_test.yaml deleted file mode 100644 index 71e5034..0000000 --- a/train/configs/saco_video_evals/saco_veval_yt1b_test.yaml +++ /dev/null @@ -1,174 +0,0 @@ -# @package _global_ -defaults: - - _self_ - -# ============================================================================ -# Paths Configuration (Chage this to your own paths) -# ============================================================================ -paths: - - dump_file_name: saco_veval_yt1b_test - experiment_log_dir: - ytvis_json: /saco_veval_yt1b_test.json - ytvis_dir : - bpe_path: # This should be under assets/bpe_simple_vocab_16e6.txt.gz - num_videos: null - -# ============================================================================ -# Different helper parameters and functions -# ============================================================================ -scratch: - vid_mask_postprocessor: - _target_: sam3.eval.postprocessors.PostProcessNullOp - - use_presence_eval: True - - video_transforms_val: - - _target_: sam3.train.transforms.basic_for_api.ComposeAPI - transforms: - - _target_: sam3.train.transforms.segmentation.DecodeRle - # resize the image to 1024x1024 resolution - - _target_: sam3.train.transforms.basic_for_api.RandomResizeAPI - sizes: ${scratch.resolution} # originally `resolution: 1024` - square: true - consistent_transform: true - - _target_: sam3.train.transforms.basic_for_api.ToTensorAPI - - _target_: sam3.train.transforms.basic_for_api.NormalizeAPI - mean: ${scratch.val_norm_mean} - std: ${scratch.val_norm_std} - - # Model parameters - d_model: 256 - - # Image processing parameters - resolution: 1008 - - # 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] - - val_batch_size: 1 - num_val_workers: 0 - max_data_epochs: 20 - hybrid_repeats: 1 - gather_pred_via_filesys: false - - -# ============================================================================ -# Trainer Configuration -# ============================================================================ - -trainer: - _target_: sam3.train.trainer.Trainer - skip_saving_ckpts: true - empty_gpu_mem_cache_after_eval: True - skip_first_val: True - max_epochs: ${scratch.max_data_epochs} - accelerator: cuda - seed_value: 123 - val_epoch_freq: 10 - mode: val - - distributed: - backend: nccl - find_unused_parameters: True - gradient_as_bucket_view: True - - loss: - all: - _target_: sam3.train.loss.sam3_loss.DummyLoss - default: - _target_: sam3.train.loss.sam3_loss.DummyLoss - - data: - train: null - val: - _target_: sam3.train.data.torch_dataset.TorchDataset - dataset: - _target_: sam3.train.data.sam3_video_dataset.VideoGroundingDataset - limit_ids: ${paths.num_videos} - img_folder: ${paths.ytvis_dir} - ann_file: ${paths.ytvis_json} - coco_json_loader: - _target_: sam3.train.data.coco_json_loaders.SAM3_VEVAL_API_FROM_JSON_NP - _partial_: true - - transforms: ${scratch.video_transforms_val} - max_ann_per_img: 100000 # filtered in transforms - max_val_queries: 100000 - multiplier: 1 - load_segmentation: true - training: false - - - shuffle: False - batch_size: ${scratch.val_batch_size} - num_workers: ${scratch.num_val_workers} - pin_memory: True - drop_last: False - collate_fn: - _target_: sam3.train.data.collator.collate_fn_api - _partial_: true - repeats: ${scratch.hybrid_repeats} - dict_key: ytvis_val - with_seg_masks: true - - - model: - _target_: sam3.model_builder.build_sam3_video_model - bpe_path: ${paths.bpe_path} - has_presence_token: True - geo_encoder_use_img_cross_attn: True - apply_temporal_disambiguation: True - - meters: - val: - ytvis_val: - pred_file: # key - _target_: sam3.eval.ytvis_eval.YTVISResultsWriter - dump_file: ${launcher.experiment_log_dir}/preds/${paths.dump_file_name}.json - postprocessor: ${scratch.vid_mask_postprocessor} - gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} - - optim: - amp: - enabled: True - amp_dtype: bfloat16 - - - checkpoint: - save_dir: ${launcher.experiment_log_dir}/checkpoints - save_freq: 0 # 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: 8 - gpus_per_node: 8 - experiment_log_dir: ${paths.experiment_log_dir} - multiprocessing_context: forkserver - -submitit: - account: null - partition: null - qos: null - timeout_hour: 72 - use_cluster: True - cpus_per_task: 10 - port_range: [10000, 65000] - constraint: null diff --git a/train/configs/saco_video_evals/saco_veval_yt1b_test_noheur.yaml b/train/configs/saco_video_evals/saco_veval_yt1b_test_noheur.yaml deleted file mode 100644 index f8df6ae..0000000 --- a/train/configs/saco_video_evals/saco_veval_yt1b_test_noheur.yaml +++ /dev/null @@ -1,174 +0,0 @@ -# @package _global_ -defaults: - - _self_ - -# ============================================================================ -# Paths Configuration (Chage this to your own paths) -# ============================================================================ -paths: - - dump_file_name: saco_veval_yt1b_test - experiment_log_dir: - ytvis_json: /saco_veval_yt1b_test.json - ytvis_dir : - bpe_path: # This should be under assets/bpe_simple_vocab_16e6.txt.gz - num_videos: null - -# ============================================================================ -# Different helper parameters and functions -# ============================================================================ -scratch: - vid_mask_postprocessor: - _target_: sam3.eval.postprocessors.PostProcessNullOp - - use_presence_eval: True - - video_transforms_val: - - _target_: sam3.train.transforms.basic_for_api.ComposeAPI - transforms: - - _target_: sam3.train.transforms.segmentation.DecodeRle - # resize the image to 1024x1024 resolution - - _target_: sam3.train.transforms.basic_for_api.RandomResizeAPI - sizes: ${scratch.resolution} # originally `resolution: 1024` - square: true - consistent_transform: true - - _target_: sam3.train.transforms.basic_for_api.ToTensorAPI - - _target_: sam3.train.transforms.basic_for_api.NormalizeAPI - mean: ${scratch.val_norm_mean} - std: ${scratch.val_norm_std} - - # Model parameters - d_model: 256 - - # Image processing parameters - resolution: 1008 - - # 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] - - val_batch_size: 1 - num_val_workers: 0 - max_data_epochs: 20 - hybrid_repeats: 1 - gather_pred_via_filesys: false - - -# ============================================================================ -# Trainer Configuration -# ============================================================================ - -trainer: - _target_: sam3.train.trainer.Trainer - skip_saving_ckpts: true - empty_gpu_mem_cache_after_eval: True - skip_first_val: True - max_epochs: ${scratch.max_data_epochs} - accelerator: cuda - seed_value: 123 - val_epoch_freq: 10 - mode: val - - distributed: - backend: nccl - find_unused_parameters: True - gradient_as_bucket_view: True - - loss: - all: - _target_: sam3.train.loss.sam3_loss.DummyLoss - default: - _target_: sam3.train.loss.sam3_loss.DummyLoss - - data: - train: null - val: - _target_: sam3.train.data.torch_dataset.TorchDataset - dataset: - _target_: sam3.train.data.sam3_video_dataset.VideoGroundingDataset - limit_ids: ${paths.num_videos} - img_folder: ${paths.ytvis_dir} - ann_file: ${paths.ytvis_json} - coco_json_loader: - _target_: sam3.train.data.coco_json_loaders.SAM3_VEVAL_API_FROM_JSON_NP - _partial_: true - - transforms: ${scratch.video_transforms_val} - max_ann_per_img: 100000 # filtered in transforms - max_val_queries: 100000 - multiplier: 1 - load_segmentation: true - training: false - - - shuffle: False - batch_size: ${scratch.val_batch_size} - num_workers: ${scratch.num_val_workers} - pin_memory: True - drop_last: False - collate_fn: - _target_: sam3.train.data.collator.collate_fn_api - _partial_: true - repeats: ${scratch.hybrid_repeats} - dict_key: ytvis_val - with_seg_masks: true - - - model: - _target_: sam3.model_builder.build_sam3_video_model - bpe_path: ${paths.bpe_path} - has_presence_token: True - geo_encoder_use_img_cross_attn: True - apply_temporal_disambiguation: False - - meters: - val: - ytvis_val: - pred_file: # key - _target_: sam3.eval.ytvis_eval.YTVISResultsWriter - dump_file: ${launcher.experiment_log_dir}/preds/${paths.dump_file_name}.json - postprocessor: ${scratch.vid_mask_postprocessor} - gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} - - optim: - amp: - enabled: True - amp_dtype: bfloat16 - - - checkpoint: - save_dir: ${launcher.experiment_log_dir}/checkpoints - save_freq: 0 # 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: 8 - gpus_per_node: 8 - experiment_log_dir: ${paths.experiment_log_dir} - multiprocessing_context: forkserver - -submitit: - account: null - partition: null - qos: null - timeout_hour: 72 - use_cluster: True - cpus_per_task: 10 - port_range: [10000, 65000] - constraint: null diff --git a/train/configs/saco_video_evals/saco_veval_yt1b_val.yaml b/train/configs/saco_video_evals/saco_veval_yt1b_val.yaml deleted file mode 100644 index 5816952..0000000 --- a/train/configs/saco_video_evals/saco_veval_yt1b_val.yaml +++ /dev/null @@ -1,174 +0,0 @@ -# @package _global_ -defaults: - - _self_ - -# ============================================================================ -# Paths Configuration (Chage this to your own paths) -# ============================================================================ -paths: - - dump_file_name: saco_veval_yt1b_val - experiment_log_dir: - ytvis_json: /saco_veval_yt1b_val.json - ytvis_dir : - bpe_path: # This should be under assets/bpe_simple_vocab_16e6.txt.gz - num_videos: null - -# ============================================================================ -# Different helper parameters and functions -# ============================================================================ -scratch: - vid_mask_postprocessor: - _target_: sam3.eval.postprocessors.PostProcessNullOp - - use_presence_eval: True - - video_transforms_val: - - _target_: sam3.train.transforms.basic_for_api.ComposeAPI - transforms: - - _target_: sam3.train.transforms.segmentation.DecodeRle - # resize the image to 1024x1024 resolution - - _target_: sam3.train.transforms.basic_for_api.RandomResizeAPI - sizes: ${scratch.resolution} # originally `resolution: 1024` - square: true - consistent_transform: true - - _target_: sam3.train.transforms.basic_for_api.ToTensorAPI - - _target_: sam3.train.transforms.basic_for_api.NormalizeAPI - mean: ${scratch.val_norm_mean} - std: ${scratch.val_norm_std} - - # Model parameters - d_model: 256 - - # Image processing parameters - resolution: 1008 - - # 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] - - val_batch_size: 1 - num_val_workers: 0 - max_data_epochs: 20 - hybrid_repeats: 1 - gather_pred_via_filesys: false - - -# ============================================================================ -# Trainer Configuration -# ============================================================================ - -trainer: - _target_: sam3.train.trainer.Trainer - skip_saving_ckpts: true - empty_gpu_mem_cache_after_eval: True - skip_first_val: True - max_epochs: ${scratch.max_data_epochs} - accelerator: cuda - seed_value: 123 - val_epoch_freq: 10 - mode: val - - distributed: - backend: nccl - find_unused_parameters: True - gradient_as_bucket_view: True - - loss: - all: - _target_: sam3.train.loss.sam3_loss.DummyLoss - default: - _target_: sam3.train.loss.sam3_loss.DummyLoss - - data: - train: null - val: - _target_: sam3.train.data.torch_dataset.TorchDataset - dataset: - _target_: sam3.train.data.sam3_video_dataset.VideoGroundingDataset - limit_ids: ${paths.num_videos} - img_folder: ${paths.ytvis_dir} - ann_file: ${paths.ytvis_json} - coco_json_loader: - _target_: sam3.train.data.coco_json_loaders.SAM3_VEVAL_API_FROM_JSON_NP - _partial_: true - - transforms: ${scratch.video_transforms_val} - max_ann_per_img: 100000 # filtered in transforms - max_val_queries: 100000 - multiplier: 1 - load_segmentation: true - training: false - - - shuffle: False - batch_size: ${scratch.val_batch_size} - num_workers: ${scratch.num_val_workers} - pin_memory: True - drop_last: False - collate_fn: - _target_: sam3.train.data.collator.collate_fn_api - _partial_: true - repeats: ${scratch.hybrid_repeats} - dict_key: ytvis_val - with_seg_masks: true - - - model: - _target_: sam3.model_builder.build_sam3_video_model - bpe_path: ${paths.bpe_path} - has_presence_token: True - geo_encoder_use_img_cross_attn: True - apply_temporal_disambiguation: True - - meters: - val: - ytvis_val: - pred_file: # key - _target_: sam3.eval.ytvis_eval.YTVISResultsWriter - dump_file: ${launcher.experiment_log_dir}/preds/${paths.dump_file_name}.json - postprocessor: ${scratch.vid_mask_postprocessor} - gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} - - optim: - amp: - enabled: True - amp_dtype: bfloat16 - - - checkpoint: - save_dir: ${launcher.experiment_log_dir}/checkpoints - save_freq: 0 # 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: 8 - gpus_per_node: 8 - experiment_log_dir: ${paths.experiment_log_dir} - multiprocessing_context: forkserver - -submitit: - account: null - partition: null - qos: null - timeout_hour: 72 - use_cluster: True - cpus_per_task: 10 - port_range: [10000, 65000] - constraint: null diff --git a/train/configs/saco_video_evals/saco_veval_yt1b_val_noheur.yaml b/train/configs/saco_video_evals/saco_veval_yt1b_val_noheur.yaml deleted file mode 100644 index 374e5e2..0000000 --- a/train/configs/saco_video_evals/saco_veval_yt1b_val_noheur.yaml +++ /dev/null @@ -1,174 +0,0 @@ -# @package _global_ -defaults: - - _self_ - -# ============================================================================ -# Paths Configuration (Chage this to your own paths) -# ============================================================================ -paths: - - dump_file_name: saco_veval_yt1b_val - experiment_log_dir: - ytvis_json: /saco_veval_yt1b_val.json - ytvis_dir : - bpe_path: # This should be under assets/bpe_simple_vocab_16e6.txt.gz - num_videos: null - -# ============================================================================ -# Different helper parameters and functions -# ============================================================================ -scratch: - vid_mask_postprocessor: - _target_: sam3.eval.postprocessors.PostProcessNullOp - - use_presence_eval: True - - video_transforms_val: - - _target_: sam3.train.transforms.basic_for_api.ComposeAPI - transforms: - - _target_: sam3.train.transforms.segmentation.DecodeRle - # resize the image to 1024x1024 resolution - - _target_: sam3.train.transforms.basic_for_api.RandomResizeAPI - sizes: ${scratch.resolution} # originally `resolution: 1024` - square: true - consistent_transform: true - - _target_: sam3.train.transforms.basic_for_api.ToTensorAPI - - _target_: sam3.train.transforms.basic_for_api.NormalizeAPI - mean: ${scratch.val_norm_mean} - std: ${scratch.val_norm_std} - - # Model parameters - d_model: 256 - - # Image processing parameters - resolution: 1008 - - # 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] - - val_batch_size: 1 - num_val_workers: 0 - max_data_epochs: 20 - hybrid_repeats: 1 - gather_pred_via_filesys: false - - -# ============================================================================ -# Trainer Configuration -# ============================================================================ - -trainer: - _target_: sam3.train.trainer.Trainer - skip_saving_ckpts: true - empty_gpu_mem_cache_after_eval: True - skip_first_val: True - max_epochs: ${scratch.max_data_epochs} - accelerator: cuda - seed_value: 123 - val_epoch_freq: 10 - mode: val - - distributed: - backend: nccl - find_unused_parameters: True - gradient_as_bucket_view: True - - loss: - all: - _target_: sam3.train.loss.sam3_loss.DummyLoss - default: - _target_: sam3.train.loss.sam3_loss.DummyLoss - - data: - train: null - val: - _target_: sam3.train.data.torch_dataset.TorchDataset - dataset: - _target_: sam3.train.data.sam3_video_dataset.VideoGroundingDataset - limit_ids: ${paths.num_videos} - img_folder: ${paths.ytvis_dir} - ann_file: ${paths.ytvis_json} - coco_json_loader: - _target_: sam3.train.data.coco_json_loaders.SAM3_VEVAL_API_FROM_JSON_NP - _partial_: true - - transforms: ${scratch.video_transforms_val} - max_ann_per_img: 100000 # filtered in transforms - max_val_queries: 100000 - multiplier: 1 - load_segmentation: true - training: false - - - shuffle: False - batch_size: ${scratch.val_batch_size} - num_workers: ${scratch.num_val_workers} - pin_memory: True - drop_last: False - collate_fn: - _target_: sam3.train.data.collator.collate_fn_api - _partial_: true - repeats: ${scratch.hybrid_repeats} - dict_key: ytvis_val - with_seg_masks: true - - - model: - _target_: sam3.model_builder.build_sam3_video_model - bpe_path: ${paths.bpe_path} - has_presence_token: True - geo_encoder_use_img_cross_attn: True - apply_temporal_disambiguation: False - - meters: - val: - ytvis_val: - pred_file: # key - _target_: sam3.eval.ytvis_eval.YTVISResultsWriter - dump_file: ${launcher.experiment_log_dir}/preds/${paths.dump_file_name}.json - postprocessor: ${scratch.vid_mask_postprocessor} - gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} - - optim: - amp: - enabled: True - amp_dtype: bfloat16 - - - checkpoint: - save_dir: ${launcher.experiment_log_dir}/checkpoints - save_freq: 0 # 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: 8 - gpus_per_node: 8 - experiment_log_dir: ${paths.experiment_log_dir} - multiprocessing_context: forkserver - -submitit: - account: null - partition: null - qos: null - timeout_hour: 72 - use_cluster: True - cpus_per_task: 10 - port_range: [10000, 65000] - constraint: null diff --git a/train/configs/silver_image_evals/sam3_silver_image_bdd100k.yaml b/train/configs/silver_image_evals/sam3_silver_image_bdd100k.yaml deleted file mode 100644 index e5587cf..0000000 --- a/train/configs/silver_image_evals/sam3_silver_image_bdd100k.yaml +++ /dev/null @@ -1,64 +0,0 @@ -# @package _global_ -defaults: - - /configs/eval_base.yaml - - _self_ - -# ============================================================================ -# Paths Configuration (you can override here, but it shouldn't require further changes if eval_base.yaml is correct -# ============================================================================ -paths: - experiment_log_dir: ${paths.base_experiment_log_dir}/silver_bdd100k/ - coco_gt: ${paths.base_annotation_path_silver}/silver_bdd100k_merged_test.json - img_path: ${paths.silver_img_path}/bdd100k/ - - - -# ============================================================================ -# Trainer Configuration -# ============================================================================ - -trainer: - data: - val: - _target_: sam3.train.data.torch_dataset.TorchDataset - dataset: - _target_: sam3.train.data.sam3_image_dataset.Sam3ImageDataset - coco_json_loader: - _target_: sam3.train.data.coco_json_loaders.SAM3_EVAL_API_FROM_JSON_NP - _partial_: true - img_folder: ${paths.img_path} - ann_file: ${paths.coco_gt} - transforms: ${scratch.base_val_transform} - 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: False - drop_last: False - collate_fn: - _target_: sam3.train.data.collator.collate_fn_api - _partial_: true - repeats: ${scratch.hybrid_repeats} - dict_key: silver_bdd100k - - meters: - val: - silver_bdd100k: # this key matches the "dict_key" in the dataloader's collate function - cgf1: - _target_: sam3.eval.coco_writer.PredictionDumper - iou_type: "segm" - dump_dir: ${launcher.experiment_log_dir}/dumps/silver_bdd100k - merge_predictions: True - postprocessor: ${scratch.mask_postprocessor_thresholded} - gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} - maxdets: 1000000 # no limit - pred_file_evaluators: - - _target_: sam3.eval.cgf1_eval.CGF1Evaluator - gt_path: ${paths.coco_gt} - iou_type: "bbox" - - _target_: sam3.eval.cgf1_eval.CGF1Evaluator - gt_path: ${paths.coco_gt} - iou_type: "segm" diff --git a/train/configs/silver_image_evals/sam3_silver_image_droid.yaml b/train/configs/silver_image_evals/sam3_silver_image_droid.yaml deleted file mode 100644 index c0d6234..0000000 --- a/train/configs/silver_image_evals/sam3_silver_image_droid.yaml +++ /dev/null @@ -1,64 +0,0 @@ -# @package _global_ -defaults: - - /configs/eval_base.yaml - - _self_ - -# ============================================================================ -# Paths Configuration (you can override here, but it shouldn't require further changes if eval_base.yaml is correct -# ============================================================================ -paths: - experiment_log_dir: ${paths.base_experiment_log_dir}/silver_droid/ - coco_gt: ${paths.base_annotation_path_silver}/silver_droid_merged_test.json - img_path: ${paths.silver_img_path}/droid/ - - - -# ============================================================================ -# Trainer Configuration -# ============================================================================ - -trainer: - data: - val: - _target_: sam3.train.data.torch_dataset.TorchDataset - dataset: - _target_: sam3.train.data.sam3_image_dataset.Sam3ImageDataset - coco_json_loader: - _target_: sam3.train.data.coco_json_loaders.SAM3_EVAL_API_FROM_JSON_NP - _partial_: true - img_folder: ${paths.img_path} - ann_file: ${paths.coco_gt} - transforms: ${scratch.base_val_transform} - 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: False - drop_last: False - collate_fn: - _target_: sam3.train.data.collator.collate_fn_api - _partial_: true - repeats: ${scratch.hybrid_repeats} - dict_key: silver_droid - - meters: - val: - silver_droid: # this key matches the "dict_key" in the dataloader's collate function - cgf1: - _target_: sam3.eval.coco_writer.PredictionDumper - iou_type: "segm" - dump_dir: ${launcher.experiment_log_dir}/dumps/silver_droid - merge_predictions: True - postprocessor: ${scratch.mask_postprocessor_thresholded} - gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} - maxdets: 1000000 # no limit - pred_file_evaluators: - - _target_: sam3.eval.cgf1_eval.CGF1Evaluator - gt_path: ${paths.coco_gt} - iou_type: "bbox" - - _target_: sam3.eval.cgf1_eval.CGF1Evaluator - gt_path: ${paths.coco_gt} - iou_type: "segm" diff --git a/train/configs/silver_image_evals/sam3_silver_image_ego4d.yaml b/train/configs/silver_image_evals/sam3_silver_image_ego4d.yaml deleted file mode 100644 index d5a036d..0000000 --- a/train/configs/silver_image_evals/sam3_silver_image_ego4d.yaml +++ /dev/null @@ -1,64 +0,0 @@ -# @package _global_ -defaults: - - /configs/eval_base.yaml - - _self_ - -# ============================================================================ -# Paths Configuration (you can override here, but it shouldn't require further changes if eval_base.yaml is correct -# ============================================================================ -paths: - experiment_log_dir: ${paths.base_experiment_log_dir}/silver_ego4d/ - coco_gt: ${paths.base_annotation_path_silver}/silver_ego4d_merged_test.json - img_path: ${paths.silver_img_path}/ego4d/ - - - -# ============================================================================ -# Trainer Configuration -# ============================================================================ - -trainer: - data: - val: - _target_: sam3.train.data.torch_dataset.TorchDataset - dataset: - _target_: sam3.train.data.sam3_image_dataset.Sam3ImageDataset - coco_json_loader: - _target_: sam3.train.data.coco_json_loaders.SAM3_EVAL_API_FROM_JSON_NP - _partial_: true - img_folder: ${paths.img_path} - ann_file: ${paths.coco_gt} - transforms: ${scratch.base_val_transform} - 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: False - drop_last: False - collate_fn: - _target_: sam3.train.data.collator.collate_fn_api - _partial_: true - repeats: ${scratch.hybrid_repeats} - dict_key: silver_ego4d - - meters: - val: - silver_ego4d: # this key matches the "dict_key" in the dataloader's collate function - cgf1: - _target_: sam3.eval.coco_writer.PredictionDumper - iou_type: "segm" - dump_dir: ${launcher.experiment_log_dir}/dumps/silver_ego4d - merge_predictions: True - postprocessor: ${scratch.mask_postprocessor_thresholded} - gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} - maxdets: 1000000 # no limit - pred_file_evaluators: - - _target_: sam3.eval.cgf1_eval.CGF1Evaluator - gt_path: ${paths.coco_gt} - iou_type: "bbox" - - _target_: sam3.eval.cgf1_eval.CGF1Evaluator - gt_path: ${paths.coco_gt} - iou_type: "segm" diff --git a/train/configs/silver_image_evals/sam3_silver_image_fathomnet.yaml b/train/configs/silver_image_evals/sam3_silver_image_fathomnet.yaml deleted file mode 100644 index b15d0c8..0000000 --- a/train/configs/silver_image_evals/sam3_silver_image_fathomnet.yaml +++ /dev/null @@ -1,64 +0,0 @@ -# @package _global_ -defaults: - - /configs/eval_base.yaml - - _self_ - -# ============================================================================ -# Paths Configuration (you can override here, but it shouldn't require further changes if eval_base.yaml is correct -# ============================================================================ -paths: - experiment_log_dir: ${paths.base_experiment_log_dir}/silver_fathomnet/ - coco_gt: ${paths.base_annotation_path_silver}/silver_fathomnet_test.json - img_path: ${paths.silver_img_path}/fathomnet/ - - - -# ============================================================================ -# Trainer Configuration -# ============================================================================ - -trainer: - data: - val: - _target_: sam3.train.data.torch_dataset.TorchDataset - dataset: - _target_: sam3.train.data.sam3_image_dataset.Sam3ImageDataset - coco_json_loader: - _target_: sam3.train.data.coco_json_loaders.SAM3_EVAL_API_FROM_JSON_NP - _partial_: true - img_folder: ${paths.img_path} - ann_file: ${paths.coco_gt} - transforms: ${scratch.base_val_transform} - 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: False - drop_last: False - collate_fn: - _target_: sam3.train.data.collator.collate_fn_api - _partial_: true - repeats: ${scratch.hybrid_repeats} - dict_key: silver_fathomnet - - meters: - val: - silver_fathomnet: # this key matches the "dict_key" in the dataloader's collate function - cgf1: - _target_: sam3.eval.coco_writer.PredictionDumper - iou_type: "segm" - dump_dir: ${launcher.experiment_log_dir}/dumps/silver_fathomnet - merge_predictions: True - postprocessor: ${scratch.mask_postprocessor_thresholded} - gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} - maxdets: 1000000 # no limit - pred_file_evaluators: - - _target_: sam3.eval.cgf1_eval.CGF1Evaluator - gt_path: ${paths.coco_gt} - iou_type: "bbox" - - _target_: sam3.eval.cgf1_eval.CGF1Evaluator - gt_path: ${paths.coco_gt} - iou_type: "segm" diff --git a/train/configs/silver_image_evals/sam3_silver_image_food_rec.yaml b/train/configs/silver_image_evals/sam3_silver_image_food_rec.yaml deleted file mode 100644 index 5158ff5..0000000 --- a/train/configs/silver_image_evals/sam3_silver_image_food_rec.yaml +++ /dev/null @@ -1,64 +0,0 @@ -# @package _global_ -defaults: - - /configs/eval_base.yaml - - _self_ - -# ============================================================================ -# Paths Configuration (you can override here, but it shouldn't require further changes if eval_base.yaml is correct -# ============================================================================ -paths: - experiment_log_dir: ${paths.base_experiment_log_dir}/silver_food_rec/ - coco_gt: ${paths.base_annotation_path_silver}/silver_food_rec_merged_test.json - img_path: ${paths.silver_img_path}/food_rec/ - - - -# ============================================================================ -# Trainer Configuration -# ============================================================================ - -trainer: - data: - val: - _target_: sam3.train.data.torch_dataset.TorchDataset - dataset: - _target_: sam3.train.data.sam3_image_dataset.Sam3ImageDataset - coco_json_loader: - _target_: sam3.train.data.coco_json_loaders.SAM3_EVAL_API_FROM_JSON_NP - _partial_: true - img_folder: ${paths.img_path} - ann_file: ${paths.coco_gt} - transforms: ${scratch.base_val_transform} - 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: False - drop_last: False - collate_fn: - _target_: sam3.train.data.collator.collate_fn_api - _partial_: true - repeats: ${scratch.hybrid_repeats} - dict_key: silver_food_rec - - meters: - val: - silver_food_rec: # this key matches the "dict_key" in the dataloader's collate function - cgf1: - _target_: sam3.eval.coco_writer.PredictionDumper - iou_type: "segm" - dump_dir: ${launcher.experiment_log_dir}/dumps/silver_food_rec - merge_predictions: True - postprocessor: ${scratch.mask_postprocessor_thresholded} - gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} - maxdets: 1000000 # no limit - pred_file_evaluators: - - _target_: sam3.eval.cgf1_eval.CGF1Evaluator - gt_path: ${paths.coco_gt} - iou_type: "bbox" - - _target_: sam3.eval.cgf1_eval.CGF1Evaluator - gt_path: ${paths.coco_gt} - iou_type: "segm" diff --git a/train/configs/silver_image_evals/sam3_silver_image_geode.yaml b/train/configs/silver_image_evals/sam3_silver_image_geode.yaml deleted file mode 100644 index 08f159f..0000000 --- a/train/configs/silver_image_evals/sam3_silver_image_geode.yaml +++ /dev/null @@ -1,64 +0,0 @@ -# @package _global_ -defaults: - - /configs/eval_base.yaml - - _self_ - -# ============================================================================ -# Paths Configuration (you can override here, but it shouldn't require further changes if eval_base.yaml is correct -# ============================================================================ -paths: - experiment_log_dir: ${paths.base_experiment_log_dir}/silver_geode/ - coco_gt: ${paths.base_annotation_path_silver}/silver_geode_merged_test.json - img_path: ${paths.silver_img_path}/geode/ - - - -# ============================================================================ -# Trainer Configuration -# ============================================================================ - -trainer: - data: - val: - _target_: sam3.train.data.torch_dataset.TorchDataset - dataset: - _target_: sam3.train.data.sam3_image_dataset.Sam3ImageDataset - coco_json_loader: - _target_: sam3.train.data.coco_json_loaders.SAM3_EVAL_API_FROM_JSON_NP - _partial_: true - img_folder: ${paths.img_path} - ann_file: ${paths.coco_gt} - transforms: ${scratch.base_val_transform} - 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: False - drop_last: False - collate_fn: - _target_: sam3.train.data.collator.collate_fn_api - _partial_: true - repeats: ${scratch.hybrid_repeats} - dict_key: silver_geode - - meters: - val: - silver_geode: # this key matches the "dict_key" in the dataloader's collate function - cgf1: - _target_: sam3.eval.coco_writer.PredictionDumper - iou_type: "segm" - dump_dir: ${launcher.experiment_log_dir}/dumps/silver_geode - merge_predictions: True - postprocessor: ${scratch.mask_postprocessor_thresholded} - gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} - maxdets: 1000000 # no limit - pred_file_evaluators: - - _target_: sam3.eval.cgf1_eval.CGF1Evaluator - gt_path: ${paths.coco_gt} - iou_type: "bbox" - - _target_: sam3.eval.cgf1_eval.CGF1Evaluator - gt_path: ${paths.coco_gt} - iou_type: "segm" diff --git a/train/configs/silver_image_evals/sam3_silver_image_inaturalist.yaml b/train/configs/silver_image_evals/sam3_silver_image_inaturalist.yaml deleted file mode 100644 index 1d56d97..0000000 --- a/train/configs/silver_image_evals/sam3_silver_image_inaturalist.yaml +++ /dev/null @@ -1,64 +0,0 @@ -# @package _global_ -defaults: - - /configs/eval_base.yaml - - _self_ - -# ============================================================================ -# Paths Configuration (you can override here, but it shouldn't require further changes if eval_base.yaml is correct -# ============================================================================ -paths: - experiment_log_dir: ${paths.base_experiment_log_dir}/silver_inaturalist/ - coco_gt: ${paths.base_annotation_path_silver}/silver_inaturalist_merged_test.json - img_path: ${paths.silver_img_path}/inaturalist/ - - - -# ============================================================================ -# Trainer Configuration -# ============================================================================ - -trainer: - data: - val: - _target_: sam3.train.data.torch_dataset.TorchDataset - dataset: - _target_: sam3.train.data.sam3_image_dataset.Sam3ImageDataset - coco_json_loader: - _target_: sam3.train.data.coco_json_loaders.SAM3_EVAL_API_FROM_JSON_NP - _partial_: true - img_folder: ${paths.img_path} - ann_file: ${paths.coco_gt} - transforms: ${scratch.base_val_transform} - 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: False - drop_last: False - collate_fn: - _target_: sam3.train.data.collator.collate_fn_api - _partial_: true - repeats: ${scratch.hybrid_repeats} - dict_key: silver_inaturalist - - meters: - val: - silver_inaturalist: # this key matches the "dict_key" in the dataloader's collate function - cgf1: - _target_: sam3.eval.coco_writer.PredictionDumper - iou_type: "segm" - dump_dir: ${launcher.experiment_log_dir}/dumps/silver_inaturalist - merge_predictions: True - postprocessor: ${scratch.mask_postprocessor_thresholded} - gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} - maxdets: 1000000 # no limit - pred_file_evaluators: - - _target_: sam3.eval.cgf1_eval.CGF1Evaluator - gt_path: ${paths.coco_gt} - iou_type: "bbox" - - _target_: sam3.eval.cgf1_eval.CGF1Evaluator - gt_path: ${paths.coco_gt} - iou_type: "segm" diff --git a/train/configs/silver_image_evals/sam3_silver_image_nga.yaml b/train/configs/silver_image_evals/sam3_silver_image_nga.yaml deleted file mode 100644 index b2de0af..0000000 --- a/train/configs/silver_image_evals/sam3_silver_image_nga.yaml +++ /dev/null @@ -1,64 +0,0 @@ -# @package _global_ -defaults: - - /configs/eval_base.yaml - - _self_ - -# ============================================================================ -# Paths Configuration (you can override here, but it shouldn't require further changes if eval_base.yaml is correct -# ============================================================================ -paths: - experiment_log_dir: ${paths.base_experiment_log_dir}/silver_nga_art/ - coco_gt: ${paths.base_annotation_path_silver}/silver_nga_art_merged_test.json - img_path: ${paths.silver_img_path}/nga/ - - - -# ============================================================================ -# Trainer Configuration -# ============================================================================ - -trainer: - data: - val: - _target_: sam3.train.data.torch_dataset.TorchDataset - dataset: - _target_: sam3.train.data.sam3_image_dataset.Sam3ImageDataset - coco_json_loader: - _target_: sam3.train.data.coco_json_loaders.SAM3_EVAL_API_FROM_JSON_NP - _partial_: true - img_folder: ${paths.img_path} - ann_file: ${paths.coco_gt} - transforms: ${scratch.base_val_transform} - 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: False - drop_last: False - collate_fn: - _target_: sam3.train.data.collator.collate_fn_api - _partial_: true - repeats: ${scratch.hybrid_repeats} - dict_key: silver_nga_art - - meters: - val: - silver_nga_art: # this key matches the "dict_key" in the dataloader's collate function - cgf1: - _target_: sam3.eval.coco_writer.PredictionDumper - iou_type: "segm" - dump_dir: ${launcher.experiment_log_dir}/dumps/silver_nga_art - merge_predictions: True - postprocessor: ${scratch.mask_postprocessor_thresholded} - gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} - maxdets: 1000000 # no limit - pred_file_evaluators: - - _target_: sam3.eval.cgf1_eval.CGF1Evaluator - gt_path: ${paths.coco_gt} - iou_type: "bbox" - - _target_: sam3.eval.cgf1_eval.CGF1Evaluator - gt_path: ${paths.coco_gt} - iou_type: "segm" diff --git a/train/configs/silver_image_evals/sam3_silver_image_sav.yaml b/train/configs/silver_image_evals/sam3_silver_image_sav.yaml deleted file mode 100644 index 7ebbb0f..0000000 --- a/train/configs/silver_image_evals/sam3_silver_image_sav.yaml +++ /dev/null @@ -1,64 +0,0 @@ -# @package _global_ -defaults: - - /configs/eval_base.yaml - - _self_ - -# ============================================================================ -# Paths Configuration (you can override here, but it shouldn't require further changes if eval_base.yaml is correct -# ============================================================================ -paths: - experiment_log_dir: ${paths.base_experiment_log_dir}/silver_sav/ - coco_gt: ${paths.base_annotation_path_silver}/silver_sav_merged_test.json - img_path: ${paths.silver_img_path}/sav/ - - - -# ============================================================================ -# Trainer Configuration -# ============================================================================ - -trainer: - data: - val: - _target_: sam3.train.data.torch_dataset.TorchDataset - dataset: - _target_: sam3.train.data.sam3_image_dataset.Sam3ImageDataset - coco_json_loader: - _target_: sam3.train.data.coco_json_loaders.SAM3_EVAL_API_FROM_JSON_NP - _partial_: true - img_folder: ${paths.img_path} - ann_file: ${paths.coco_gt} - transforms: ${scratch.base_val_transform} - 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: False - drop_last: False - collate_fn: - _target_: sam3.train.data.collator.collate_fn_api - _partial_: true - repeats: ${scratch.hybrid_repeats} - dict_key: silver_sav - - meters: - val: - silver_sav: # this key matches the "dict_key" in the dataloader's collate function - cgf1: - _target_: sam3.eval.coco_writer.PredictionDumper - iou_type: "segm" - dump_dir: ${launcher.experiment_log_dir}/dumps/silver_sav - merge_predictions: True - postprocessor: ${scratch.mask_postprocessor_thresholded} - gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} - maxdets: 1000000 # no limit - pred_file_evaluators: - - _target_: sam3.eval.cgf1_eval.CGF1Evaluator - gt_path: ${paths.coco_gt} - iou_type: "bbox" - - _target_: sam3.eval.cgf1_eval.CGF1Evaluator - gt_path: ${paths.coco_gt} - iou_type: "segm" diff --git a/train/configs/silver_image_evals/sam3_silver_image_yt1b.yaml b/train/configs/silver_image_evals/sam3_silver_image_yt1b.yaml deleted file mode 100644 index 901bd3a..0000000 --- a/train/configs/silver_image_evals/sam3_silver_image_yt1b.yaml +++ /dev/null @@ -1,64 +0,0 @@ -# @package _global_ -defaults: - - /configs/eval_base.yaml - - _self_ - -# ============================================================================ -# Paths Configuration (you can override here, but it shouldn't require further changes if eval_base.yaml is correct -# ============================================================================ -paths: - experiment_log_dir: ${paths.base_experiment_log_dir}/silver_yt1b/ - coco_gt: ${paths.base_annotation_path_silver}/silver_yt1b_merged_test.json - img_path: ${paths.silver_img_path}/yt1b/ - - - -# ============================================================================ -# Trainer Configuration -# ============================================================================ - -trainer: - data: - val: - _target_: sam3.train.data.torch_dataset.TorchDataset - dataset: - _target_: sam3.train.data.sam3_image_dataset.Sam3ImageDataset - coco_json_loader: - _target_: sam3.train.data.coco_json_loaders.SAM3_EVAL_API_FROM_JSON_NP - _partial_: true - img_folder: ${paths.img_path} - ann_file: ${paths.coco_gt} - transforms: ${scratch.base_val_transform} - 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: False - drop_last: False - collate_fn: - _target_: sam3.train.data.collator.collate_fn_api - _partial_: true - repeats: ${scratch.hybrid_repeats} - dict_key: silver_yt1b - - meters: - val: - silver_yt1b: # this key matches the "dict_key" in the dataloader's collate function - cgf1: - _target_: sam3.eval.coco_writer.PredictionDumper - iou_type: "segm" - dump_dir: ${launcher.experiment_log_dir}/dumps/silver_yt1b - merge_predictions: True - postprocessor: ${scratch.mask_postprocessor_thresholded} - gather_pred_via_filesys: ${scratch.gather_pred_via_filesys} - maxdets: 1000000 # no limit - pred_file_evaluators: - - _target_: sam3.eval.cgf1_eval.CGF1Evaluator - gt_path: ${paths.coco_gt} - iou_type: "bbox" - - _target_: sam3.eval.cgf1_eval.CGF1Evaluator - gt_path: ${paths.coco_gt} - iou_type: "segm" diff --git a/train/data/__init__.py b/train/data/__init__.py deleted file mode 100644 index 46d37d2..0000000 --- a/train/data/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved diff --git a/train/data/coco_json_loaders.py b/train/data/coco_json_loaders.py deleted file mode 100644 index 4bdf487..0000000 --- a/train/data/coco_json_loaders.py +++ /dev/null @@ -1,465 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved - -import json -from collections import defaultdict -from typing import Dict, List, Tuple - -import torch -from pycocotools import mask as mask_util - - -# ============================================================================ -# Utility Functions -# ============================================================================ - - -def convert_boxlist_to_normalized_tensor(box_list, image_width, image_height): - """ - Converts a list of bounding boxes to a normalized PyTorch tensor. - - Args: - box_list (list of list or tuples): Each box is [x_min, y_min, x_max, y_max]. - image_width (int or float): Width of the image. - image_height (int or float): Height of the image. - - Returns: - torch.Tensor: Normalized tensor of shape (N, 4), values in [0, 1]. - """ - boxes = torch.tensor(box_list, dtype=torch.float32) - boxes[:, [0, 2]] /= image_width # x_min, x_max - boxes[:, [1, 3]] /= image_height # y_min, y_max - boxes = boxes.clamp(0, 1) - return boxes - - -def load_coco_and_group_by_image(json_path: str) -> Tuple[List[Dict], Dict[int, str]]: - """ - Load COCO JSON file and group annotations by image. - - Args: - json_path (str): Path to COCO JSON file. - - Returns: - Tuple containing: - - List of dicts with 'image' and 'annotations' keys - - Dict mapping category IDs to category names - """ - with open(json_path, "r") as f: - coco = json.load(f) - - images = {img["id"]: img for img in coco["images"]} - - anns_by_image = defaultdict(list) - for ann in coco["annotations"]: - anns_by_image[ann["image_id"]].append(ann) - - sorted_image_ids = sorted(images.keys()) - - grouped = [] - for image_id in sorted_image_ids: - image_info = images[image_id] - grouped.append( - {"image": image_info, "annotations": anns_by_image.get(image_id, [])} - ) - - cat_id_to_name = {cat["id"]: cat["name"] for cat in coco["categories"]} - - return grouped, cat_id_to_name - - -def ann_to_rle(segm, im_info: Dict) -> Dict: - """ - Convert annotation which can be polygons or uncompressed RLE to RLE. - - Args: - segm: Segmentation data (polygon list or RLE dict) - im_info (dict): Image info containing 'height' and 'width' - - Returns: - RLE encoded segmentation - """ - h, w = im_info["height"], im_info["width"] - - if isinstance(segm, list): - # Polygon - merge all parts into one mask RLE code - rles = mask_util.frPyObjects(segm, h, w) - rle = mask_util.merge(rles) - elif isinstance(segm["counts"], list): - # Uncompressed RLE - rle = mask_util.frPyObjects(segm, h, w) - else: - # Already RLE - rle = segm - - return rle - - -# ============================================================================ -# COCO Training API -# ============================================================================ - - -class COCO_FROM_JSON: - """ - COCO training API for loading box-only annotations from JSON. - Groups all annotations per image and creates queries per category. - """ - - def __init__( - self, - annotation_file, - prompts=None, - include_negatives=True, - category_chunk_size=None, - ): - """ - Initialize the COCO training API. - - Args: - annotation_file (str): Path to COCO JSON annotation file - prompts: Optional custom prompts for categories - include_negatives (bool): Whether to include negative examples (categories with no instances) - """ - self._raw_data, self._cat_idx_to_text = load_coco_and_group_by_image( - annotation_file - ) - self._sorted_cat_ids = sorted(list(self._cat_idx_to_text.keys())) - self.prompts = None - self.include_negatives = include_negatives - self.category_chunk_size = ( - category_chunk_size - if category_chunk_size is not None - else len(self._sorted_cat_ids) - ) - self.category_chunks = [ - self._sorted_cat_ids[i : i + self.category_chunk_size] - for i in range(0, len(self._sorted_cat_ids), self.category_chunk_size) - ] - if prompts is not None: - prompts = eval(prompts) - self.prompts = {} - for loc_dict in prompts: - self.prompts[int(loc_dict["id"])] = loc_dict["name"] - assert len(self.prompts) == len( - self._sorted_cat_ids - ), "Number of prompts must match number of categories" - - def getDatapointIds(self): - """Return all datapoint indices for training.""" - return list(range(len(self._raw_data) * len(self.category_chunks))) - - def loadQueriesAndAnnotationsFromDatapoint(self, idx): - """ - Load queries and annotations for a specific datapoint. - - Args: - idx (int): Datapoint index - - Returns: - Tuple of (queries, annotations) lists - """ - img_idx = idx // len(self.category_chunks) - chunk_idx = idx % len(self.category_chunks) - cat_chunk = self.category_chunks[chunk_idx] - - queries = [] - annotations = [] - - query_template = { - "id": None, - "original_cat_id": None, - "object_ids_output": None, - "query_text": None, - "query_processing_order": 0, - "ptr_x_query_id": None, - "ptr_y_query_id": None, - "image_id": 0, # Single image per datapoint - "input_box": None, - "input_box_label": None, - "input_points": None, - "is_exhaustive": True, - } - - annot_template = { - "image_id": 0, - "bbox": None, # Normalized bbox in xywh - "area": None, # Unnormalized area - "segmentation": None, # RLE encoded - "object_id": None, - "is_crowd": None, - "id": None, - } - - raw_annotations = self._raw_data[img_idx]["annotations"] - image_info = self._raw_data[img_idx]["image"] - width, height = image_info["width"], image_info["height"] - - # Group annotations by category - cat_id_to_anns = defaultdict(list) - for ann in raw_annotations: - cat_id_to_anns[ann["category_id"]].append(ann) - - annotations_by_cat_sorted = [ - (cat_id, cat_id_to_anns[cat_id]) for cat_id in cat_chunk - ] - - for cat_id, anns in annotations_by_cat_sorted: - if len(anns) == 0 and not self.include_negatives: - continue - - cur_ann_ids = [] - - # Create annotations for this category - for ann in anns: - annotation = annot_template.copy() - annotation["id"] = len(annotations) - annotation["object_id"] = annotation["id"] - annotation["is_crowd"] = ann["iscrowd"] - - normalized_boxes = convert_boxlist_to_normalized_tensor( - [ann["bbox"]], width, height - ) - bbox = normalized_boxes[0] - - annotation["area"] = (bbox[2] * bbox[3]).item() - annotation["bbox"] = bbox - - if ( - "segmentation" in ann - and ann["segmentation"] is not None - and ann["segmentation"] != [] - ): - annotation["segmentation"] = ann_to_rle( - ann["segmentation"], im_info=image_info - ) - - annotations.append(annotation) - cur_ann_ids.append(annotation["id"]) - - # Create query for this category - query = query_template.copy() - query["id"] = len(queries) - query["original_cat_id"] = cat_id - query["query_text"] = ( - self._cat_idx_to_text[cat_id] - if self.prompts is None - else self.prompts[cat_id] - ) - query["object_ids_output"] = cur_ann_ids - queries.append(query) - - return queries, annotations - - def loadImagesFromDatapoint(self, idx): - """ - Load image information for a specific datapoint. - - Args: - idx (int): Datapoint index - - Returns: - List containing image info dict - """ - img_idx = idx // len(self.category_chunks) - img_data = self._raw_data[img_idx]["image"] - images = [ - { - "id": 0, - "file_name": img_data["file_name"], - "original_img_id": img_data["id"], - "coco_img_id": img_data["id"], - } - ] - return images - - -# ============================================================================ -# SAM3 Evaluation APIs -# ============================================================================ - - -class SAM3_EVAL_API_FROM_JSON_NP: - """ - SAM3 evaluation API for loading noun phrase queries from JSON. - """ - - def __init__(self, annotation_file): - """ - Initialize the SAM3 evaluation API. - - Args: - annotation_file (str): Path to SAM3 JSON annotation file - """ - with open(annotation_file, "r") as f: - data = json.load(f) - self._image_data = data["images"] - - def getDatapointIds(self): - """Return all datapoint indices.""" - return list(range(len(self._image_data))) - - def loadQueriesAndAnnotationsFromDatapoint(self, idx): - """ - Load queries and annotations for a specific datapoint. - - Args: - idx (int): Datapoint index - - Returns: - Tuple of (queries, annotations) lists - """ - cur_img_data = self._image_data[idx] - queries = [] - annotations = [] - - query_template = { - "id": None, - "original_cat_id": None, - "object_ids_output": None, - "query_text": None, - "query_processing_order": 0, - "ptr_x_query_id": None, - "ptr_y_query_id": None, - "image_id": 0, - "input_box": None, - "input_box_label": None, - "input_points": None, - "is_exhaustive": True, - } - - # Create query - query = query_template.copy() - query["id"] = len(queries) - query["original_cat_id"] = int(cur_img_data["queried_category"]) - query["query_text"] = cur_img_data["text_input"] - query["object_ids_output"] = [] - queries.append(query) - - return queries, annotations - - def loadImagesFromDatapoint(self, idx): - """ - Load image information for a specific datapoint. - - Args: - idx (int): Datapoint index - - Returns: - List containing image info dict - """ - img_data = self._image_data[idx] - images = [ - { - "id": 0, - "file_name": img_data["file_name"], - "original_img_id": img_data["id"], - "coco_img_id": img_data["id"], - } - ] - return images - - -class SAM3_VEVAL_API_FROM_JSON_NP: - """ - SAM3 video evaluation API for loading noun phrase queries from JSON. - """ - - def __init__(self, annotation_file): - """ - Initialize the SAM3 video evaluation API. - - Args: - annotation_file (str): Path to SAM3 video JSON annotation file - """ - with open(annotation_file, "r") as f: - data = json.load(f) - - assert "video_np_pairs" in data, "Incorrect data format" - - self._video_data = data["videos"] - self._video_id_to_np_ids = defaultdict(list) - self._cat_id_to_np = {} - - for cat_dict in data["categories"]: - self._cat_id_to_np[cat_dict["id"]] = cat_dict["name"] - - for video_np_dict in data["video_np_pairs"]: - self._video_id_to_np_ids[video_np_dict["video_id"]].append( - video_np_dict["category_id"] - ) - assert ( - self._cat_id_to_np[video_np_dict["category_id"]] - == video_np_dict["noun_phrase"] - ), "Category name does not match text input" - - def getDatapointIds(self): - """Return all datapoint indices.""" - return list(range(len(self._video_data))) - - def loadQueriesAndAnnotationsFromDatapoint(self, idx): - """ - Load queries and annotations for a specific video datapoint. - - Args: - idx (int): Datapoint index - - Returns: - Tuple of (queries, annotations) lists - """ - cur_vid_data = self._video_data[idx] - queries = [] - annotations = [] - - query_template = { - "id": None, - "original_cat_id": None, - "object_ids_output": None, - "query_text": None, - "query_processing_order": 0, - "ptr_x_query_id": None, - "ptr_y_query_id": None, - "image_id": 0, - "input_box": None, - "input_box_label": None, - "input_points": None, - "is_exhaustive": True, - } - - all_np_ids = self._video_id_to_np_ids[cur_vid_data["id"]] - - for np_id in all_np_ids: - text_input = self._cat_id_to_np[np_id] - - for i, image_path in enumerate(cur_vid_data["file_names"]): - query = query_template.copy() - query["id"] = len(queries) - query["original_cat_id"] = np_id - query["query_text"] = text_input - query["image_id"] = i - query["query_processing_order"] = i - query["object_ids_output"] = [] - queries.append(query) - - return queries, annotations - - def loadImagesFromDatapoint(self, idx): - """ - Load image information for a specific video datapoint. - - Args: - idx (int): Datapoint index - - Returns: - List containing image info dicts for all frames - """ - video_data = self._video_data[idx] - images = [ - { - "id": i, - "file_name": file_name, - "original_img_id": video_data["id"], - "coco_img_id": video_data["id"], - } - for i, file_name in enumerate(video_data["file_names"]) - ] - return images diff --git a/train/data/collator.py b/train/data/collator.py deleted file mode 100644 index e0f7e2e..0000000 --- a/train/data/collator.py +++ /dev/null @@ -1,360 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved - -from dataclasses import dataclass, field as field_ptr_behaviour, fields, is_dataclass -from typing import Any, get_args, get_origin, List, Union - -import torch - -from sam3.model.data_misc import ( - BatchedDatapoint, - BatchedFindTarget, - BatchedInferenceMetadata, - FindStage, -) - -from .sam3_image_dataset import Datapoint - - -MyTensor = Union[torch.Tensor, List[Any]] - - -def convert_my_tensors(obj): - def is_optional_field(field) -> bool: - return get_origin(field) is Union and type(None) in get_args(field) - - for field in fields(obj): - if is_dataclass(getattr(obj, field.name)): - convert_my_tensors(getattr(obj, field.name)) - continue - - field_type = field.type - if is_optional_field(field.type): - field_type = Union[get_args(field.type)[:-1]] # Get the Optional field type - - if field_type != MyTensor or getattr(obj, field.name) is None: - continue - - elif len(getattr(obj, field.name)) and isinstance( - getattr(obj, field.name)[0], torch.Tensor - ): - stack_dim = 0 - if field.name in [ - "input_boxes", - "input_boxes_label", - ]: - stack_dim = 1 - setattr( - obj, - field.name, - torch.stack(getattr(obj, field.name), dim=stack_dim).to( - getattr(obj, field.name + "__type") - ), - ) - else: - setattr( - obj, - field.name, - torch.as_tensor( - getattr(obj, field.name), dtype=getattr(obj, field.name + "__type") - ), - ) - return obj - - -def packed_to_padded_naive(boxes_packed, num_boxes, fill_value=0): - """ - Convert a packed tensor of bounding boxes to a padded tensor of bounding - boxes. Naive implementation using a loop. - - Inputs: - - boxes_packed: Tensor of shape (N_1 + ... + N_B, 4) - - num_boxes: Tensor of shape (B,) where num_boxes[i] = N_i - - Returns: - - boxes_padded: Tensor of shape (B, N_max, 4) where N_max = max_i N_i - """ - B = num_boxes.shape[0] - Ns = num_boxes.tolist() - - boxes_padded = boxes_packed.new_zeros(B, max(Ns), *boxes_packed.shape[1:]) - if fill_value != 0: - boxes_padded[...] = fill_value - prev_idx = 0 - for i in range(B): - next_idx = prev_idx + Ns[i] - boxes_padded[i, : Ns[i]] = boxes_packed[prev_idx:next_idx] - prev_idx = next_idx - return boxes_padded - - -def pad_tensor_list_to_longest( - tensors: List[torch.Tensor], dim=0, pad_val=0 -) -> List[torch.Tensor]: - # Edits the list in-place - if not tensors: - return tensors - pad_len = max(t.shape[dim] for t in tensors) - for i in range(len(tensors)): - n_dims = len(tensors[i].shape) - n_right_dims = (n_dims - 1) - (n_dims + dim) % n_dims - n_pad = pad_len - tensors[i].shape[dim] - pad_tuple = tuple([0] * 2 * n_right_dims + [0, n_pad]) - tensors[i] = torch.nn.functional.pad(tensors[i], pad_tuple, value=pad_val) - return tensors - - -def collate_fn_api_with_chunking( - batch, - num_chunks, - dict_key, - with_seg_masks=False, - input_points_embedding_dim=257, - repeats: int = 0, - load_image_in_fp16: bool = False, -): - assert num_chunks >= 1, "num_chunks must be >= 1" - - # split the batch into num_chunks chunks - batch_chunks = [batch[i::num_chunks] for i in range(num_chunks)] - - # collate each chunk - collated_chunks = [ - collate_fn_api( - chunk, - dict_key, - with_seg_masks, - input_points_embedding_dim, - repeats, - # ptr_behaviour, - load_image_in_fp16, - ) - for chunk in batch_chunks - ] - return collated_chunks - - -def collate_fn_api( - batch: List[Datapoint], - dict_key, - with_seg_masks=False, - input_points_embedding_dim=257, - repeats: int = 0, - load_image_in_fp16: bool = False, -): - # img_batch = torch.stack(sum([[img.data for img in v.images] for v in batch], [])) - img_batch = [] - text_batch = [] - raw_images = None - - num_stages = ( - max(q.query_processing_order for data in batch for q in data.find_queries) + 1 - ) - - stages = [ - FindStage( - img_ids=[], - text_ids=[], - input_boxes=[], - input_boxes_label=[], - input_boxes_mask=[], - input_points=[], - input_points_mask=[], - object_ids=[], - ) - for _ in range(num_stages) - ] - find_targets = [ - BatchedFindTarget( - num_boxes=[], - boxes=[], - boxes_padded=[], - is_exhaustive=[], - segments=[], - semantic_segments=[], - is_valid_segment=[], - repeated_boxes=[], - object_ids=[], - object_ids_padded=[], - ) - for _ in range(num_stages) - ] - find_metadatas = [ - BatchedInferenceMetadata( - coco_image_id=[], - original_size=[], - object_id=[], - frame_index=[], - original_image_id=[], - original_category_id=[], - is_conditioning_only=[], - ) - for _ in range(num_stages) - ] - - offset_img_id = 0 - offset_query_id = [0 for _ in range(num_stages)] - for i, data in enumerate(batch): - img_batch.extend([img.data for img in data.images]) - - if data.raw_images is not None: - if raw_images is None: - raw_images = [] - raw_images.extend(data.raw_images) - - # Conversion of query_ids indexing in a datapoint to query_ids indexing in a stage - datapoint_query_id_2_stage_query_id = [] - for q in data.find_queries: - stage_id = q.query_processing_order - datapoint_query_id_2_stage_query_id.append(offset_query_id[stage_id]) - offset_query_id[stage_id] += 1 - - for j, q in enumerate(data.find_queries): - stage_id = q.query_processing_order - stages[stage_id].img_ids.append(q.image_id + offset_img_id) - if q.query_text not in text_batch: - text_batch.append(q.query_text) - stages[stage_id].text_ids.append(text_batch.index(q.query_text)) - - assert ( - q.inference_metadata is not None - ), "inference_metadata must be provided when FindQueryLoaded is created." - for f in fields(q.inference_metadata): - getattr(find_metadatas[stage_id], f.name).append( - getattr(q.inference_metadata, f.name) - ) - - if q.input_bbox is not None: - assert q.input_bbox.numel() % 4 == 0 - assert q.input_bbox_label is not None - nb_boxes = q.input_bbox.numel() // 4 - assert len(q.input_bbox_label) == nb_boxes - stages[stage_id].input_boxes.append(q.input_bbox.view(nb_boxes, 4)) - stages[stage_id].input_boxes_label.append( - q.input_bbox_label.view(nb_boxes) - ) - stages[stage_id].input_boxes_mask.append( - torch.zeros(nb_boxes, dtype=torch.bool) - ) - else: - stages[stage_id].input_boxes.append(torch.zeros(0, 4)) - stages[stage_id].input_boxes_label.append( - torch.zeros(0, dtype=torch.bool) - ) - stages[stage_id].input_boxes_mask.append( - torch.ones(0, dtype=torch.bool) - ) - - if q.input_points is not None: - stages[stage_id].input_points.append( - q.input_points.squeeze(0) # Strip a trivial batch index - ) - # All masks will be padded up to the longest length - # with 1s before final conversion to batchd tensors - stages[stage_id].input_points_mask.append( - torch.zeros(q.input_points.shape[1]) - ) - else: - stages[stage_id].input_points.append( - torch.empty(0, input_points_embedding_dim) - ) - stages[stage_id].input_points_mask.append(torch.empty(0)) - - current_out_boxes = [] - current_out_object_ids = [] - # Set the object ids referred to by this query - stages[stage_id].object_ids.append(q.object_ids_output) - for object_id in q.object_ids_output: - current_out_boxes.append( - data.images[q.image_id].objects[object_id].bbox - ) - current_out_object_ids.append(object_id) - find_targets[stage_id].boxes.extend(current_out_boxes) - find_targets[stage_id].object_ids.extend(current_out_object_ids) - if repeats > 0: - for _ in range(repeats): - find_targets[stage_id].repeated_boxes.extend(current_out_boxes) - find_targets[stage_id].num_boxes.append(len(current_out_boxes)) - find_targets[stage_id].is_exhaustive.append(q.is_exhaustive) - - if with_seg_masks: - current_seg_mask = [] - current_is_valid_segment = [] - for object_id in q.object_ids_output: - seg_mask = data.images[q.image_id].objects[object_id].segment - if seg_mask is not None: - current_seg_mask.append(seg_mask) - current_is_valid_segment.append(1) - else: - dummy_mask = torch.zeros( - data.images[q.image_id].data.shape[-2:], dtype=torch.bool - ) - current_seg_mask.append(dummy_mask) - current_is_valid_segment.append(0) - find_targets[stage_id].segments.extend(current_seg_mask) - find_targets[stage_id].is_valid_segment.extend(current_is_valid_segment) - else: - # We are not loading segmentation masks - find_targets[stage_id].segments = None - find_targets[stage_id].is_valid_segment = None - - if q.semantic_target is not None: - find_targets[stage_id].semantic_segments.append(q.semantic_target) - - offset_img_id += len(data.images) - - # Pad input points to equal sequence lengths - for i in range(len(stages)): - stages[i].input_points = pad_tensor_list_to_longest( - stages[i].input_points, dim=0, pad_val=0 - ) - # Masked-out regions indicated by 1s. - stages[i].input_points_mask = pad_tensor_list_to_longest( - stages[i].input_points_mask, dim=0, pad_val=1 - ) - - # Pad input boxes to equal sequence lengths - for i in range(len(stages)): - stages[i].input_boxes = pad_tensor_list_to_longest( - stages[i].input_boxes, dim=0, pad_val=0 - ) - stages[i].input_boxes_label = pad_tensor_list_to_longest( - stages[i].input_boxes_label, dim=0, pad_val=0 - ) - # Masked-out regions indicated by 1s. - stages[i].input_boxes_mask = pad_tensor_list_to_longest( - stages[i].input_boxes_mask, dim=0, pad_val=1 - ) - - # Convert to tensors - for i in range(len(stages)): - stages[i] = convert_my_tensors(stages[i]) - find_targets[i] = convert_my_tensors(find_targets[i]) - find_metadatas[i] = convert_my_tensors(find_metadatas[i]) - # get padded representation for the boxes - find_targets[i].boxes_padded = packed_to_padded_naive( - find_targets[i].boxes.view(-1, 4), find_targets[i].num_boxes - ) - find_targets[i].object_ids_padded = packed_to_padded_naive( - find_targets[i].object_ids, find_targets[i].num_boxes, fill_value=-1 - ) - - # Finalize the image batch - # check sizes - for img in img_batch[1:]: - assert img.shape == img_batch[0].shape, "All images must have the same size" - image_batch = torch.stack(img_batch) - if load_image_in_fp16: - # Optionally, cast the image tensors to fp16, which helps save GPU memory on - # long videos with thousands of frames (where image tensors could be several GBs) - image_batch = image_batch.half() - - return { - dict_key: BatchedDatapoint( - img_batch=image_batch, - find_text_batch=text_batch, - find_inputs=stages, - find_targets=find_targets, - find_metadatas=find_metadatas, - raw_images=raw_images, - ) - } diff --git a/train/data/sam3_image_dataset.py b/train/data/sam3_image_dataset.py deleted file mode 100644 index 97efb1d..0000000 --- a/train/data/sam3_image_dataset.py +++ /dev/null @@ -1,528 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved - -"""Dataset class for modulated detection""" - -import json -import os -import random -import sys -import traceback -from collections import Counter -from dataclasses import dataclass -from enum import Enum -from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union - -import torch -import torch.utils.data -import torchvision -from decord import cpu, VideoReader -from iopath.common.file_io import g_pathmgr - -from PIL import Image as PILImage -from PIL.Image import DecompressionBombError - -from sam3.model.box_ops import box_xywh_to_xyxy -from torchvision.datasets.vision import VisionDataset - -from .coco_json_loaders import COCO_FROM_JSON - - -@dataclass -class InferenceMetadata: - """Metadata required for postprocessing""" - - # Coco id that corresponds to the "image" for evaluation by the coco evaluator - # This is used for our own "class agnostic" evaluation - coco_image_id: int - - # id in the original dataset, such that we can use the original evaluator - original_image_id: int - - # Original category id (if we want to use the original evaluator) - original_category_id: int - - # Size of the raw image (height, width) - original_size: Tuple[int, int] - - # Id of the object in the media - object_id: int - - # Index of the frame in the media (0 if single image) - frame_index: int - - # Whether it is for conditioning only, e.g., 0-th frame in TA is for conditioning - # as we assume GT available in frame 0. - is_conditioning_only: Optional[bool] = False - - -@dataclass -class FindQuery: - query_text: str - - image_id: int - - # In case of a find query, the list of object ids that have to be predicted - object_ids_output: List[int] - - # This is "instance exhaustivity". - # true iff all instances are separable and annotated - # See below the slightly different "pixel exhaustivity" - is_exhaustive: bool - - # The order in which the queries are processed (only meaningful for video) - query_processing_order: int = 0 - - # Input geometry, initially in denormalized XYXY format. Then - # 1. converted to normalized CxCyWH by the Normalize transform - input_bbox: Optional[torch.Tensor] = None - input_bbox_label: Optional[torch.Tensor] = None - - # Only for the PVS task - input_points: Optional[torch.Tensor] = None - - semantic_target: Optional[torch.Tensor] = None - - # pixel exhaustivity: true iff the union of all segments (including crowds) - # covers every pixel belonging to the target class - # Note that instance_exhaustive implies pixel_exhaustive - is_pixel_exhaustive: Optional[bool] = None - - -@dataclass -class FindQueryLoaded(FindQuery): - # Must have default value since FindQuery has entries with default values - inference_metadata: Optional[InferenceMetadata] = None - - -@dataclass -class Object: - # Initially in denormalized XYXY format, gets converted to normalized CxCyWH by the Normalize transform - bbox: torch.Tensor - area: float - - # Id of the object in the media - object_id: Optional[int] = -1 - - # Index of the frame in the media (0 if single image) - frame_index: Optional[int] = -1 - - segment: Optional[Union[torch.Tensor, dict]] = None # RLE dict or binary mask - - is_crowd: bool = False - - source: Optional[str] = None - - -@dataclass -class Image: - data: Union[torch.Tensor, PILImage.Image] - objects: List[Object] - size: Tuple[int, int] # (height, width) - - # For blurring augmentation - blurring_mask: Optional[Dict[str, Any]] = None - - -@dataclass -class Datapoint: - """Refers to an image/video and all its annotations""" - - find_queries: List[FindQueryLoaded] - images: List[Image] - raw_images: Optional[List[PILImage.Image]] = None - - -class CustomCocoDetectionAPI(VisionDataset): - """`MS Coco Detection `_ Dataset. - - Args: - root (string): Root directory where images are downloaded to. - annFile (string): Path to json annotation file. - transform (callable, optional): A function/transform that takes in an PIL image - and returns a transformed version. E.g, ``transforms.ToTensor`` - target_transform (callable, optional): A function/transform that takes in the - target and transforms it. - transforms (callable, optional): A function/transform that takes input sample and its target as entry - and returns a transformed version. - """ - - def __init__( - self, - root: str, - annFile: str, - load_segmentation: bool, - fix_fname: bool = False, - training: bool = True, - blurring_masks_path: Optional[str] = None, - use_caching: bool = True, - zstd_dict_path=None, - filter_query=None, - coco_json_loader: Callable = COCO_FROM_JSON, - limit_ids: int = None, - ) -> None: - super().__init__(root) - - self.annFile = annFile - self.use_caching = use_caching - self.zstd_dict_path = zstd_dict_path - - self.curr_epoch = 0 # Used in case data loader behavior changes across epochs - self.load_segmentation = load_segmentation - self.fix_fname = fix_fname - self.filter_query = filter_query - - self.coco = None - self.coco_json_loader = coco_json_loader - self.limit_ids = limit_ids - self.set_sharded_annotation_file(0) - self.training = training - self.blurring_masks_path = blurring_masks_path - - def _load_images( - self, datapoint_id: int, img_ids_to_load: Optional[Set[int]] = None - ) -> Tuple[List[Tuple[int, PILImage.Image]], List[Dict[str, Any]]]: - all_images = [] - all_img_metadata = [] - for current_meta in self.coco.loadImagesFromDatapoint(datapoint_id): - img_id = current_meta["id"] - if img_ids_to_load is not None and img_id not in img_ids_to_load: - continue - if self.fix_fname: - current_meta["file_name"] = current_meta["file_name"].split("/")[-1] - path = current_meta["file_name"] - if self.blurring_masks_path is not None: - mask_fname = os.path.basename(path).replace(".jpg", "-mask.json") - mask_path = os.path.join(self.blurring_masks_path, mask_fname) - if os.path.exists(mask_path): - with open(mask_path, "r") as fopen: - current_meta["blurring_mask"] = json.load(fopen) - - all_img_metadata.append(current_meta) - path = os.path.join(self.root, path) - try: - if ".mp4" in path and path[-4:] == ".mp4": - # Going to load a video frame - video_path, frame = path.split("@") - video = VideoReader(video_path, ctx=cpu(0)) - # Convert to PIL image - all_images.append( - ( - img_id, - torchvision.transforms.ToPILImage()( - video[int(frame)].asnumpy() - ), - ) - ) - else: - with g_pathmgr.open(path, "rb") as fopen: - all_images.append((img_id, PILImage.open(fopen).convert("RGB"))) - except FileNotFoundError as e: - print(f"File not found: {path} from dataset: {self.annFile}") - raise e - - return all_images, all_img_metadata - - def set_curr_epoch(self, epoch: int): - self.curr_epoch = epoch - - def set_epoch(self, epoch: int): - pass - - def set_sharded_annotation_file(self, data_epoch: int): - if self.coco is not None: - return - - assert g_pathmgr.isfile( - self.annFile - ), f"please provide valid annotation file. Missing: {self.annFile}" - annFile = g_pathmgr.get_local_path(self.annFile) - - if self.coco is not None: - del self.coco - - self.coco = self.coco_json_loader(annFile) - # Use a torch tensor here to optimize memory usage when using several dataloaders - ids_list = list(sorted(self.coco.getDatapointIds())) - if self.limit_ids is not None: - local_random = random.Random(len(ids_list)) - local_random.shuffle(ids_list) - ids_list = ids_list[: self.limit_ids] - self.ids = torch.as_tensor(ids_list, dtype=torch.long) - - def __getitem__(self, index: int) -> Datapoint: - return self._load_datapoint(index) - - def _load_datapoint(self, index: int) -> Datapoint: - """A separate method for easy overriding in subclasses.""" - id = self.ids[index].item() - pil_images, img_metadata = self._load_images(id) - queries, annotations = self.coco.loadQueriesAndAnnotationsFromDatapoint(id) - return self.load_queries(pil_images, annotations, queries, img_metadata) - - def load_queries(self, pil_images, annotations, queries, img_metadata): - """Transform the raw image and queries into a Datapoint sample.""" - images: List[Image] = [] - id2index_img = {} - id2index_obj = {} - id2index_find_query = {} - id2imsize = {} - assert len(pil_images) == len(img_metadata) - for i in range(len(pil_images)): - w, h = pil_images[i][1].size - blurring_mask = None - if "blurring_mask" in img_metadata[i]: - blurring_mask = img_metadata[i]["blurring_mask"] - images.append( - Image( - data=pil_images[i][1], - objects=[], - size=(h, w), - blurring_mask=blurring_mask, - ) - ) - id2index_img[pil_images[i][0]] = i - id2imsize[pil_images[i][0]] = (h, w) - - for annotation in annotations: - image_id = id2index_img[annotation["image_id"]] - bbox = box_xywh_to_xyxy(torch.as_tensor(annotation["bbox"])).view(1, 4) - h, w = id2imsize[annotation["image_id"]] - bbox[:, 0::2].mul_(w).clamp_(min=0, max=w) - bbox[:, 1::2].mul_(h).clamp_(min=0, max=h) - segment = None - if self.load_segmentation and "segmentation" in annotation: - # We're not decoding the RLE here, a transform will do it lazily later - segment = annotation["segmentation"] - images[image_id].objects.append( - Object( - bbox=bbox[0], - area=annotation["area"], - object_id=( - annotation["object_id"] if "object_id" in annotation else -1 - ), - frame_index=( - annotation["frame_index"] if "frame_index" in annotation else -1 - ), - segment=segment, - is_crowd=( - annotation["is_crowd"] if "is_crowd" in annotation else None - ), - source=annotation["source"] if "source" in annotation else "", - ) - ) - id2index_obj[annotation["id"]] = len(images[image_id].objects) - 1 - - find_queries = [] - stage2num_queries = Counter() - for i, query in enumerate(queries): - stage2num_queries[query["query_processing_order"]] += 1 - id2index_find_query[query["id"]] = i - - # Sanity check: all the stages should have the same number of queries - if len(stage2num_queries) == 0: - num_queries_per_stage = 0 - else: - num_queries_per_stage = stage2num_queries.most_common(1)[0][1] - for stage, num_queries in stage2num_queries.items(): - assert ( - num_queries == num_queries_per_stage - ), f"Number of queries in stage {stage} is {num_queries}, expected {num_queries_per_stage}" - - for query_id, query in enumerate(queries): - h, w = id2imsize[query["image_id"]] - if ( - "input_box" in query - and query["input_box"] is not None - and len(query["input_box"]) > 0 - ): - bbox = box_xywh_to_xyxy(torch.as_tensor(query["input_box"])).view(-1, 4) - bbox[:, 0::2].mul_(w).clamp_(min=0, max=w) - bbox[:, 1::2].mul_(h).clamp_(min=0, max=h) - if "input_box_label" in query and query["input_box_label"] is not None: - bbox_label = torch.as_tensor( - query["input_box_label"], dtype=torch.long - ).view(-1) - assert len(bbox_label) == len(bbox) - else: - # assume the boxes are positives - bbox_label = torch.ones(len(bbox), dtype=torch.long) - else: - bbox = None - bbox_label = None - - if "input_points" in query and query["input_points"] is not None: - points = torch.as_tensor(query["input_points"]).view(1, -1, 3) - points[:, :, 0:1].mul_(w).clamp_(min=0, max=w) - points[:, :, 1:2].mul_(h).clamp_(min=0, max=h) - else: - points = None - - try: - original_image_id = int( - img_metadata[id2index_img[query["image_id"]]]["original_img_id"] - ) - except ValueError: - original_image_id = -1 - - try: - img_metadata_query = img_metadata[id2index_img[query["image_id"]]] - coco_image_id = ( - int(img_metadata_query["coco_img_id"]) - if "coco_img_id" in img_metadata_query - else query["id"] - ) - except KeyError: - coco_image_id = -1 - - try: - original_category_id = int(query["original_cat_id"]) - except (ValueError, KeyError): - original_category_id = -1 - - # For evaluation, we associate the ids of the object to be tracked to the query - if query["object_ids_output"]: - obj_id = query["object_ids_output"][0] - obj_idx = id2index_obj[obj_id] - image_idx = id2index_img[query["image_id"]] - object_id = images[image_idx].objects[obj_idx].object_id - frame_index = images[image_idx].objects[obj_idx].frame_index - else: - object_id = -1 - frame_index = -1 - - find_queries.append( - FindQueryLoaded( - # id=query["id"], - # query_type=qtype, - query_text=( - query["query_text"] if query["query_text"] is not None else "" - ), - image_id=id2index_img[query["image_id"]], - input_bbox=bbox, - input_bbox_label=bbox_label, - input_points=points, - object_ids_output=[ - id2index_obj[obj_id] for obj_id in query["object_ids_output"] - ], - is_exhaustive=query["is_exhaustive"], - is_pixel_exhaustive=( - query["is_pixel_exhaustive"] - if "is_pixel_exhaustive" in query - else ( - query["is_exhaustive"] if query["is_exhaustive"] else None - ) - ), - query_processing_order=query["query_processing_order"], - inference_metadata=InferenceMetadata( - coco_image_id=-1 if self.training else coco_image_id, - original_image_id=(-1 if self.training else original_image_id), - frame_index=frame_index, - original_category_id=original_category_id, - original_size=(h, w), - object_id=object_id, - ), - ) - ) - - return Datapoint( - find_queries=find_queries, - images=images, - raw_images=[p[1] for p in pil_images], - ) - - def __len__(self) -> int: - return len(self.ids) - - -class Sam3ImageDataset(CustomCocoDetectionAPI): - def __init__( - self, - img_folder, - ann_file, - transforms, - max_ann_per_img: int, - multiplier: int, - training: bool, - load_segmentation: bool = False, - max_train_queries: int = 81, - max_val_queries: int = 300, - fix_fname: bool = False, - is_sharded_annotation_dir: bool = False, - blurring_masks_path: Optional[str] = None, - use_caching: bool = True, - zstd_dict_path=None, - filter_query=None, - coco_json_loader: Callable = COCO_FROM_JSON, - limit_ids: int = None, - ): - super(Sam3ImageDataset, self).__init__( - img_folder, - ann_file, - fix_fname=fix_fname, - load_segmentation=load_segmentation, - training=training, - blurring_masks_path=blurring_masks_path, - use_caching=use_caching, - zstd_dict_path=zstd_dict_path, - filter_query=filter_query, - coco_json_loader=coco_json_loader, - limit_ids=limit_ids, - ) - - self._transforms = transforms - self.training = training - self.max_ann_per_img = max_ann_per_img - self.max_train_queries = max_train_queries - self.max_val_queries = max_val_queries - - self.repeat_factors = torch.ones(len(self.ids), dtype=torch.float32) - - self.repeat_factors *= multiplier - print(f"Raw dataset length = {len(self.ids)}") - - self._MAX_RETRIES = 100 - - def __getitem__(self, idx): - return self.__orig_getitem__(idx) - - def __orig_getitem__(self, idx): - for _ in range(self._MAX_RETRIES): - try: - datapoint = super(Sam3ImageDataset, self).__getitem__(idx) - - # This can be done better by filtering the offending find queries - # However, this requires care: - # - Delete any find/get query that may depend on the deleted one - # - Re-compute the indexes in the pointers to account for the deleted finds - for q in datapoint.find_queries: - if len(q.object_ids_output) > self.max_ann_per_img: - raise DecompressionBombError( - f"Too many outputs ({len(q.object_ids_output)})" - ) - - max_queries = ( - self.max_train_queries if self.training else self.max_val_queries - ) - - if len(datapoint.find_queries) > max_queries: - raise DecompressionBombError( - f"Too many find queries ({len(datapoint.find_queries)})" - ) - - if len(datapoint.find_queries) == 0: - raise DecompressionBombError("No find queries") - for transform in self._transforms: - datapoint = transform(datapoint, epoch=self.curr_epoch) - - break - except (DecompressionBombError, OSError, ValueError) as error: - sys.stderr.write(f"ERROR: got loading error on datapoint {idx}\n") - sys.stderr.write(f"Exception: {error}\n") - sys.stderr.write(traceback.format_exc()) - idx = (idx + 1) % len(self) - else: - raise RuntimeError( - f"Failed {self._MAX_RETRIES} times trying to load an image." - ) - - return datapoint diff --git a/train/data/sam3_video_dataset.py b/train/data/sam3_video_dataset.py deleted file mode 100644 index 8ee8a4f..0000000 --- a/train/data/sam3_video_dataset.py +++ /dev/null @@ -1,327 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved - -import copy - -import io -import json -import logging -import math -import os -import pickle -import random -import sys -from typing import Any, Dict, List, Optional, Set, Tuple, Union - -import torch -import torchvision - -# from decord import cpu, VideoReader - -from iopath.common.file_io import PathManager -from PIL import Image as PILImage - -from .sam3_image_dataset import Datapoint, Sam3ImageDataset - - -SEED = 42 - - -class VideoGroundingDataset(Sam3ImageDataset): - def __init__( - self, - num_stages_sample: int = 4, - stage_stride_min: int = 1, - stage_stride_max: int = 5, - random_reverse_time_axis: bool = True, - is_tiling_single_image: bool = False, - # By default, we remove find those queries with geometric inputs (input_box or input_points) - # when creating synthetic videos from frames (since they are not *video-level* text prompts). - # If we need them later, we can sample them on-the-fly via transforms or inside the model. - tile_img_keep_find_queries_with_geo_inputs: bool = False, - tile_img_keep_get_queries: bool = False, - # the maximum number of find queries (for each frame) to keep in a video; if the datapoint - # contains more queries per frame than this limit, we subsample them to avoid OOM errors - max_query_num: int = -1, # the default -1 means no limit - # whether to override the "is_exhaustive" flag of the loaded find queries to True - # (by default, our video datasets are ingested with is_exhaustive=False, since the YTVIS format - # annotations doesn't involve an "is_exhaustive" flag; this means that those unmatched (negative) - # detection queries or tracking queries do not receive a classification loss given that we have - # weak_loss=True in IABCEMdetr -- this could lead to false positives for both image detection - # and video association.) - override_query_is_exhaustive_to_true: bool = False, - # the maximum number of masklets in a video; if the datapoint contains more masklets - # than this limit, we skip the datapoint to avoid OOM errors (this is useful for - # training with large videos that contain many objects) - max_masklet_num_in_video: int = 300, # 300 masklets is usually OK to avoid OOM - **kwargs, - ): - """ - Loading video grounding data - - Video frame sampling parameters (for training only): - - num_stages_sample: number of frames to sample from the video during training - - stage_stride_min: minimum stride between sampled frames during training - - stage_stride_max: maximum stride between sampled frames during training (if it's - greater than stage_stride_min, the actual stride is sampled uniformly between min - and max; during inference, we always use all frames in the video with stride=1) - - random_reverse_time_axis: whether to randomly invert the video's temporal axis - (i.e. playing it backwards) during training - """ - super().__init__(**kwargs) - assert num_stages_sample >= 1 - assert stage_stride_min >= 1 - assert stage_stride_max >= stage_stride_min - self.num_stages_sample = num_stages_sample - self.stage_stride_min = stage_stride_min - self.stage_stride_max = stage_stride_max - self.random_reverse_time_axis = random_reverse_time_axis - self.is_tiling_single_image = is_tiling_single_image - self.tile_img_keep_find_queries_with_geo_inputs = ( - tile_img_keep_find_queries_with_geo_inputs - ) - self.tile_img_keep_get_queries = tile_img_keep_get_queries - self.max_query_num = max_query_num - self.override_query_is_exhaustive_to_true = override_query_is_exhaustive_to_true - self.max_masklet_num_in_video = max_masklet_num_in_video - self.rng = random.Random() - self.set_curr_epoch(0) - - def set_curr_epoch(self, epoch: int): - super().set_curr_epoch(epoch) - self.rng.seed(SEED + epoch) - - def _load_datapoint(self, index: int) -> Datapoint: - id = self.ids[index].item() - queries, annotations = self.coco.loadQueriesAndAnnotationsFromDatapoint(id) - - # we subsample the video frames during training - if self.training and not self.is_tiling_single_image: - # pick a random stride for sampling query stages (`randint` includes both ends) - stage_stride = self.rng.randint( - self.stage_stride_min, self.stage_stride_max - ) - stage_ids_to_keep = self._sample_stage_ids( - queries, self.num_stages_sample, stage_stride - ) - # filter the queries and annotations to keep only the selected stages - # (also remap the stage ids so that they are contiguous and start from 0) - reverse_time_axis = ( - self.rng.random() < 0.5 if self.random_reverse_time_axis else False - ) - queries, annotations, kept_img_ids = self._filter_query_and_anns( - queries, - annotations, - stage_ids_to_keep, - remap_stage_id=True, - reverse_time_axis=reverse_time_axis, - ) - pil_images, img_metadata = self._load_images(id, kept_img_ids) - if reverse_time_axis: - # reverse the temporal ordering of the images and their metadata - # so that the image order matches the query order - pil_images = pil_images[::-1] - img_metadata = img_metadata[::-1] - else: - pil_images, img_metadata = self._load_images(id) - - # check that all the images have the same image size (they are expected - # to have the same image size since they are frames from the same video) - assert all(p.size == pil_images[0][1].size for _, p in pil_images) - - queries.sort(key=lambda q: q["query_processing_order"]) - if self.override_query_is_exhaustive_to_true: - for query in queries: - query["is_exhaustive"] = True - datapoint = self.load_queries(pil_images, annotations, queries, img_metadata) - - # skip datapoints with too many masklets to avoid OOM errors - num_masklets_in_video = len(datapoint.images[0].objects) - if num_masklets_in_video > self.max_masklet_num_in_video > 0: - logging.warning( - f"Datapoint {id} has ({num_masklets_in_video=}), exceeding " - f"the maximum allowed ({self.max_masklet_num_in_video}). " - "Skipping this datapoint." - ) - next_index = (index + 1) % len(self) - return self._load_datapoint(next_index) # move to the next datapoint - - if self.is_tiling_single_image: - datapoint = self._tile_single_image_data(datapoint, self.num_stages_sample) - if self.max_query_num > 0: - datapoint = self._subsample_queries(datapoint, self.max_query_num) - - # ensure that all find queries have the same processing order as their image id - for query in datapoint.find_queries: - assert query.image_id == query.query_processing_order, ( - f"find query has inconsistent image_id and " - f"query_processing_order: {query.image_id=} vs " - f"{query.query_processing_order=}" - ) - return datapoint - - def _sample_stage_ids(self, queries, num_stages_sample, stage_stride): - """Sample a subset of stage ids from all queries.""" - # Later we can perhaps turn it into a Sampler class to be more flexible. - all_stage_ids = sorted(set(q["query_processing_order"] for q in queries)) - num_stages_total = len(all_stage_ids) - if num_stages_total < num_stages_sample: - raise ValueError("Not enough stages to sample") - - # the difference in index between the first and the last sampled stage ids - b_e_gap = (num_stages_sample - 1) * stage_stride - if b_e_gap > num_stages_total - 1: - # In this case, it's not possible to sample with the provide stride, - # so we use the maximum possible stride. - prev_stage_stride = stage_stride - stage_stride = math.floor((num_stages_total - 1) / (num_stages_sample - 1)) - logging.info( - f"lowering stride from {prev_stage_stride} to {stage_stride} to " - f"sample {num_stages_sample} stages (from {num_stages_total} total)" - ) - b_e_gap = (num_stages_sample - 1) * stage_stride - - # randomly select a starting stage id (`randint` includes both ends) - b_max = len(all_stage_ids) - 1 - b_e_gap - b = self.rng.randint(0, b_max) - e = b + b_e_gap - stage_ids_to_keep = all_stage_ids[b : e + 1 : stage_stride] - return stage_ids_to_keep - - def _filter_query_and_anns( - self, queries, annotations, stage_ids_to_keep, remap_stage_id, reverse_time_axis - ): - """Filter queries and annotations to only keep those in `stage_ids_to_keep`.""" - stage_ids_to_keep = set(stage_ids_to_keep) - kept_img_ids = set() - kept_stage_ids = set() - - # Filter queries -- keep those queries with stage_id in `stage_ids_to_keep` - filtered_queries = [] - for query in queries: - input_box = query.get("input_box", None) - input_points = query.get("input_points", None) - has_geo_input = input_box is not None or input_points is not None - if has_geo_input and not self.tile_img_keep_find_queries_with_geo_inputs: - continue - stage_id = query["query_processing_order"] - if stage_id in stage_ids_to_keep: - kept_img_ids.add(query["image_id"]) - kept_stage_ids.add(stage_id) - filtered_queries.append(query) - # Check that all frames in `stage_ids_to_keep` are present after filtering - all_frame_present = kept_stage_ids == stage_ids_to_keep - assert all_frame_present, f"{kept_stage_ids=} vs {stage_ids_to_keep=}" - if remap_stage_id: - # Remap those kept stage ids to be contiguous and starting from 0 - old_stage_ids = sorted(kept_stage_ids, reverse=reverse_time_axis) - stage_id_old2new = {old: new for new, old in enumerate(old_stage_ids)} - for query in filtered_queries: - ptr_x_is_empty = query["ptr_x_query_id"] in [None, -1] - ptr_y_is_empty = query["ptr_y_query_id"] in [None, -1] - assert ( - ptr_x_is_empty and ptr_y_is_empty - ), "Remapping stage ids is not supported for queries with non-empty ptr_x or ptr_y pointers" - query["query_processing_order"] = stage_id_old2new[ - query["query_processing_order"] - ] - - # Filter annotations -- keep those annotations with image_id in `kept_img_ids` - filtered_annotations = [ - ann for ann in annotations if ann["image_id"] in kept_img_ids - ] - - return filtered_queries, filtered_annotations, kept_img_ids - - def _tile_single_image_data(self, datapoint: Datapoint, num_stages_sample: int): - """ - Tile a single image and its queries to simulate video frames. The output is a - datapoint with *identical video frames* (i.e. the same static image) and needs - further transforms (e.g. affine) to get video frames with different content. - """ - # tile `images: List[Image]` - assert len(datapoint.images) == 1, "Expected only one single image" - tiled_images = [ - copy.deepcopy(datapoint.images[0]) for _ in range(num_stages_sample) - ] - for stage_id, img in enumerate(tiled_images): - for obj in img.objects: - obj.frame_index = stage_id - - # tile `raw_images: Optional[List[PILImage.Image]] = None` - tiled_raw_images = None - if datapoint.raw_images is not None: - assert len(datapoint.raw_images) == 1, "Expected only one single image" - tiled_raw_images = [ - datapoint.raw_images[0].copy() for _ in range(num_stages_sample) - ] - - # tile `find_queries: List[FindQueryLoaded]` - tiled_find_queries_per_stage = [[] for _ in range(num_stages_sample)] - for query in datapoint.find_queries: - assert query.image_id == 0 - assert query.query_processing_order == 0 - # check and make sure that a query doesn't contain pointers or references - # to other queries (that cannot be tiled) - assert query.ptr_x is None and query.ptr_y is None - assert query.ptr_mem is None - # assert query.wkdata_qid is None - # assert query.other_positive_qids is None - # assert query.negative_qids is None - has_geo_input = ( - query.input_bbox is not None or query.input_points is not None - ) - if has_geo_input and not self.tile_img_keep_find_queries_with_geo_inputs: - continue - for stage_id in range(num_stages_sample): - # copy the query and update the image_id - new_query = copy.deepcopy(query) - new_query.image_id = stage_id - new_query.query_processing_order = stage_id - if new_query.inference_metadata is not None: - new_query.inference_metadata.frame_index = stage_id - tiled_find_queries_per_stage[stage_id].append(new_query) - - tiled_find_queries = sum(tiled_find_queries_per_stage, []) - - # tile `get_queries: List[GetQuery]` -- we skip them for now (since they involve - # a pointer to a find query that is complicated to tile, and there is not an - # imminent use case for them in the video grounding task in the near future) - if self.tile_img_keep_get_queries: - raise NotImplementedError("Tiling get queries is not implemented yet") - else: - tiled_get_queries = [] - - return Datapoint( - images=tiled_images, - raw_images=tiled_raw_images, - find_queries=tiled_find_queries, - get_queries=tiled_get_queries, - ) - - def _subsample_queries(self, datapoint: Datapoint, max_query_num: int): - """Subsample to keep at most `max_query_num` queries per frame in a datapoint.""" - # aggregate the find queries per stage - num_frames = max(q.query_processing_order for q in datapoint.find_queries) + 1 - find_queries_per_stage = [[] for _ in range(num_frames)] - for query in datapoint.find_queries: - find_queries_per_stage[query.query_processing_order].append(query) - - # verify that all the stages have the same number of queries - num_queries_per_stage = len(find_queries_per_stage[0]) - for queries in find_queries_per_stage: - assert len(queries) == num_queries_per_stage - if max_query_num <= 0 or num_queries_per_stage <= max_query_num: - return datapoint - - # subsample the queries to keep only `max_query_num` queries - sampled_inds = self.rng.sample(range(num_queries_per_stage), max_query_num) - sampled_find_queries_per_stage = [ - [queries[idx] for idx in sampled_inds] for queries in find_queries_per_stage - ] - sampled_find_queries = sum(sampled_find_queries_per_stage, []) - return Datapoint( - images=datapoint.images, - raw_images=datapoint.raw_images, - find_queries=sampled_find_queries, - get_queries=datapoint.get_queries, - ) diff --git a/train/data/torch_dataset.py b/train/data/torch_dataset.py deleted file mode 100644 index 0555617..0000000 --- a/train/data/torch_dataset.py +++ /dev/null @@ -1,52 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved - -from typing import Callable, Iterable, Optional - -from torch.utils.data import DataLoader, Dataset, DistributedSampler, IterableDataset - - -class TorchDataset: - def __init__( - self, - dataset: Dataset, - batch_size: int, - num_workers: int, - shuffle: bool, - pin_memory: bool, - drop_last: bool, - collate_fn: Optional[Callable] = None, - worker_init_fn: Optional[Callable] = None, - enable_distributed_sampler=True, - ) -> None: - self.dataset = dataset - self.batch_size = batch_size - self.num_workers = num_workers - self.shuffle = shuffle - self.pin_memory = pin_memory - self.drop_last = drop_last - self.collate_fn = collate_fn - self.worker_init_fn = worker_init_fn - assert not isinstance(self.dataset, IterableDataset), "Not supported yet" - if enable_distributed_sampler: - self.sampler = DistributedSampler(self.dataset, shuffle=self.shuffle) - else: - self.sampler = None - - def get_loader(self, epoch) -> Iterable: - if self.sampler: - self.sampler.set_epoch(epoch) - if hasattr(self.dataset, "epoch"): - self.dataset.epoch = epoch - if hasattr(self.dataset, "set_epoch"): - self.dataset.set_epoch(epoch) - - return DataLoader( - self.dataset, - batch_size=self.batch_size, - num_workers=self.num_workers, - pin_memory=self.pin_memory, - drop_last=self.drop_last, - sampler=self.sampler, - collate_fn=self.collate_fn, - worker_init_fn=self.worker_init_fn, - ) diff --git a/train/loss/__init__.py b/train/loss/__init__.py deleted file mode 100644 index 46d37d2..0000000 --- a/train/loss/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved diff --git a/train/loss/loss_fns.py b/train/loss/loss_fns.py deleted file mode 100644 index 3b61d7f..0000000 --- a/train/loss/loss_fns.py +++ /dev/null @@ -1,1319 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved - -import logging -import warnings - -import torch -import torch.distributed -import torch.nn.functional as F -import torchmetrics - -from sam3.model import box_ops - -from sam3.model.data_misc import interpolate - -from sam3.train.loss.sigmoid_focal_loss import ( - triton_sigmoid_focal_loss, - triton_sigmoid_focal_loss_reduce, -) -from torch import nn - -from .mask_sampling import ( - calculate_uncertainty, - get_uncertain_point_coords_with_randomness, - point_sample, -) - - -CORE_LOSS_KEY = "core_loss" - - -def instance_masks_to_semantic_masks( - instance_masks: torch.Tensor, num_instances: torch.Tensor -) -> torch.Tensor: - """This function converts instance masks to semantic masks. - It accepts a collapsed batch of instances masks (ie all instance masks are concatenated in a single tensor) and - the number of instances in each image of the batch. - It returns a mask with the same spatial dimensions as the input instance masks, where for each batch element the - semantic mask is the union of all the instance masks in the batch element. - - If for a given batch element there are no instances (ie num_instances[i]==0), the corresponding semantic mask will be a tensor of zeros. - - Args: - instance_masks (torch.Tensor): A tensor of shape (N, H, W) where N is the number of instances in the batch. - num_instances (torch.Tensor): A tensor of shape (B,) where B is the batch size. It contains the number of instances - in each image of the batch. - - Returns: - torch.Tensor: A tensor of shape (B, H, W) where B is the batch size and H, W are the spatial dimensions of the - input instance masks. - """ - if num_instances.sum() == 0: - # all negative batch, create a tensor of zeros (B, 1, 1) - return num_instances.unsqueeze(-1).unsqueeze(-1) - - masks_per_query = torch.split(instance_masks, num_instances.tolist()) - - return torch.stack([torch.any(masks, dim=0) for masks in masks_per_query], dim=0) - - -@torch.no_grad() -def accuracy(output, target, topk=(1,)): - """Computes the precision@k for the specified values of k""" - if target.numel() == 0: - return [torch.zeros([], device=output.device)] - maxk = max(topk) - batch_size = target.size(0) - - _, pred = output.topk(maxk, 1, True, True) - pred = pred.t() - correct = pred.eq(target.view(1, -1).expand_as(pred)) - - res = [] - for k in topk: - correct_k = correct[:k].view(-1).float().sum(0) - res.append(correct_k.mul_(100.0 / batch_size)) - return res - - -def dice_loss(inputs, targets, num_boxes, loss_on_multimask=False, reduce=True): - """ - Compute the DICE loss, similar to generalized IOU for masks - Args: - inputs: A float tensor of arbitrary shape. - The predictions for each example. - targets: A float tensor with the same shape as inputs. Stores the binary - classification label for each element in inputs - (0 for the negative class and 1 for the positive class). - """ - try: - loss = _dice_loss(inputs, targets, num_boxes, loss_on_multimask, reduce) - except torch.OutOfMemoryError: - logging.error("GPU OOM, computing dice loss on CPU") - # try to recover from GPU OOM by moving tensors to CPU and computing loss there - orig_device = inputs.device - inputs = inputs.cpu() - targets = targets.cpu() - if isinstance(num_boxes, torch.Tensor): - num_boxes = num_boxes.cpu() - loss = _dice_loss(inputs, targets, num_boxes, loss_on_multimask, reduce) - loss = loss.to(orig_device) - - return loss - - -def _dice_loss(inputs, targets, num_boxes, loss_on_multimask=False, reduce=True): - inputs = inputs.sigmoid() - if loss_on_multimask: - # inputs and targets are [N, M, H, W] where M corresponds to multiple predicted masks - assert inputs.dim() == 4 and targets.dim() == 4 - # flatten spatial dimension while keeping multimask channel dimension - inputs = inputs.flatten(2) - targets = targets.flatten(2) - numerator = 2 * (inputs * targets).sum(-1) - else: - inputs = inputs.flatten(1) - numerator = 2 * (inputs * targets).sum(1) - denominator = inputs.sum(-1) + targets.sum(-1) - loss = 1 - (numerator + 1) / (denominator + 1) - if loss_on_multimask: - return loss / num_boxes - if not reduce: - return loss - return loss.sum() / num_boxes - - -def sigmoid_focal_loss( - inputs, - targets, - num_boxes, - alpha: float = 0.25, - gamma: float = 2, - loss_on_multimask=False, - reduce=True, - triton=True, -): - """ - Loss used in RetinaNet for dense detection: https://arxiv.org/abs/1708.02002. - Args: - inputs: A float tensor of arbitrary shape. - The predictions for each example. - targets: A float tensor with the same shape as inputs. Stores the binary - classification label for each element in inputs - (0 for the negative class and 1 for the positive class). - alpha: (optional) Weighting factor in range (0,1) to balance - positive vs negative examples. Default = -1 (no weighting). - gamma: Exponent of the modulating factor (1 - p_t) to - balance easy vs hard examples. - Returns: - Loss tensor - """ - if not (0 <= alpha <= 1) and triton: - raise RuntimeError(f"Alpha should be in [0,1], got {alpha}") - if triton: - if reduce and not loss_on_multimask: - loss = triton_sigmoid_focal_loss_reduce(inputs, targets, alpha, gamma) - return loss / (num_boxes * inputs.shape[1]) - - loss = triton_sigmoid_focal_loss(inputs, targets, alpha, gamma) - else: - prob = inputs.sigmoid() - ce_loss = F.binary_cross_entropy_with_logits(inputs, targets, reduction="none") - p_t = prob * targets + (1 - prob) * (1 - targets) - loss = ce_loss * ((1 - p_t) ** gamma) - - if alpha >= 0: - alpha_t = alpha * targets + (1 - alpha) * (1 - targets) - loss = alpha_t * loss - - if not reduce: - return loss - - if loss_on_multimask: - # loss is [N, M, H, W] where M corresponds to multiple predicted masks - assert loss.dim() == 4 - return loss.flatten(2).mean(-1) / num_boxes # average over spatial dims - return loss.mean(1).sum() / num_boxes - - -def iou_loss( - inputs, targets, pred_ious, num_boxes, loss_on_multimask=False, use_l1_loss=False -): - """MSE loss between predicted IoUs and actual IoUs between inputs and targets.""" - assert inputs.dim() == 4 and targets.dim() == 4 - pred_mask = inputs.flatten(2) > 0 - gt_mask = targets.flatten(2) > 0 - area_i = torch.sum(pred_mask & gt_mask, dim=-1).float() - area_u = torch.sum(pred_mask | gt_mask, dim=-1).float() - actual_ious = area_i / torch.clamp(area_u, min=1.0) - - if use_l1_loss: - loss = F.l1_loss(pred_ious, actual_ious, reduction="none") - else: - loss = F.mse_loss(pred_ious, actual_ious, reduction="none") - if loss_on_multimask: - return loss / num_boxes - return loss.sum() / num_boxes - - -@torch.jit.script -def _contrastive_align(logits, positive_map): - positive_logits = -logits.masked_fill(~positive_map, 0) - negative_logits = logits # .masked_fill(positive_map, -1000000) - - boxes_with_pos = positive_map.any(2) - pos_term = positive_logits.sum(2) - neg_term = negative_logits.logsumexp(2) - - nb_pos = positive_map.sum(2) + 1e-6 - - box_to_token_loss = ( - (pos_term / nb_pos + neg_term).masked_fill(~boxes_with_pos, 0).sum() - ) - - tokens_with_pos = positive_map.any(1) - pos_term = positive_logits.sum(1) - neg_term = negative_logits.logsumexp(1) - - nb_pos = positive_map.sum(1) + 1e-6 - - tokens_to_boxes_loss = ( - (pos_term / nb_pos + neg_term).masked_fill(~tokens_with_pos, 0).sum() - ) - return (box_to_token_loss + tokens_to_boxes_loss) / 2 - - -def _get_src_permutation_idx(indices): - # permute predictions following indices - batch_idx = torch.cat( - [torch.full_like(src, i) for i, (src, _) in enumerate(indices)] - ) - src_idx = torch.cat([src for (src, _) in indices]) - return batch_idx, src_idx - - -class LossWithWeights(nn.Module): - def __init__(self, weight_dict, compute_aux, supports_o2m_loss=True): - super().__init__() - # weights for each computed loss key (those losses not in weight_dict - # will not be aggregated in the final reduced core loss) - self.weight_dict = weight_dict if weight_dict is not None else {} - # whether this loss will be applied on auxiliary outputs - self.compute_aux = compute_aux - self.supports_o2m_loss = supports_o2m_loss - self.target_keys = [] - - def forward(self, *args, is_aux=False, **kwargs): - if is_aux and not self.compute_aux: - return {CORE_LOSS_KEY: 0.0} - losses = self.get_loss(*args, **kwargs) - losses[CORE_LOSS_KEY] = self.reduce_loss(losses) - return losses - - def get_loss(self, **kwargs): - raise NotImplementedError() - - def reduce_loss(self, losses): - reduced_loss = 0.0 - for loss_key, weight in self.weight_dict.items(): - if loss_key not in losses: - raise ValueError(f"{type(self)} doesn't compute {loss_key}") - if weight != 0: - reduced_loss += losses[loss_key] * weight - - return reduced_loss - - -class IABCEMdetr(LossWithWeights): - def __init__( - self, - pos_weight, - weight_dict=None, - compute_aux=True, - gamma=0, - weak_loss=True, - alpha=0.25, - pad_n_queries=None, - pad_scale_pos=1.0, - use_separate_loss_for_det_and_trk=False, - num_det_queries=None, - det_exhaustive_loss_scale_pos=1.0, - det_exhaustive_loss_scale_neg=1.0, - det_non_exhaustive_loss_scale_pos=1.0, - det_non_exhaustive_loss_scale_neg=1.0, - trk_loss_scale_pos=1.0, - trk_loss_scale_neg=1.0, - no_loss_for_fp_propagation=False, - apply_loss_to_det_queries_in_video_grounding=True, - use_presence=False, - use_presence_semgseg=False, # If True, use presence scores from the semgseg head. - presence_alpha=0.5, - presence_gamma=0.0, - pos_focal: bool = False, # for box scores, use focal loss for positives as well - ): - super().__init__(weight_dict, compute_aux) - self.pos_weight = pos_weight - self.gamma = gamma - self.weak_loss = weak_loss - self.alpha = alpha - self.target_keys.append("boxes_xyxy") - self.no_loss_for_fp_propagation = no_loss_for_fp_propagation - if self.weak_loss: - self.target_keys.append("is_exhaustive") - # NOTE: This is hacky solution to have the same CE loss scale across datasets where the model might predict different number of object queries for different tasks. - # If not None, we assume there are a total pad_n_queries object queries. - # For example, if the model predicts only 1 object query and pad_n_queries=100, we pad the predictions with 99 zero preds. - # Currently this only affects the BCE loss and not the F1 score. - self.pad_n_queries = pad_n_queries - self.pad_scale_pos = pad_scale_pos - if self.pad_scale_pos != 1.0: - assert self.pad_n_queries is not None - # whether to use presence scores - self.use_presence = use_presence - self.use_presence_semgseg = use_presence_semgseg - if self.use_presence_semgseg: - assert self.use_presence - self.presence_alpha = presence_alpha - self.presence_gamma = presence_gamma - self.pos_focal = pos_focal - - # Decoupled loss for detection and tracking queries - self.apply_loss_to_det_queries_in_video_grounding = ( - apply_loss_to_det_queries_in_video_grounding - ) - self.use_separate_loss_for_det_and_trk = use_separate_loss_for_det_and_trk - if num_det_queries is not None: - logging.warning("note: it's not needed to set num_det_queries anymore") - if self.use_separate_loss_for_det_and_trk: - assert not self.weak_loss, "Do not use weak_loss in this case -- set separate loss for detection and tracking queries instead" - self.det_exhaustive_loss_scale_pos = det_exhaustive_loss_scale_pos - self.det_exhaustive_loss_scale_neg = det_exhaustive_loss_scale_neg - self.det_non_exhaustive_loss_scale_pos = det_non_exhaustive_loss_scale_pos - self.det_non_exhaustive_loss_scale_neg = det_non_exhaustive_loss_scale_neg - self.trk_loss_scale_pos = trk_loss_scale_pos - self.trk_loss_scale_neg = trk_loss_scale_neg - else: - assert ( - det_exhaustive_loss_scale_pos == 1.0 - and det_exhaustive_loss_scale_neg == 1.0 - and det_non_exhaustive_loss_scale_pos == 1.0 - and det_non_exhaustive_loss_scale_neg == 1.0 - and trk_loss_scale_pos == 1.0 - and trk_loss_scale_neg == 1.0 - ), "If not using separate loss for detection and tracking queries, separate detection and tracking loss scales should all be 1.0" - - def get_loss(self, outputs, targets, indices, num_boxes): - assert len(outputs["pred_logits"].shape) > 2, "Incorrect predicted logits shape" - assert outputs["pred_logits"].shape[-1] == 1, "Incorrect predicted logits shape" - src_logits = outputs["pred_logits"].squeeze(-1) - prob = src_logits.sigmoid() - - with torch.no_grad(): - target_classes = torch.full( - src_logits.shape[:2], - 0, - dtype=torch.float, - device=src_logits.device, - ) - target_classes[(indices[0], indices[1])] = 1 - src_boxes_xyxy = outputs["pred_boxes_xyxy"][(indices[0], indices[1])] - target_boxes_giou = ( - targets["boxes_xyxy"][indices[2]] - if indices[2] is not None - else targets["boxes_xyxy"] - ) - - iou = box_ops.fast_diag_box_iou(src_boxes_xyxy, target_boxes_giou) - t = prob[(indices[0], indices[1])] ** self.alpha * iou ** (1 - self.alpha) - t = torch.clamp(t, 0.01).detach() - positive_target_classes = target_classes.clone() - positive_target_classes[(indices[0], indices[1])] = t - - # Soft loss on positives - if self.pos_focal: - loss_bce = sigmoid_focal_loss( - src_logits.contiguous(), - positive_target_classes, - num_boxes=1, - alpha=0.5, - gamma=self.gamma, - reduce=False, - ) - else: - loss_bce = F.binary_cross_entropy_with_logits( - src_logits, positive_target_classes, reduction="none" - ) - loss_bce = loss_bce * target_classes * self.pos_weight - - if ( - self.pad_n_queries is not None - and isinstance(self.pad_n_queries, int) - and loss_bce.size(1) < self.pad_n_queries - ): - loss_bce = loss_bce * self.pad_scale_pos - # Negatives - loss_bce = loss_bce + F.binary_cross_entropy_with_logits( - src_logits, target_classes, reduction="none" - ) * (1 - target_classes) * (prob**self.gamma) - - # Optionally, not applying IABCEMdetr loss to detection queries in video. - is_video_grounding = outputs.get("is_video_grounding_batch", False) - if is_video_grounding and not self.apply_loss_to_det_queries_in_video_grounding: - Q_det = outputs["Q_det"] - loss_bce[:, :Q_det] *= 0.0 - presence_loss = torch.tensor(0.0, device=src_logits.device) - presence_dec_acc = torch.tensor(0.0, device=src_logits.device) - if self.use_presence: - # no classifiction loss for individual tokens if no target gt - # cannot directly use targets["num_boxes"] to check if some - # GT box exists as there may be dummy boxes for "invisible objects" - # in video grounding data - - gt_padded_object_ids = targets["object_ids_padded"] # (B, H) - gt_padded_boxes = targets["boxes_padded"] # (B, H, 4) shape, CxCyWH - gt_padded_is_visible = ( - (gt_padded_object_ids >= 0) - & (gt_padded_boxes[..., 2] > 0) # width > 0 - & (gt_padded_boxes[..., 3] > 0) # height > 0 - ) - keep_loss = (gt_padded_is_visible.sum(dim=-1)[..., None] != 0).float() - - loss_bce = loss_bce * keep_loss - - if self.use_presence_semgseg: - # no loss here, has it's own separate loss computation - assert "presence_logit_dec" not in outputs - elif "presence_logit_dec" in outputs: - presence_logits = outputs["presence_logit_dec"].view_as(keep_loss) - bs = presence_logits.shape[0] - presence_loss = sigmoid_focal_loss( - presence_logits, - keep_loss, - # not num_boxes, but we'll use it to normalize by bs - num_boxes=bs, - alpha=self.presence_alpha, - gamma=self.presence_gamma, - ) - pred = (presence_logits.sigmoid() > 0.5).float() - presence_dec_acc = (pred == keep_loss).float().mean() - else: - # for o2m, nothing to do - pass - - if self.weak_loss: - assert not self.use_separate_loss_for_det_and_trk, "Do not use weak_loss in this case -- set separate loss for detection and tracking queries instead" - - # nullify the negative loss for the non-exhaustive classes - assert loss_bce.shape[0] == targets["is_exhaustive"].shape[0] - assert targets["is_exhaustive"].ndim == 1 - - loss_mask = (~targets["is_exhaustive"]).view(-1, 1).expand_as(loss_bce) - # restrict the mask to the negative supervision - loss_mask = loss_mask & (target_classes < 0.5) - loss_mask = ~loss_mask - # Mask the loss - loss_bce = loss_bce * loss_mask.float() - # Average - loss_bce = loss_bce.sum() / (loss_mask.sum() + 1e-6) - else: - # apply separate loss weights to detection and tracking queries - if self.use_separate_loss_for_det_and_trk: - Q_det = outputs["Q_det"] - assert loss_bce.size(1) >= Q_det - is_positive = target_classes > 0.5 - is_positive_det = is_positive[:, :Q_det] - is_positive_trk = is_positive[:, Q_det:] - assert loss_bce.size(0) == targets["is_exhaustive"].size(0) - is_exhaustive = targets["is_exhaustive"].unsqueeze(1).bool() - loss_scales = torch.zeros_like(loss_bce) - # detection query loss weights - loss_scales[:, :Q_det] = ( - (is_exhaustive & is_positive_det).float() - * self.det_exhaustive_loss_scale_pos - + (is_exhaustive & ~is_positive_det).float() - * self.det_exhaustive_loss_scale_neg - + (~is_exhaustive & is_positive_det).float() - * self.det_non_exhaustive_loss_scale_pos - + (~is_exhaustive & ~is_positive_det).float() - * self.det_non_exhaustive_loss_scale_neg - ) - # tracking query weights - loss_scales[:, Q_det:] = ( - is_positive_trk.float() * self.trk_loss_scale_pos - + (~is_positive_trk).float() * self.trk_loss_scale_neg - ) - # apply the loss weights - - # if the id is -2 means it is a fp propagation , we don't apply the loss to them - if self.no_loss_for_fp_propagation: - is_original_queries = outputs["pred_old_obj_ids"] != -2 - loss_scales *= (is_exhaustive | is_original_queries).float() - - loss_bce = loss_bce * loss_scales - - if self.pad_n_queries is None or loss_bce.size(1) >= self.pad_n_queries: - loss_bce = loss_bce.mean() - else: - assert isinstance(self.pad_n_queries, int) - assert ( - loss_bce.size(1) < self.pad_n_queries - ), f"The number of predictions is more than the expected total after padding. Got {loss_bce.size(1)} predictions." - loss_bce = loss_bce.sum() / (self.pad_n_queries * loss_bce.size(0)) - - bce_f1 = torchmetrics.functional.f1_score( - src_logits.sigmoid().flatten(), - target=target_classes.flatten().long(), - task="binary", - ) - - losses = { - "loss_ce": loss_bce, - "ce_f1": bce_f1, - "presence_loss": presence_loss, - "presence_dec_acc": presence_dec_acc, - } - return losses - - -class Boxes(LossWithWeights): - def __init__( - self, - weight_dict=None, - compute_aux=True, - apply_loss_to_det_queries_in_video_grounding=True, - ): - super().__init__(weight_dict, compute_aux) - self.apply_loss_to_det_queries_in_video_grounding = ( - apply_loss_to_det_queries_in_video_grounding - ) - self.target_keys.extend(["boxes", "boxes_xyxy"]) - - def get_loss(self, outputs, targets, indices, num_boxes): - """Compute the losses related to the bounding boxes, the L1 regression loss and the GIoU loss - targets dicts must contain the key "boxes" containing a tensor of dim [nb_target_boxes, 4] - The target boxes are expected in format (center_x, center_y, h, w), normalized by the image size. - """ - # Optionally, not applying Boxes loss to detection queries in video. - is_video_grounding = outputs.get("is_video_grounding_batch", False) - if is_video_grounding and not self.apply_loss_to_det_queries_in_video_grounding: - indices = _keep_only_trk_queries_in_match_inds( - indices, Q_det=outputs["Q_det"] - ) - - assert "pred_boxes" in outputs - # idx = self._get_src_permutation_idx(indices) - src_boxes = outputs["pred_boxes"][(indices[0], indices[1])] - src_boxes_xyxy = outputs["pred_boxes_xyxy"][(indices[0], indices[1])] - target_boxes = ( - targets["boxes"] if indices[2] is None else targets["boxes"][indices[2]] - ) - target_boxes_giou = ( - targets["boxes_xyxy"] - if indices[2] is None - else targets["boxes_xyxy"][indices[2]] - ) - - loss_bbox = F.l1_loss(src_boxes, target_boxes, reduction="none") - - losses = {} - losses["loss_bbox"] = loss_bbox.sum() / num_boxes - - loss_giou = 1 - box_ops.fast_diag_generalized_box_iou( - src_boxes_xyxy, target_boxes_giou - ) - losses["loss_giou"] = loss_giou.sum() / num_boxes - return losses - - -class Masks(LossWithWeights): - def __init__( - self, - weight_dict=None, - compute_aux=False, - focal_alpha=0.25, - focal_gamma=2, - num_sample_points=None, - oversample_ratio=None, - importance_sample_ratio=None, - apply_loss_to_det_queries_in_video_grounding=True, - ): - super().__init__(weight_dict, compute_aux) - if compute_aux: - warnings.warn("Masks loss usually shouldn't be applied to aux outputs") - self.focal_alpha = focal_alpha - self.focal_gamma = focal_gamma - self.num_sample_points = num_sample_points - self.oversample_ratio = oversample_ratio - self.importance_sample_ratio = importance_sample_ratio - self.apply_loss_to_det_queries_in_video_grounding = ( - apply_loss_to_det_queries_in_video_grounding - ) - self.target_keys.extend(["masks", "is_valid_mask"]) - - def _sampled_loss(self, src_masks, target_masks, num_boxes): - assert len(src_masks.shape) == 3 and len(target_masks.shape) == 3 - src_masks = src_masks[:, None] - target_masks = target_masks[:, None] - with torch.no_grad(): - # Sample point_coords - point_coords = get_uncertain_point_coords_with_randomness( - src_masks, - calculate_uncertainty, - self.num_sample_points, - self.oversample_ratio, - self.importance_sample_ratio, - ) - - # get GT labels - sampled_target_masks = point_sample( - target_masks, - point_coords, - align_corners=False, - ).squeeze(1) - - sampled_src_masks = point_sample( - src_masks, - point_coords, - align_corners=False, - ).squeeze(1) - - losses = { - "loss_mask": sigmoid_focal_loss( - sampled_src_masks, - sampled_target_masks, - num_boxes, - alpha=self.focal_alpha, - gamma=self.focal_gamma, - ), - "loss_dice": dice_loss(sampled_src_masks, sampled_target_masks, num_boxes), - } - # Not needed for backward - del src_masks - del target_masks - - return losses - - def get_loss(self, outputs, targets, indices, num_boxes): - """Compute the losses related to the masks: the focal loss and the dice loss. - targets dicts must contain the key "masks" containing a tensor of dim [nb_target_boxes, h, w] - """ - assert "pred_masks" in outputs - assert "is_valid_mask" in targets - # Optionally, not applying Masks loss to detection queries in video. - is_video_grounding = outputs.get("is_video_grounding_batch", False) - if is_video_grounding and not self.apply_loss_to_det_queries_in_video_grounding: - indices = _keep_only_trk_queries_in_match_inds( - indices, Q_det=outputs["Q_det"] - ) - - src_masks = outputs["pred_masks"] - - # Dataset doesn't have segmentation masks - if targets["masks"] is None: - return { - "loss_mask": torch.tensor(0.0, device=src_masks.device), - "loss_dice": torch.tensor(0.0, device=src_masks.device), - } - - target_masks = ( - targets["masks"] if indices[2] is None else targets["masks"][indices[2]] - ) - target_masks = target_masks.to(src_masks) - keep = ( - targets["is_valid_mask"] - if indices[2] is None - else targets["is_valid_mask"][indices[2]] - ) - - src_masks = src_masks[(indices[0], indices[1])] - - # Remove invalid masks from loss - src_masks = src_masks[keep] - target_masks = target_masks[keep] - - if self.num_sample_points is not None: - # Compute loss on sampled points for the Mask - losses = self._sampled_loss(src_masks, target_masks, num_boxes) - - else: - # upsample predictions to the target size - if target_masks.shape[0] == 0 and src_masks.shape[0] == 0: - src_masks = src_masks.flatten(1) - target_masks = target_masks.reshape(src_masks.shape) - else: - if len(src_masks.shape) == 3: - src_masks = src_masks[:, None] - if src_masks.dtype == torch.bfloat16: - # Bilinear interpolation does not support bf16 - src_masks = src_masks.to(dtype=torch.float32) - src_masks = interpolate( - src_masks, - size=target_masks.shape[-2:], - mode="bilinear", - align_corners=False, - ) - src_masks = src_masks[:, 0].flatten(1) - target_masks = target_masks.flatten(1) - - losses = { - "loss_mask": sigmoid_focal_loss( - src_masks, - target_masks, - num_boxes, - alpha=self.focal_alpha, - gamma=self.focal_gamma, - ), - "loss_dice": dice_loss(src_masks, target_masks, num_boxes), - } - - return losses - - -# class MultiStepIteractiveMasks(LossWithWeights): -# def __init__( -# self, -# weight_dict=None, -# compute_aux=False, -# focal_alpha=0.25, -# focal_gamma=2, -# ): -# warnings.warn( -# "MultiStepIteractiveMasks is deprecated. Please use MultiStepMultiMasksAndIous", -# DeprecationWarning, -# ) -# super().__init__(weight_dict, compute_aux) -# self.focal_alpha = focal_alpha -# self.focal_gamma = focal_gamma -# self.target_keys.extend(["masks"]) - -# def get_loss(self, outputs, targets, indices, num_boxes): -# """Compute the losses related to the masks: the focal loss and the dice loss. -# targets dicts must contain the key "masks" containing a tensor of dim [nb_target_boxes, h, w] - -# Unlike `Masks`, here the "multistep_pred_masks" can have multiple channels, each -# corresponding to one iterative prediction step in SAM-style training. We treat each -# channel as a mask prediction and sum the loss across channels. -# """ -# src_masks = outputs["multistep_pred_masks"] -# target_masks = targets["masks"] -# assert src_masks.size(0) == target_masks.size(0) -# assert src_masks.dim() == 4 -# assert target_masks.dim() == 3 - -# # tile target_masks according to the number of -# # channels `src_masks`. -# num_steps = src_masks.size(1) -# target_masks = target_masks.unsqueeze(1).to(src_masks.dtype) -# if num_steps > 1: -# target_masks = target_masks.repeat(1, num_steps, 1, 1) - -# # resize `src_masks` to target mask resolution -# if src_masks.shape != target_masks.shape: -# src_masks = interpolate( -# src_masks, -# size=target_masks.shape[-2:], -# mode="bilinear", -# align_corners=False, -# ) -# assert src_masks.shape == target_masks.shape - -# # flatten the multiple steps in to the batch dimension -# src_masks = src_masks.flatten(0, 1).flatten(1) -# target_masks = target_masks.flatten(0, 1).flatten(1) -# losses = { -# "loss_mask": sigmoid_focal_loss( -# src_masks, -# target_masks, -# num_boxes, -# alpha=self.focal_alpha, -# gamma=self.focal_gamma, -# ), -# "loss_dice": dice_loss(src_masks, target_masks, num_boxes), -# } - -# return losses - - -# class MultiStepMultiMasksAndIous(LossWithWeights): -# def __init__( -# self, -# weight_dict=None, -# compute_aux=False, -# focal_alpha=0.25, -# focal_gamma=2, -# # if True, back-prop on all predicted ious -# # not just the one with lowest loss_combo -# supervise_all_iou=False, -# # Less slack vs MSE loss in [-1, 1] error range -# iou_use_l1_loss=False, -# # Settings for obj score prediction -# pred_obj_scores=False, -# focal_gamma_obj_score=0.0, -# focal_alpha_obj_score=-1, -# ): -# super().__init__(weight_dict, compute_aux) -# self.focal_alpha = focal_alpha -# self.focal_gamma = focal_gamma -# self.target_keys.extend(["masks"]) -# assert "loss_mask" in self.weight_dict -# assert "loss_dice" in self.weight_dict -# assert "loss_iou" in self.weight_dict -# if "loss_class" not in self.weight_dict: -# self.weight_dict["loss_class"] = 0.0 -# self.focal_alpha_obj_score = focal_alpha_obj_score -# self.focal_gamma_obj_score = focal_gamma_obj_score -# self.supervise_all_iou = supervise_all_iou -# self.iou_use_l1_loss = iou_use_l1_loss -# self.pred_obj_scores = pred_obj_scores - -# def get_loss(self, outputs, targets, indices, num_boxes): -# """ -# Compute the losses related to the masks: the focal loss and the dice loss. -# and also the MSE loss between predicted IoUs and actual IoUs. - -# Here "multistep_pred_multimasks_high_res" is a list of multimasks (tensors -# of shape [N, M, H, W], where M could be 1 or larger, corresponding to -# one or multiple predicted masks from a click. - -# We back-propagate focal, dice and iou losses only on the prediction channel -# with the lowest focal+dice loss between predicted mask and ground-truth. -# """ - -# target_masks = targets["masks"].unsqueeze(1).float() -# assert target_masks.dim() == 4 # [N, 1, H, W] -# src_masks_list = outputs["multistep_pred_multimasks_high_res"] -# ious_list = outputs["multistep_pred_ious"] -# object_score_logits_list = outputs["multistep_object_score_logits"] - -# assert len(src_masks_list) == len(ious_list) -# assert len(object_score_logits_list) == len(ious_list) - -# # Remove invalid masks from loss -# keep = targets["is_valid_mask"] -# target_masks = target_masks[keep] - -# # accumulate the loss over prediction steps -# losses = {"loss_mask": 0, "loss_dice": 0, "loss_iou": 0, "loss_class": 0} -# for src_masks, ious, object_score_logits in zip( -# src_masks_list, ious_list, object_score_logits_list -# ): -# object_score_logits = object_score_logits[keep] -# ious = ious[keep] -# src_masks = src_masks[keep] -# self._update_losses( -# losses, src_masks, target_masks, ious, num_boxes, object_score_logits -# ) -# return losses - -# def _update_losses( -# self, losses, src_masks, target_masks, ious, num_boxes, object_score_logits -# ): -# target_masks = target_masks.expand_as(src_masks) -# # get focal, dice and iou loss on all output masks in a prediction step -# loss_multimask = sigmoid_focal_loss( -# src_masks, -# target_masks, -# num_boxes, -# alpha=self.focal_alpha, -# gamma=self.focal_gamma, -# loss_on_multimask=True, -# triton=False, # only use triton if alpha > 0 -# ) -# loss_multidice = dice_loss( -# src_masks, target_masks, num_boxes, loss_on_multimask=True -# ) -# if not self.pred_obj_scores: -# loss_class = torch.tensor( -# 0.0, dtype=loss_multimask.dtype, device=loss_multimask.device -# ) -# target_obj = torch.ones( -# loss_multimask.shape[0], -# 1, -# dtype=loss_multimask.dtype, -# device=loss_multimask.device, -# ) -# else: -# target_obj = torch.any((target_masks[:, 0] > 0).flatten(1), dim=-1)[ -# ..., None -# ].float() -# loss_class = sigmoid_focal_loss( -# object_score_logits, -# target_obj, -# num_boxes, -# alpha=self.focal_alpha_obj_score, -# gamma=self.focal_gamma_obj_score, -# triton=False, -# ) - -# loss_multiiou = iou_loss( -# src_masks, -# target_masks, -# ious, -# num_boxes, -# loss_on_multimask=True, -# use_l1_loss=self.iou_use_l1_loss, -# ) -# assert loss_multimask.dim() == 2 -# assert loss_multidice.dim() == 2 -# assert loss_multiiou.dim() == 2 -# if loss_multimask.size(1) > 1: -# # take the mask indices with the smallest focal + dice loss for back propagation -# loss_combo = ( -# loss_multimask * self.weight_dict["loss_mask"] -# + loss_multidice * self.weight_dict["loss_dice"] -# ) -# best_loss_inds = torch.argmin(loss_combo, dim=-1) -# batch_inds = torch.arange(loss_combo.size(0), device=loss_combo.device) -# loss_mask = loss_multimask[batch_inds, best_loss_inds].unsqueeze(1) -# loss_dice = loss_multidice[batch_inds, best_loss_inds].unsqueeze(1) -# # calculate the iou prediction and slot losses only in the index -# # with the minimum loss for each mask (to be consistent w/ SAM) -# if self.supervise_all_iou: -# loss_iou = loss_multiiou.mean(dim=-1).unsqueeze(1) -# else: -# loss_iou = loss_multiiou[batch_inds, best_loss_inds].unsqueeze(1) -# else: -# loss_mask = loss_multimask -# loss_dice = loss_multidice -# loss_iou = loss_multiiou - -# # backprop focal, dice and iou loss only if obj present -# loss_mask = loss_mask * target_obj -# loss_dice = loss_dice * target_obj -# loss_iou = loss_iou * target_obj - -# # sum over batch dimension (note that the losses are already divided by num_boxes) -# losses["loss_mask"] += loss_mask.sum() -# losses["loss_dice"] += loss_dice.sum() -# losses["loss_iou"] += loss_iou.sum() -# losses["loss_class"] += loss_class - - -# class TextCriterion(LossWithWeights): -# def __init__( -# self, -# pad_token, -# max_seq_len=100, -# weight_dict=None, -# compute_aux=False, -# ): -# super().__init__(weight_dict, compute_aux) -# self.pad_token = pad_token -# self.max_seq_len = max_seq_len -# self.in_lengths = None - -# def get_loss(self, outputs, **kwargs): -# nb_tokens = outputs["captioning_tokenized_target"].input_ids.numel() -# bs, seq_len = outputs["captioning_tokenized_target"].input_ids.shape -# ce = F.cross_entropy( -# outputs["captioning_pred_text"].flatten(0, -2), -# outputs["captioning_tokenized_target"].input_ids.flatten(), -# ignore_index=self.pad_token, -# reduction="sum", -# ) - -# not_pad = ( -# outputs["captioning_tokenized_target"] -# .input_ids.reshape(-1) -# .ne(self.pad_token) -# ) - -# if nb_tokens > 0: -# nb_non_pad = not_pad.numel() -# ce = ce / nb_non_pad - -# preds = outputs["captioning_pred_text"].flatten(0, -2).argmax(-1)[not_pad] -# targets = outputs["captioning_tokenized_target"].input_ids.flatten()[not_pad] -# correct = preds == targets -# correct = correct.sum() / (correct.numel() + 1e-5) - -# correct_sequence_level = torch.all( -# ( -# outputs["captioning_pred_text"] -# .flatten(0, -2) -# .argmax(-1) -# .reshape(bs, seq_len) -# == outputs["captioning_tokenized_target"].input_ids -# ) -# | (~not_pad).view(bs, seq_len), -# dim=1, -# ) -# seq_level_acc = correct_sequence_level.float().mean() - -# return {"loss_text": ce, "text_acc": correct, "text_seq_acc": seq_level_acc} - - -def segment_miou(source, target): - """Compute the mean IoU between two sets of masks""" - assert source.shape == target.shape, "The two masks must have the same shape" - assert source.ndim == 3, "The masks must be 3D" - - valid_targets = (target.sum(dim=(1, 2)) > 0).sum() - if valid_targets == 0: - return torch.tensor(1.0, device=source.device) - intersection = (source.bool() & target.bool()).sum(dim=(1, 2)) - union = (source.bool() | target.bool()).sum(dim=(1, 2)) - iou = intersection / (union + 1e-8) - return iou.sum() / valid_targets - - -class SemanticSegCriterion(LossWithWeights): - def __init__( - self, - weight_dict, - focal: bool = False, - focal_alpha: float = 0.6, - focal_gamma: float = 1.6, - downsample: bool = True, - presence_head: bool = False, - # Option to turn off presence loss, if some other component - # is already doing it, e.g. decoder - in which case, - # we could still set presence_head to True so that - # losses are not propogated to masks when there is no GT mask - presence_loss: bool = True, - ): - super().__init__(weight_dict, False) - self.focal = focal - self.focal_alpha = focal_alpha - self.focal_gamma = focal_gamma - self.downsample = downsample - self.presence_head = presence_head - self.presence_loss = presence_loss - - def get_loss(self, out_dict, targets): - outputs = out_dict["semantic_seg"] - presence_logit = out_dict["presence_logit"] - if ( - "semantic_masks" in targets - and targets["semantic_masks"] is not None - and targets["semantic_masks"].size(0) > 0 - ): - semantic_targets = targets["semantic_masks"] - with torch.no_grad(): - if self.downsample: - # downsample targets to the size of predictions - size = outputs.shape[-2:] - semantic_targets = ( - F.interpolate( - semantic_targets.float().unsqueeze(1), - size=size, - mode="bilinear", - align_corners=False, - ) - .squeeze(1) - .bool() - ) - else: - with torch.no_grad(): - if self.downsample: - # downsample targets to the size of predictions - size = outputs.shape[-2:] - segments = ( - F.interpolate( - targets["masks"].float().unsqueeze(1), - size=size, - mode="bilinear", - align_corners=False, - ) - .squeeze(1) - .bool() - ) - else: - segments = targets["masks"].bool() - - # the annotations are for instance segmentation, so we merge them to get semantic segmentation - semantic_targets = instance_masks_to_semantic_masks( - segments, targets["num_boxes"] - ) - - if not self.downsample: - # upsample predictions to the target size - size = semantic_targets.shape[-2:] - outputs = F.interpolate( - outputs.float(), - size=size, - mode="bilinear", - align_corners=False, - ) - - if self.focal: - loss = sigmoid_focal_loss( - outputs.squeeze(1).flatten(-2), - semantic_targets.float().flatten(-2), - num_boxes=len(semantic_targets), - alpha=self.focal_alpha, - gamma=self.focal_gamma, - reduce=not self.presence_head, - ) - if self.presence_head: - loss = loss.mean(1) - else: - loss = F.binary_cross_entropy_with_logits( - outputs.squeeze(1), - semantic_targets.float(), - reduction="none" if self.presence_head else "mean", - ) - if self.presence_head: - loss = loss.flatten(1).mean(1) - - loss_dice = dice_loss( - outputs.squeeze(1).flatten(1), - semantic_targets.flatten(1), - len(semantic_targets), - reduce=not self.presence_head, - ) - - miou = segment_miou(outputs.sigmoid().squeeze(1) > 0.5, semantic_targets) - - loss_dict = {} - - if self.presence_head: - presence_target = semantic_targets.flatten(1).any(-1) - if self.presence_loss: - loss_presence = F.binary_cross_entropy_with_logits( - presence_logit.flatten(), - presence_target.float(), - ) - presence_acc = ( - ((presence_logit.flatten().sigmoid() > 0.5) == presence_target) - .float() - .mean() - ) - else: - # Dummy values - loss_presence = torch.tensor(0.0, device=loss.device) - # Whichever component is computing the presence loss, - # should also track presence_acc - presence_acc = torch.tensor(0.0, device=loss.device) - - loss_dict["loss_semantic_presence"] = loss_presence - loss_dict["presence_acc"] = presence_acc - - # reduce the other losses, skipping the negative ones - bs = loss.shape[0] - assert presence_target.numel() == bs - - mask = presence_target - nb_valid = presence_target.sum().item() - - loss = (loss * mask.float()).sum() / (nb_valid + 1e-6) - loss_dice = (loss_dice * mask.float()).sum() / (nb_valid + 1e-6) - - loss_dict.update( - { - "loss_semantic_seg": loss, - "loss_semantic_dice": loss_dice, - "miou_semantic_seg": miou, - } - ) - - return loss_dict - - -class Det2TrkAssoc(LossWithWeights): - def __init__( - self, - weight_dict, - use_fp_loss=False, - fp_loss_on_exhaustive_only=True, - treat_fp_as_new_obj=False, - ): - super().__init__(weight_dict, compute_aux=False) - self.use_fp_loss = use_fp_loss - self.fp_loss_on_exhaustive_only = fp_loss_on_exhaustive_only - self.treat_fp_as_new_obj = treat_fp_as_new_obj - if self.use_fp_loss: - self.target_keys.append("is_exhaustive") - - def get_loss(self, outputs, targets, indices, num_boxes): - det2trk_assoc_logits = outputs["det2trk_assoc_logits"] - device = det2trk_assoc_logits.device - B, Q_det, Q_trk_plus_2 = det2trk_assoc_logits.shape - assert Q_trk_plus_2 >= 2 - Q_trk = Q_trk_plus_2 - 2 - - # We only apply association losses to those detection queries that either match - # a GT instance or have score > 0 (i.e. those TP, FN and FP detection queries) - matched_object_ids = outputs["matched_object_ids"] - assert matched_object_ids.shape == (B, Q_det + Q_trk) - matched_obj_ids_det = matched_object_ids[:, :Q_det] - matched_obj_ids_trk = matched_object_ids[:, Q_det:] - det_is_matched_to_gt = matched_obj_ids_det >= 0 - trk_is_matched_to_gt = matched_obj_ids_trk >= 0 - - # note: -1 label is ignored in the (softmax) cross_entropy loss below - det2trk_assoc_labels = -torch.ones(B, Q_det, dtype=torch.long, device=device) - # a) If a detection query is matched to a same object ID as a tracking query, - # we assign it the index of the tracking query as a label - det_is_same_obj_id_as_trk = ( - det_is_matched_to_gt[:, :, None] - & trk_is_matched_to_gt[:, None, :] - & (matched_obj_ids_det[:, :, None] == matched_obj_ids_trk[:, None, :]) - ) - batch_idx, det_idx, trk_idx = det_is_same_obj_id_as_trk.nonzero(as_tuple=True) - det2trk_assoc_labels[batch_idx, det_idx] = trk_idx - - # b) If a detection query is matched to GT but not to any tracking query, - # we assign it a "new_object" label - det_is_new_obj = det_is_matched_to_gt & ~det_is_same_obj_id_as_trk.any(dim=-1) - det2trk_assoc_labels[det_is_new_obj] = Q_trk - - # c) If a detection query is not matched to GT but have score > 0, - # we assign it a "false_positive" label - if self.use_fp_loss: - det_is_above_thresh = outputs["pred_logits"][:, :Q_det].squeeze(2) > 0 - det_is_fp = ~det_is_matched_to_gt & det_is_above_thresh - if self.treat_fp_as_new_obj: - det2trk_assoc_labels[det_is_fp] = Q_trk - else: - if self.fp_loss_on_exhaustive_only: - # only count FP detections on batches that are exhaustively annotated - det_is_fp &= targets["is_exhaustive"].unsqueeze(1).bool() - det2trk_assoc_labels[det_is_fp] = Q_trk + 1 - - # softmax cross-entropy loss for detection-to-tracking association - loss_det2trk_assoc = F.cross_entropy( - input=det2trk_assoc_logits.flatten(0, 1), # (B * Q_det, Q_trk + 2) - target=det2trk_assoc_labels.flatten(0, 1), # (B * Q_det) - ignore_index=-1, - reduction="none", - ).view(B, Q_det) - # skip det2trk assocation loss on frames w/o any (non-padding) tracking queries - frame_has_valid_trk = trk_is_matched_to_gt.any(dim=-1, keepdims=True) # (B, 1) - loss_det2trk_assoc = loss_det2trk_assoc * frame_has_valid_trk.float() - - loss_det2trk_assoc = loss_det2trk_assoc.sum() / (B * num_boxes) - return {"loss_det2trk_assoc": loss_det2trk_assoc} - - -class TrackingByDetectionAssoc(LossWithWeights): - def __init__(self, weight_dict): - super().__init__(weight_dict, compute_aux=False, supports_o2m_loss=False) - assert "loss_det2trk_assoc" in self.weight_dict - assert "loss_trk2det_assoc" in self.weight_dict - - def get_loss(self, outputs, targets, indices, num_boxes): - # Part A: gather object id matching between detection and tracking - det2trk_assoc_logits = outputs["det2trk_assoc_logits"] # (B, Q_det+1, Q_trk+1) - B, Q_det_plus_1, Q_trk_plus_1 = det2trk_assoc_logits.shape - assert Q_det_plus_1 >= 1 and Q_trk_plus_1 >= 1 - Q_det = Q_det_plus_1 - 1 - Q_trk = Q_trk_plus_1 - 1 - device = det2trk_assoc_logits.device - - matched_obj_ids_det = outputs["matched_object_ids"] - assert matched_obj_ids_det.shape == (B, Q_det) - det_is_matched_to_gt = matched_obj_ids_det >= 0 - matched_obj_ids_trk = outputs["prev_trk_object_ids"] - assert matched_obj_ids_trk.shape == (B, Q_trk) - trk_is_matched_to_gt = matched_obj_ids_trk >= 0 - frame_has_valid_trk = trk_is_matched_to_gt.any(dim=-1, keepdims=True) # (B, 1) - - # check whether a detection object is the same as a tracking object - det_is_same_obj_id_as_trk = ( - det_is_matched_to_gt[:, :, None] - & trk_is_matched_to_gt[:, None, :] - & (matched_obj_ids_det[:, :, None] == matched_obj_ids_trk[:, None, :]) - ) # (B, Q_det, Q_trk) - # there should be at most one match for each detection and each previous tracked object - torch._assert_async(torch.all(det_is_same_obj_id_as_trk.sum(dim=2) <= 1)) - torch._assert_async(torch.all(det_is_same_obj_id_as_trk.sum(dim=1) <= 1)) - batch_idx, det_idx, trk_idx = det_is_same_obj_id_as_trk.nonzero(as_tuple=True) - - # Part B: Detection-to-tracking association loss - # assign detection-to-tracking labels (note: -1 label is ignored in the loss below) - det2trk_assoc_labels = -torch.ones(B, Q_det, dtype=torch.long, device=device) - det2trk_assoc_labels[batch_idx, det_idx] = trk_idx - # if a detection is matched to GT but not to any tracking, assign it a "new-object" label - det_is_new_obj = det_is_matched_to_gt & ~det_is_same_obj_id_as_trk.any(dim=2) - det2trk_assoc_labels[det_is_new_obj] = Q_trk # "Q_trk" label is "new-object" - - # softmax cross-entropy loss for detection-to-tracking association - loss_det2trk_assoc = F.cross_entropy( - input=det2trk_assoc_logits[:, :-1].flatten(0, 1), # (B*Q_det, Q_trk+1) - target=det2trk_assoc_labels.flatten(0, 1), # (B*Q_det) - ignore_index=-1, - reduction="none", - ).view(B, Q_det) - # skip det2trk assocation loss on frames w/o any (non-padding) tracking queries - loss_det2trk_assoc = loss_det2trk_assoc * frame_has_valid_trk.float() - loss_det2trk_assoc = loss_det2trk_assoc.sum() / (B * num_boxes) - loss_dict = {"loss_det2trk_assoc": loss_det2trk_assoc} - - # Part C: tracking-to-detection association loss - trk2det_assoc_logits = det2trk_assoc_logits.transpose(1, 2) - assert trk2det_assoc_logits.shape == (B, Q_trk + 1, Q_det + 1) - # assign tracking-to-detection labels (note: -1 label is ignored in the loss below) - trk2det_assoc_labels = -torch.ones(B, Q_trk, dtype=torch.long, device=device) - trk2det_assoc_labels[batch_idx, trk_idx] = det_idx - # if a tracking is matched to GT but not to any detection, assign it a "occluded" label - trk_is_occluded = trk_is_matched_to_gt & ~det_is_same_obj_id_as_trk.any(dim=1) - trk2det_assoc_labels[trk_is_occluded] = Q_det # "Q_det" label is "occluded" - - # softmax cross-entropy loss for tracking-to-detection association - loss_trk2det_assoc = F.cross_entropy( - input=trk2det_assoc_logits[:, :-1].flatten(0, 1), # (B*Q_trk, Q_det+1) - target=trk2det_assoc_labels.flatten(0, 1), # (B*Q_trk) - ignore_index=-1, - reduction="none", - ).view(B, Q_trk) - # skip trk2det association loss on frames w/o any (non-padding) tracking queries - loss_trk2det_assoc = loss_trk2det_assoc * frame_has_valid_trk.float() - loss_trk2det_assoc = loss_trk2det_assoc.sum() / (B * num_boxes) - loss_dict["loss_trk2det_assoc"] = loss_trk2det_assoc - - return loss_dict - - -def _keep_only_trk_queries_in_match_inds(inds, Q_det): - """Keep only the tracking query indices in the indices tuple""" - batch_idx, src_idx, tgt_idx = inds - if batch_idx.numel() == 0: - return (batch_idx, src_idx, tgt_idx) # empty indices, nothing to filter - - # keep only the tracking query indices - is_trk_query = src_idx >= Q_det - batch_idx_trk = batch_idx[is_trk_query] - src_idx_trk = src_idx[is_trk_query] - tgt_idx_trk = tgt_idx[is_trk_query] if tgt_idx is not None else None - return (batch_idx_trk, src_idx_trk, tgt_idx_trk) diff --git a/train/loss/mask_sampling.py b/train/loss/mask_sampling.py deleted file mode 100644 index aeba3fe..0000000 --- a/train/loss/mask_sampling.py +++ /dev/null @@ -1,113 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved - -from typing import Callable - -import torch -from torch.nn import functional as F - - -# Adapted from https://github.com/facebookresearch/detectron2/blob/main/projects/PointRend/point_rend/point_features.py -def point_sample(input, point_coords, **kwargs): - """ - A wrapper around :function:`torch.nn.functional.grid_sample` to support 3D point_coords tensors. - Unlike :function:`torch.nn.functional.grid_sample` it assumes `point_coords` to lie inside - [0, 1] x [0, 1] square. - - Args: - input (Tensor): A tensor of shape (N, C, H, W) that contains features map on a H x W grid. - point_coords (Tensor): A tensor of shape (N, P, 2) or (N, Hgrid, Wgrid, 2) that contains - [0, 1] x [0, 1] normalized point coordinates. - - Returns: - output (Tensor): A tensor of shape (N, C, P) or (N, C, Hgrid, Wgrid) that contains - features for points in `point_coords`. The features are obtained via bilinear - interplation from `input` the same way as :function:`torch.nn.functional.grid_sample`. - """ - add_dim = False - if point_coords.dim() == 3: - add_dim = True - point_coords = point_coords.unsqueeze(2) - normalized_point_coords = 2.0 * point_coords - 1.0 # Normalize to [-1,1] - output = F.grid_sample(input, normalized_point_coords, **kwargs) - if add_dim: - output = output.squeeze(3) - return output - - -# Adapted from https://github.com/facebookresearch/detectron2/blob/main/projects/PointRend/point_rend/point_features.py -def get_uncertain_point_coords_with_randomness( - logits: torch.Tensor, - uncertainty_func: Callable, - num_points: int, - oversample_ratio: int, - importance_sample_ratio: float, -) -> torch.Tensor: - """ - Sample points in [0, 1] x [0, 1] coordinate space based on their uncertainty. The unceratinties - are calculated for each point using 'uncertainty_func' function that takes point's logit - prediction as input. - See PointRend paper for details. - - Args: - logits (Tensor): A tensor of shape (N, C, Hmask, Wmask) or (N, 1, Hmask, Wmask) for - class-specific or class-agnostic prediction. - uncertainty_func: A function that takes a Tensor of shape (N, C, P) or (N, 1, P) that - contains logit predictions for P points and returns their uncertainties as a Tensor of - shape (N, 1, P). - num_points (int): The number of points P to sample. - oversample_ratio (int): Oversampling parameter. - importance_sample_ratio (float): Ratio of points that are sampled via importnace sampling. - - Returns: - point_coords (Tensor): A tensor of shape (N, P, 2) that contains the coordinates of P - sampled points. - """ - assert oversample_ratio >= 1 - assert importance_sample_ratio <= 1 and importance_sample_ratio >= 0 - num_boxes = logits.shape[0] - num_sampled = int(num_points * oversample_ratio) - point_coords = torch.rand(num_boxes, num_sampled, 2, device=logits.device) - point_logits = point_sample(logits, point_coords, align_corners=False) - # It is crucial to calculate uncertainty based on the sampled prediction value for the points. - # Calculating uncertainties of the predictions first and sampling them for points leads - # to incorrect results. - # To illustrate this: assume uncertainty_func(logits)=-abs(logits), a sampled point between - # two predictions with -1 and 1 logits has 0 logits, and therefore 0 uncertainty value. - # However, if we calculate uncertainties for the predictions first, - # both will have -1 uncertainty, and the sampled point will get -1 uncertainty. - point_uncertainties = uncertainty_func(point_logits) - num_uncertain_points = int(importance_sample_ratio * num_points) - num_random_points = num_points - num_uncertain_points - idx = torch.topk(point_uncertainties[:, 0, :], k=num_uncertain_points, dim=1)[1] - # Flatten the indices - shift = num_sampled * torch.arange( - num_boxes, dtype=torch.long, device=logits.device - ) - idx += shift[:, None] - point_coords = point_coords.view(-1, 2)[idx.view(-1), :].view( - num_boxes, num_uncertain_points, 2 - ) - if num_random_points > 0: - point_coords = torch.cat( - [ - point_coords, - torch.rand(num_boxes, num_random_points, 2, device=logits.device), - ], - dim=1, - ) - return point_coords - - -# Adapted from https://github.com/facebookresearch/Mask2Former/blob/main/mask2former/modeling/criterion.py -def calculate_uncertainty(logits: torch.Tensor) -> torch.Tensor: - """ - Estimates uncerainty as L1 distance between 0.0 and the logit prediction. - Args: - logits (Tensor): A tensor of shape (R, 1, ...) for class-agnostic - predicted masks - Returns: - scores (Tensor): A tensor of shape (R, 1, ...) that contains uncertainty scores with - the most uncertain locations having the highest uncertainty score. - """ - assert logits.shape[1] == 1 - return -(torch.abs(logits)) diff --git a/train/loss/sam3_loss.py b/train/loss/sam3_loss.py deleted file mode 100644 index 7ef59dc..0000000 --- a/train/loss/sam3_loss.py +++ /dev/null @@ -1,203 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved - -import torch - -from sam3.model.model_misc import SAM3Output - -from sam3.train.utils.distributed import get_world_size - -from .loss_fns import CORE_LOSS_KEY, Det2TrkAssoc, Masks - - -class DummyLoss(torch.nn.Module): - """A dummy loss that always returns 0 (as a placeholder for eval)""" - - def __init__( - self, - core_loss_key: str = CORE_LOSS_KEY, - device: str = "cuda", - **kwargs, - ): - super().__init__() - self.core_loss_key = core_loss_key - self.device = torch.device(device) - - def forward(self, *args, **kwargs): - return {self.core_loss_key: torch.tensor(0.0, device=self.device)} - - def accumulate(self, out_dict): - """ - Called by iterative losses. - """ - if self.core_loss_key not in out_dict: - out_dict[self.core_loss_key] = torch.tensor(0.0, device=self.device) - return out_dict - - -class Sam3LossWrapper(torch.nn.Module): - def __init__( - self, - loss_fns_find, - normalization="global", - matcher=None, - o2m_matcher=None, - o2m_weight=1.0, - use_o2m_matcher_on_o2m_aux=True, - loss_fn_semantic_seg=None, - normalize_by_valid_object_num=False, - normalize_by_stage_num=False, - scale_by_find_batch_size=False, - ): - super().__init__() - self.loss_fns_find = loss_fns_find - assert normalization in ["global", "local", "none"] - self.normalization = normalization - self.normalize_by_valid_object_num = normalize_by_valid_object_num - self.normalize_by_stage_num = normalize_by_stage_num - self.matcher = matcher - self.o2m_matcher = o2m_matcher - self.o2m_weight = o2m_weight - # whether to use the o2m matcher on the o2m queries in auxiliary outputs - self.use_o2m_matcher_on_o2m_aux = use_o2m_matcher_on_o2m_aux - self.loss_fn_semantic_seg = loss_fn_semantic_seg - self.scale_by_find_batch_size = scale_by_find_batch_size - - def _get_num_boxes(self, targets): - # the average number of target boxes for loss normalization - if self.normalize_by_valid_object_num: - # valid boxes are those with non-zero height and width - # (while padded invisible boxes are ) - boxes_hw = targets["boxes"].view(-1, 4) # cx, cy, w, h - num_boxes = (boxes_hw[:, 2:] > 0).all(dim=-1).sum().float() - else: - num_boxes = targets["num_boxes"].sum().float() - if self.normalization == "global": - torch.distributed.all_reduce(num_boxes) - num_boxes = torch.clamp(num_boxes / get_world_size(), min=1) - elif self.normalization == "local": - num_boxes = torch.clamp(num_boxes, min=1) - elif self.normalization == "none": - num_boxes = 1 - return num_boxes - - def compute_loss(self, nested_out, targets): - num_boxes = self._get_num_boxes(targets) - o2m_out_is_valid = nested_out.get("o2m_out_is_valid", None) - o2m_target_is_valid_padded = nested_out.get("o2m_target_is_valid_padded", None) - - # Get a list of outputs, including auxiliary and first stage outputs - output_list = [(nested_out, "", False)] # (out, suffix, is_aux) - if "aux_outputs" in nested_out: - output_list.extend( - (aux_out, f"_aux_{i}", True) - for i, aux_out in enumerate(nested_out["aux_outputs"]) - ) - if "first_stage" in nested_out: - output_list.append((nested_out["first_stage"], "_fs", True)) - - # Compute all the requested losses - losses = {} - total_core_loss = 0.0 - for out, suffix, is_aux in output_list: - # o2o matcher indices need to be computed by the model (as the video model requires - # a specific way of matching free and locked indices beyond just calling the matcher) - indices = out["indices"] - has_o2m_out = "pred_logits_o2m" in out - if has_o2m_out: - o2m_out = { - k[: -len("_o2m")]: v for k, v in out.items() if k.endswith("_o2m") - } - # o2m targets are the same as the o2o targets (assuming repeat=1) - o2m_targets = targets - if self.use_o2m_matcher_on_o2m_aux or not is_aux: - o2m_indices = self.o2m_matcher( - o2m_out, - o2m_targets, - out_is_valid=o2m_out_is_valid, - target_is_valid_padded=o2m_target_is_valid_padded, - ) - else: - o2m_indices = self.matcher( - o2m_out, - o2m_targets, - out_is_valid=o2m_out_is_valid, - target_is_valid_padded=o2m_target_is_valid_padded, - ) - - for loss_fn in self.loss_fns_find: - l_dict = loss_fn( - outputs=out, - targets=targets, - indices=indices, - num_boxes=num_boxes, - is_aux=is_aux, - ) - total_core_loss += l_dict.pop(CORE_LOSS_KEY) - losses.update({f"{k}{suffix}": v for k, v in l_dict.items()}) - - compute_o2m_loss = has_o2m_out - # a special handling to allow turning off mask loss in o2m - # (to be compatible with the original implementation) - if isinstance(loss_fn, Masks): - compute_o2m_loss = compute_o2m_loss and "pred_masks" in o2m_out - if isinstance(loss_fn, Det2TrkAssoc): - compute_o2m_loss = False # Det2TrkAssoc does not support o2m - if compute_o2m_loss: - l_dict = loss_fn( - outputs=o2m_out, - targets=o2m_targets, - indices=o2m_indices, - num_boxes=num_boxes, - is_aux=is_aux, - ) - for k in l_dict: - l_dict[k] *= self.o2m_weight - total_core_loss += l_dict.pop(CORE_LOSS_KEY) - losses.update({f"{k}{suffix}_o2m": v for k, v in l_dict.items()}) - - losses[CORE_LOSS_KEY] = total_core_loss - return losses - - def forward(self, find_stages: SAM3Output, find_targets): - if find_stages.loss_stages is not None: - find_targets = [find_targets[i] for i in find_stages.loss_stages] - with SAM3Output.iteration_mode( - find_stages, iter_mode=SAM3Output.IterMode.ALL_STEPS_PER_STAGE - ) as find_stages: - assert len(find_stages) == len(find_targets) - total_losses = {} - for stage_outputs, stage_targets in zip(find_stages, find_targets): - stage_targets = [stage_targets] * len(stage_outputs) - # If there are multiple steps within a stage, compute the loss for all of them (e.g. interactivity) - for outputs, targets in zip(stage_outputs, stage_targets): - cur_losses = self.compute_loss(outputs, targets) - - if self.loss_fn_semantic_seg is not None: - cur_losses_semantic = self.loss_fn_semantic_seg( - outputs, targets - ) - cur_losses[CORE_LOSS_KEY] += cur_losses_semantic.pop( - CORE_LOSS_KEY - ) - # make sure the semantic losses don't overlap with the find losses - assert set(cur_losses).isdisjoint(set(cur_losses_semantic)) - cur_losses.update(cur_losses_semantic) - - # Optionally, normalize the loss by the number of find stages (training video frames) so that - # image batches and video batches have similar loss scales. (Otherwise video batches would - # have a much higher loss scale due to summing the losses over all the find stages.) - if self.normalize_by_stage_num: - cur_losses[CORE_LOSS_KEY] /= len(find_stages) - - if self.scale_by_find_batch_size: - bs = targets["num_boxes"].shape[0] - # sqrt scaling based on the "effective" batch size - cur_losses[CORE_LOSS_KEY] *= bs**0.5 - - for k, v in cur_losses.items(): - if k not in total_losses: - total_losses[k] = v - else: - total_losses[k] += v - - return total_losses diff --git a/train/loss/sigmoid_focal_loss.py b/train/loss/sigmoid_focal_loss.py deleted file mode 100644 index 15e6db4..0000000 --- a/train/loss/sigmoid_focal_loss.py +++ /dev/null @@ -1,321 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved - -"""Triton kernel for faster and memory efficient sigmoid focal loss""" - -import torch -import triton -import triton.language as tl -from torch._inductor.runtime.triton_helpers import libdevice - -""" - -The sigmoid focal loss is defined as: - - prob = inputs.sigmoid() - ce_loss = F.binary_cross_entropy_with_logits(inputs, targets, reduction="none") - p_t = prob * targets + (1 - prob) * (1 - targets) - alpha_t = alpha * targets + (1 - alpha) * (1 - targets) - loss = alpha_t * ce_loss * ((1 - p_t) ** gamma) - -Where alpha and gamma are scalar parameters, inputs are the logits, targets the float targets. - -We implement two versions of the sigmoid focal loss: with and without sum reduction. -The latter is implemented with built-in reduction to avoid materializing wrt the output of the loss. -This can help save a bit of peak memory. - -The reduction version is implemented using somewhat of a hack. Pytorch's generated kernels usually do the point-wise operation in a first kernel, and implement the reduction another kernel launched on a grid of size 1, where the reduction happens as a for loop in the triton kernel. -Since we want to fuse those two kernels, that is not a good idea: we'd have to launch the overall kernel on a grid of size 1, which is obviously inefficient. -On the other hand, typical CUDA algorithms for reduction (eg reduction tree) are hard to implement in triton due to the lack of thread sync primitives. -We settle for a version that abuses triton's atomic_add: we can have all threads simply add to the same location. -In practice, this is not good, since it creates a massive bottleneck on the semaphore for that single memory location. So instead, we create M reduction locations. Each thread will simply write to thread_id%M. The python code can finally sum over the M reductions. -M = 32 works fine in benchmarking tests. The forward is a tiny bit slower compared to the non-reduced kernel, but the backward breaks even due to one less memory allocation. -""" - - -@triton.jit -def _inner_focal_loss_fwd(inputs, targets, alpha, gamma): - inv_targets = 1 - targets - # Sigmoid - sig = tl.sigmoid(inputs) - - # Binary cross entropy with logits - # In practice, we want the following: - # bce_loss = -targets * tl.log(sig) - (1 - targets) * tl.log(1 - sig) - # However, the above is not numerically stable. - # We're also not directly taking the sum here, so the usual log-sum-exp trick doesn't apply - # The bce can be reformulated, after algebraic manipulation, to - # bce_loss = log(1 + exp(-x)) + x * (1-y) - # This is still not stable, because for large (-x) the exponential will blow up. - # We'll use the following alternate formulation: - # bce_loss = max(x, 0) - x * y + log(1 + exp(-abs(x))) - # Let's show that it's equivalent: - # Case x>=0: abs(x) = x , max(x, 0) = x - # so we get x - x * y + log(1 + exp(-x)) which is equivalent - # Case x<0: abs(x) = -x, max(x, 0) = 0 - # we have log(1 + exp(-abs(x))) = log(1 + exp(x)) = log(exp(x)(1 + exp(-x))) = x+log(1 + exp(-x)) - # plugging it in, we get - # 0 - x * y + x + log(1 + exp(-x)), which is also equivalent - # Note that this is stable because now the exponent are guaranteed to be below 0. - max_val = tl.clamp(inputs, min=0, max=1e9) - bce_loss = max_val - inputs * targets + tl.log(1 + tl.exp(-tl.abs(inputs))) - - # Modulating factor - p_t = sig * targets + (1 - sig) * inv_targets - mod_factor = libdevice.pow(1 - p_t, gamma) - - # Alpha factor - alpha_t = alpha * targets + (1 - alpha) * inv_targets - - # Final loss calculation - return alpha_t * mod_factor * bce_loss - - -# Non-reduced version -@triton.jit -def sigmoid_focal_loss_fwd_kernel( - inputs_ptr, - targets_ptr, - loss_ptr, - alpha: float, - gamma: float, - n_elements: int, - BLOCK_SIZE: tl.constexpr, -): - pid = tl.program_id(axis=0) - block_start = pid * BLOCK_SIZE - offset = block_start + tl.arange(0, BLOCK_SIZE) - mask = offset < n_elements - - # Load data - inputs = tl.load(inputs_ptr + offset, mask=mask).to(tl.float32) - targets = tl.load(targets_ptr + offset, mask=mask) - - final_loss = _inner_focal_loss_fwd(inputs, targets, alpha, gamma) - - # Store result - tl.store(loss_ptr + offset, final_loss, mask=mask) - - -# version with reduction -@triton.jit -def sigmoid_focal_loss_fwd_kernel_reduce( - inputs_ptr, - targets_ptr, - loss_ptr, - alpha: float, - gamma: float, - n_elements: int, - BLOCK_SIZE: tl.constexpr, - REDUCE_SIZE: tl.constexpr, -): - pid = tl.program_id(axis=0) - block_start = pid * BLOCK_SIZE - reduce_loc = pid % REDUCE_SIZE - offset = block_start + tl.arange(0, BLOCK_SIZE) - mask = offset < n_elements - # Load data - inputs = tl.load(inputs_ptr + offset, mask=mask).to(tl.float32) - targets = tl.load(targets_ptr + offset, mask=mask) - - final_loss = _inner_focal_loss_fwd(inputs, targets, alpha, gamma) * mask - - fl = tl.sum(final_loss) - - # Store result - tl.atomic_add(loss_ptr + reduce_loc, fl) - - -@triton.jit -def _inner_focal_loss_bwd(inputs, targets, alpha, gamma): - inv_targets = 1 - targets - - # Recompute forward - max_val = tl.clamp(inputs, min=0, max=1e9) - bce_loss = max_val - inputs * targets + tl.log(1 + tl.exp(-tl.abs(inputs))) - - # Sigmoid - sig = tl.sigmoid(inputs) - inv_sig = 1 - sig - - # Modulating factor - p_t = sig * targets + inv_sig * inv_targets - tmp = libdevice.pow(1 - p_t, gamma - 1) - mod_factor = tmp * (1 - p_t) - - # Alpha factor - alpha_t = alpha * targets + (1 - alpha) * inv_targets - - # Now computing the derivatives - d_pt = (2 * targets - 1) * sig * inv_sig - d_mod_factor = -gamma * d_pt * tmp - - d_bce_loss = sig - targets - - return alpha_t * (d_bce_loss * mod_factor + d_mod_factor * bce_loss) - - -@triton.jit -def sigmoid_focal_loss_bwd_kernel( - inputs_ptr, - targets_ptr, - grad_inputs_ptr, - grad_out_ptr, - alpha: float, - gamma: float, - n_elements: int, - BLOCK_SIZE: tl.constexpr, -): - pid = tl.program_id(axis=0) - block_start = pid * BLOCK_SIZE - offset = block_start + tl.arange(0, BLOCK_SIZE) - mask = offset < n_elements - input_ptrs = inputs_ptr + offset - target_ptrs = targets_ptr + offset - grad_input_ptrs = grad_inputs_ptr + offset - grad_out_ptrs = grad_out_ptr + offset - # Load data - inputs = tl.load(input_ptrs, mask=mask).to(tl.float32) - targets = tl.load(target_ptrs, mask=mask) - grad_out = tl.load(grad_out_ptrs, mask=mask) - d_loss = grad_out * _inner_focal_loss_bwd(inputs, targets, alpha, gamma) - tl.store(grad_input_ptrs, d_loss, mask=mask) - - -@triton.jit -def sigmoid_focal_loss_bwd_kernel_reduce( - inputs_ptr, - targets_ptr, - grad_inputs_ptr, - grad_out_ptr, - alpha: float, - gamma: float, - n_elements: int, - BLOCK_SIZE: tl.constexpr, -): - # The only difference is that the gradient is now a single scalar - pid = tl.program_id(axis=0) - block_start = pid * BLOCK_SIZE - offset = block_start + tl.arange(0, BLOCK_SIZE) - mask = offset < n_elements - input_ptrs = inputs_ptr + offset - target_ptrs = targets_ptr + offset - grad_input_ptrs = grad_inputs_ptr + offset - # Load data - inputs = tl.load(input_ptrs, mask=mask).to(tl.float32) - targets = tl.load(target_ptrs, mask=mask) - grad_out = tl.load(grad_out_ptr) - d_loss = grad_out * _inner_focal_loss_bwd(inputs, targets, alpha, gamma) - tl.store(grad_input_ptrs, d_loss, mask=mask) - - -class SigmoidFocalLoss(torch.autograd.Function): - BLOCK_SIZE = 256 - - @staticmethod - def forward(ctx, inputs, targets, alpha=0.25, gamma=2): - n_elements = inputs.numel() - assert targets.numel() == n_elements - input_shape = inputs.shape - inputs = inputs.view(-1).contiguous() - targets = targets.view(-1).contiguous() - loss = torch.empty(inputs.shape, dtype=torch.float32, device=inputs.device) - grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),) - sigmoid_focal_loss_fwd_kernel[grid]( - inputs, targets, loss, alpha, gamma, n_elements, SigmoidFocalLoss.BLOCK_SIZE - ) - ctx.save_for_backward(inputs.view(input_shape), targets.view(input_shape)) - ctx.alpha = alpha - ctx.gamma = gamma - return loss.view(input_shape) - - @staticmethod - def backward(ctx, grad_output): - inputs, targets = ctx.saved_tensors - alpha = ctx.alpha - gamma = ctx.gamma - n_elements = inputs.numel() - input_shape = inputs.shape - grad_inputs = torch.empty( - inputs.shape, dtype=grad_output.dtype, device=grad_output.device - ) - inputs_ptr = inputs.view(-1).contiguous() - targets_ptr = targets.view(-1).contiguous() - grad_output_ptr = grad_output.view(-1).contiguous() - grad_inputs_ptr = grad_inputs - assert grad_output.numel() == n_elements - grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),) - sigmoid_focal_loss_bwd_kernel[grid]( - inputs_ptr, - targets_ptr, - grad_inputs_ptr, - grad_output_ptr, - alpha, - gamma, - n_elements, - SigmoidFocalLoss.BLOCK_SIZE, - ) - return grad_inputs.view(input_shape), None, None, None - - -triton_sigmoid_focal_loss = SigmoidFocalLoss.apply - - -class SigmoidFocalLossReduced(torch.autograd.Function): - BLOCK_SIZE = 256 - REDUCE_SIZE = 32 - - @staticmethod - def forward(ctx, inputs, targets, alpha=0.25, gamma=2): - n_elements = inputs.numel() - input_shape = inputs.shape - inputs = inputs.view(-1).contiguous() - targets = targets.view(-1).contiguous() - loss = torch.zeros( - SigmoidFocalLossReduced.REDUCE_SIZE, - device=inputs.device, - dtype=torch.float32, - ) - grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),) - sigmoid_focal_loss_fwd_kernel_reduce[grid]( - inputs, - targets, - loss, - alpha, - gamma, - n_elements, - SigmoidFocalLossReduced.BLOCK_SIZE, - SigmoidFocalLossReduced.REDUCE_SIZE, - ) - ctx.save_for_backward(inputs.view(input_shape), targets.view(input_shape)) - ctx.alpha = alpha - ctx.gamma = gamma - return loss.sum() - - @staticmethod - def backward(ctx, grad_output): - inputs, targets = ctx.saved_tensors - alpha = ctx.alpha - gamma = ctx.gamma - n_elements = inputs.numel() - input_shape = inputs.shape - grad_inputs = torch.empty( - inputs.shape, dtype=grad_output.dtype, device=grad_output.device - ) - inputs_ptr = inputs.view(-1).contiguous() - targets_ptr = targets.reshape(-1).contiguous() - assert grad_output.numel() == 1 - grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),) - sigmoid_focal_loss_bwd_kernel_reduce[grid]( - inputs_ptr, - targets_ptr, - grad_inputs, - grad_output, - alpha, - gamma, - n_elements, - SigmoidFocalLossReduced.BLOCK_SIZE, - ) - return grad_inputs.view(input_shape), None, None, None - - -triton_sigmoid_focal_loss_reduce = SigmoidFocalLossReduced.apply diff --git a/train/masks_ops.py b/train/masks_ops.py deleted file mode 100644 index f9d2fd7..0000000 --- a/train/masks_ops.py +++ /dev/null @@ -1,272 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved - -"""Utilities for masks manipulation""" - -import numpy as np -import pycocotools.mask as maskUtils -import torch -from pycocotools import mask as mask_util - - -def instance_masks_to_semantic_masks( - instance_masks: torch.Tensor, num_instances: torch.Tensor -) -> torch.Tensor: - """This function converts instance masks to semantic masks. - It accepts a collapsed batch of instances masks (ie all instance masks are concatenated in a single tensor) and - the number of instances in each image of the batch. - It returns a mask with the same spatial dimensions as the input instance masks, where for each batch element the - semantic mask is the union of all the instance masks in the batch element. - - If for a given batch element there are no instances (ie num_instances[i]==0), the corresponding semantic mask will be a tensor of zeros. - - Args: - instance_masks (torch.Tensor): A tensor of shape (N, H, W) where N is the number of instances in the batch. - num_instances (torch.Tensor): A tensor of shape (B,) where B is the batch size. It contains the number of instances - in each image of the batch. - - Returns: - torch.Tensor: A tensor of shape (B, H, W) where B is the batch size and H, W are the spatial dimensions of the - input instance masks. - """ - - masks_per_query = torch.split(instance_masks, num_instances.tolist()) - - return torch.stack([torch.any(masks, dim=0) for masks in masks_per_query], dim=0) - - -def mask_intersection(masks1, masks2, block_size=16): - """Compute the intersection of two sets of masks, without blowing the memory""" - - assert masks1.shape[1:] == masks2.shape[1:] - assert masks1.dtype == torch.bool and masks2.dtype == torch.bool - - result = torch.zeros( - masks1.shape[0], masks2.shape[0], device=masks1.device, dtype=torch.long - ) - for i in range(0, masks1.shape[0], block_size): - for j in range(0, masks2.shape[0], block_size): - intersection = ( - (masks1[i : i + block_size, None] * masks2[None, j : j + block_size]) - .flatten(-2) - .sum(-1) - ) - result[i : i + block_size, j : j + block_size] = intersection - return result - - -def mask_iom(masks1, masks2): - """ - Similar to IoU, except the denominator is the area of the smallest mask - """ - assert masks1.shape[1:] == masks2.shape[1:] - assert masks1.dtype == torch.bool and masks2.dtype == torch.bool - - # intersection = (masks1[:, None] * masks2[None]).flatten(-2).sum(-1) - intersection = mask_intersection(masks1, masks2) - area1 = masks1.flatten(-2).sum(-1) - area2 = masks2.flatten(-2).sum(-1) - min_area = torch.min(area1[:, None], area2[None, :]) - return intersection / (min_area + 1e-8) - - -def compute_boundary(seg): - """ - Adapted from https://github.com/JonathonLuiten/TrackEval/blob/master/trackeval/metrics/j_and_f.py#L148 - Return a 1pix wide boundary of the given mask - """ - assert seg.ndim >= 2 - e = torch.zeros_like(seg) - s = torch.zeros_like(seg) - se = torch.zeros_like(seg) - - e[..., :, :-1] = seg[..., :, 1:] - s[..., :-1, :] = seg[..., 1:, :] - se[..., :-1, :-1] = seg[..., 1:, 1:] - - b = seg ^ e | seg ^ s | seg ^ se - b[..., -1, :] = seg[..., -1, :] ^ e[..., -1, :] - b[..., :, -1] = seg[..., :, -1] ^ s[..., :, -1] - b[..., -1, -1] = 0 - return b - - -def dilation(mask, kernel_size): - """ - Implements the dilation operation. If the input is on cpu, we call the cv2 version. - Otherwise, we implement it using a convolution - - The kernel is assumed to be a square kernel - - """ - - assert mask.ndim == 3 - kernel_size = int(kernel_size) - assert ( - kernel_size % 2 == 1 - ), f"Dilation expects a odd kernel size, got {kernel_size}" - - if mask.is_cuda: - m = mask.unsqueeze(1).to(torch.float16) - k = torch.ones(1, 1, kernel_size, 1, dtype=m.dtype, device=m.device) - - result = torch.nn.functional.conv2d(m, k, padding="same") - result = torch.nn.functional.conv2d(result, k.transpose(-1, -2), padding="same") - return result.view_as(mask) > 0 - - all_masks = mask.view(-1, mask.size(-2), mask.size(-1)).numpy().astype(np.uint8) - kernel = np.ones((kernel_size, kernel_size), dtype=np.uint8) - - import cv2 - - processed = [torch.from_numpy(cv2.dilate(m, kernel)) for m in all_masks] - return torch.stack(processed).view_as(mask).to(mask) - - -def compute_F_measure( - gt_boundary_rle, gt_dilated_boundary_rle, dt_boundary_rle, dt_dilated_boundary_rle -): - """Adapted from https://github.com/JonathonLuiten/TrackEval/blob/master/trackeval/metrics/j_and_f.py#L207 - - Assumes the boundary and dilated boundaries have already been computed and converted to RLE - """ - gt_match = maskUtils.merge([gt_boundary_rle, dt_dilated_boundary_rle], True) - dt_match = maskUtils.merge([dt_boundary_rle, gt_dilated_boundary_rle], True) - - n_dt = maskUtils.area(dt_boundary_rle) - n_gt = maskUtils.area(gt_boundary_rle) - # % Compute precision and recall - if n_dt == 0 and n_gt > 0: - precision = 1 - recall = 0 - elif n_dt > 0 and n_gt == 0: - precision = 0 - recall = 1 - elif n_dt == 0 and n_gt == 0: - precision = 1 - recall = 1 - else: - precision = maskUtils.area(dt_match) / float(n_dt) - recall = maskUtils.area(gt_match) / float(n_gt) - - # Compute F measure - if precision + recall == 0: - f_val = 0 - else: - f_val = 2 * precision * recall / (precision + recall) - - return f_val - - -@torch.no_grad() -def rle_encode(orig_mask, return_areas=False): - """Encodes a collection of masks in RLE format - - This function emulates the behavior of the COCO API's encode function, but - is executed partially on the GPU for faster execution. - - Args: - mask (torch.Tensor): A mask of shape (N, H, W) with dtype=torch.bool - return_areas (bool): If True, add the areas of the masks as a part of - the RLE output dict under the "area" key. Default is False. - - Returns: - str: The RLE encoded masks - """ - assert orig_mask.ndim == 3, "Mask must be of shape (N, H, W)" - assert orig_mask.dtype == torch.bool, "Mask must have dtype=torch.bool" - - if orig_mask.numel() == 0: - return [] - - # First, transpose the spatial dimensions. - # This is necessary because the COCO API uses Fortran order - mask = orig_mask.transpose(1, 2) - - # Flatten the mask - flat_mask = mask.reshape(mask.shape[0], -1) - if return_areas: - mask_areas = flat_mask.sum(-1).tolist() - # Find the indices where the mask changes - differences = torch.ones( - mask.shape[0], flat_mask.shape[1] + 1, device=mask.device, dtype=torch.bool - ) - differences[:, 1:-1] = flat_mask[:, :-1] != flat_mask[:, 1:] - differences[:, 0] = flat_mask[:, 0] - _, change_indices = torch.where(differences) - - try: - boundaries = torch.cumsum(differences.sum(-1), 0).cpu() - except RuntimeError as _: - boundaries = torch.cumsum(differences.cpu().sum(-1), 0) - - change_indices_clone = change_indices.clone() - # First pass computes the RLEs on GPU, in a flatten format - for i in range(mask.shape[0]): - # Get the change indices for this batch item - beg = 0 if i == 0 else boundaries[i - 1].item() - end = boundaries[i].item() - change_indices[beg + 1 : end] -= change_indices_clone[beg : end - 1] - - # Now we can split the RLES of each batch item, and convert them to strings - # No more gpu at this point - change_indices = change_indices.tolist() - - batch_rles = [] - # Process each mask in the batch separately - for i in range(mask.shape[0]): - beg = 0 if i == 0 else boundaries[i - 1].item() - end = boundaries[i].item() - run_lengths = change_indices[beg:end] - - uncompressed_rle = {"counts": run_lengths, "size": list(orig_mask.shape[1:])} - h, w = uncompressed_rle["size"] - rle = mask_util.frPyObjects(uncompressed_rle, h, w) - rle["counts"] = rle["counts"].decode("utf-8") - if return_areas: - rle["area"] = mask_areas[i] - batch_rles.append(rle) - - return batch_rles - - -def robust_rle_encode(masks): - """Encodes a collection of masks in RLE format. Uses the gpu version fist, falls back to the cpu version if it fails""" - - assert masks.ndim == 3, "Mask must be of shape (N, H, W)" - assert masks.dtype == torch.bool, "Mask must have dtype=torch.bool" - - try: - return rle_encode(masks) - except RuntimeError as _: - masks = masks.cpu().numpy() - rles = [ - mask_util.encode( - np.array(mask[:, :, np.newaxis], dtype=np.uint8, order="F") - )[0] - for mask in masks - ] - for rle in rles: - rle["counts"] = rle["counts"].decode("utf-8") - return rles - - -def ann_to_rle(segm, im_info): - """Convert annotation which can be polygons, uncompressed RLE to RLE. - Args: - ann (dict) : annotation object - Returns: - ann (rle) - """ - h, w = im_info["height"], im_info["width"] - if isinstance(segm, list): - # polygon -- a single object might consist of multiple parts - # we merge all parts into one mask rle code - rles = mask_util.frPyObjects(segm, h, w) - rle = mask_util.merge(rles) - elif isinstance(segm["counts"], list): - # uncompressed RLE - rle = mask_util.frPyObjects(segm, h, w) - else: - # rle - rle = segm - return rle diff --git a/train/matcher.py b/train/matcher.py deleted file mode 100644 index b0b8d62..0000000 --- a/train/matcher.py +++ /dev/null @@ -1,806 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved - -""" -Modules to compute the matching cost and solve the corresponding LSAP. -""" - -import numpy as np -import torch - -from sam3.model.box_ops import box_cxcywh_to_xyxy, box_iou, generalized_box_iou -from scipy.optimize import linear_sum_assignment -from torch import nn - - -def _do_matching(cost, repeats=1, return_tgt_indices=False, do_filtering=False): - if repeats > 1: - cost = np.tile(cost, (1, repeats)) - - i, j = linear_sum_assignment(cost) - if do_filtering: - # filter out invalid entries (i.e. those with cost > 1e8) - valid_thresh = 1e8 - valid_ijs = [(ii, jj) for ii, jj in zip(i, j) if cost[ii, jj] < valid_thresh] - i, j = zip(*valid_ijs) if len(valid_ijs) > 0 else ([], []) - i, j = np.array(i, dtype=np.int64), np.array(j, dtype=np.int64) - if return_tgt_indices: - return i, j - order = np.argsort(j) - return i[order] - - -class HungarianMatcher(nn.Module): - """This class computes an assignment between the targets and the predictions of the network - - For efficiency reasons, the targets don't include the no_object. Because of this, in general, - there are more predictions than targets. In this case, we do a 1-to-1 matching of the best predictions, - while the others are un-matched (and thus treated as non-objects). - """ - - def __init__( - self, - cost_class: float = 1, - cost_bbox: float = 1, - cost_giou: float = 1, - focal_loss: bool = False, - focal_alpha: float = 0.25, - focal_gamma: float = 2, - ): - """Creates the matcher - - Params: - cost_class: This is the relative weight of the classification error in the matching cost - cost_bbox: This is the relative weight of the L1 error of the bounding box coordinates in the matching cost - cost_giou: This is the relative weight of the giou loss of the bounding box in the matching cost - """ - super().__init__() - self.cost_class = cost_class - self.cost_bbox = cost_bbox - self.cost_giou = cost_giou - self.norm = nn.Sigmoid() if focal_loss else nn.Softmax(-1) - assert ( - cost_class != 0 or cost_bbox != 0 or cost_giou != 0 - ), "all costs cant be 0" - self.focal_loss = focal_loss - self.focal_alpha = focal_alpha - self.focal_gamma = focal_gamma - - @torch.no_grad() - def forward(self, outputs, batched_targets): - """Performs the matching - - Params: - outputs: This is a dict that contains at least these entries: - "pred_logits": Tensor of dim [batch_size, num_queries, num_classes] with the classification logits - "pred_boxes": Tensor of dim [batch_size, num_queries, 4] with the predicted box coordinates - - targets: This is a list of targets (len(targets) = batch_size), where each target is a dict containing: - "labels": Tensor of dim [num_target_boxes] (where num_target_boxes is the number of ground-truth - objects in the target) containing the class labels - "boxes": Tensor of dim [num_target_boxes, 4] containing the target box coordinates - - Returns: - A list of size batch_size, containing tuples of (index_i, index_j) where: - - index_i is the indices of the selected predictions (in order) - - index_j is the indices of the corresponding selected targets (in order) - For each batch element, it holds: - len(index_i) = len(index_j) = min(num_queries, num_target_boxes) - """ - bs, num_queries = outputs["pred_logits"].shape[:2] - - # We flatten to compute the cost matrices in a batch - out_prob = self.norm( - outputs["pred_logits"].flatten(0, 1) - ) # [batch_size * num_queries, num_classes] - out_bbox = outputs["pred_boxes"].flatten(0, 1) # [batch_size * num_queries, 4] - - # Also concat the target labels and boxes - tgt_bbox = batched_targets["boxes"] - - if "positive_map" in batched_targets: - # In this case we have a multi-hot target - positive_map = batched_targets["positive_map"] - assert len(tgt_bbox) == len(positive_map) - - if self.focal_loss: - positive_map = positive_map > 1e-4 - alpha = self.focal_alpha - gamma = self.focal_gamma - neg_cost_class = ( - (1 - alpha) * (out_prob**gamma) * (-(1 - out_prob + 1e-8).log()) - ) - pos_cost_class = ( - alpha * ((1 - out_prob) ** gamma) * (-(out_prob + 1e-8).log()) - ) - cost_class = ( - (pos_cost_class - neg_cost_class).unsqueeze(1) - * positive_map.unsqueeze(0) - ).sum(-1) - else: - # Compute the soft-cross entropy between the predicted token alignment and the GT one for each box - cost_class = -(out_prob.unsqueeze(1) * positive_map.unsqueeze(0)).sum( - -1 - ) - else: - # In this case we are doing a "standard" cross entropy - tgt_ids = batched_targets["labels"] - assert len(tgt_bbox) == len(tgt_ids) - - if self.focal_loss: - alpha = self.focal_alpha - gamma = self.focal_gamma - neg_cost_class = ( - (1 - alpha) * (out_prob**gamma) * (-(1 - out_prob + 1e-8).log()) - ) - pos_cost_class = ( - alpha * ((1 - out_prob) ** gamma) * (-(out_prob + 1e-8).log()) - ) - cost_class = pos_cost_class[:, tgt_ids] - neg_cost_class[:, tgt_ids] - else: - # Compute the classification cost. Contrary to the loss, we don't use the NLL, - # but approximate it in 1 - proba[target class]. - # The 1 is a constant that doesn't change the matching, it can be omitted. - cost_class = -out_prob[:, tgt_ids] - - # Compute the L1 cost between boxes - cost_bbox = torch.cdist(out_bbox, tgt_bbox, p=1) - assert cost_class.shape == cost_bbox.shape - - # Compute the giou cost betwen boxes - cost_giou = -generalized_box_iou( - box_cxcywh_to_xyxy(out_bbox), box_cxcywh_to_xyxy(tgt_bbox) - ) - - # Final cost matrix - C = ( - self.cost_bbox * cost_bbox - + self.cost_class * cost_class - + self.cost_giou * cost_giou - ) - C = C.view(bs, num_queries, -1).cpu().numpy() - - sizes = torch.cumsum(batched_targets["num_boxes"], -1)[:-1] - costs = [c[i] for i, c in enumerate(np.split(C, sizes.cpu().numpy(), axis=-1))] - indices = [_do_matching(c) for c in costs] - batch_idx = torch.as_tensor( - sum([[i] * len(src) for i, src in enumerate(indices)], []), dtype=torch.long - ) - src_idx = torch.from_numpy(np.concatenate(indices)).long() - return batch_idx, src_idx - - -class BinaryHungarianMatcher(nn.Module): - """This class computes an assignment between the targets and the predictions of the network - - For efficiency reasons, the targets don't include the no_object. Because of this, in general, - there are more predictions than targets. In this case, we do a 1-to-1 matching of the best predictions, - while the others are un-matched (and thus treated as non-objects). - """ - - def __init__( - self, - cost_class: float = 1, - cost_bbox: float = 1, - cost_giou: float = 1, - ): - """Creates the matcher - - Params: - cost_class: This is the relative weight of the classification error in the matching cost - cost_bbox: This is the relative weight of the L1 error of the bounding box coordinates in the matching cost - cost_giou: This is the relative weight of the giou loss of the bounding box in the matching cost - """ - super().__init__() - self.cost_class = cost_class - self.cost_bbox = cost_bbox - self.cost_giou = cost_giou - self.norm = nn.Sigmoid() - assert ( - cost_class != 0 or cost_bbox != 0 or cost_giou != 0 - ), "all costs cant be 0" - - @torch.no_grad() - def forward(self, outputs, batched_targets, repeats=0, repeat_batch=1): - """Performs the matching - - Params: - outputs: This is a dict that contains at least these entries: - "pred_logits": Tensor of dim [batch_size, num_queries, num_classes] with the classification logits - "pred_boxes": Tensor of dim [batch_size, num_queries, 4] with the predicted box coordinates - - targets: This is a list of targets (len(targets) = batch_size), where each target is a dict containing: - "labels": Tensor of dim [num_target_boxes] (where num_target_boxes is the number of ground-truth - objects in the target) containing the class labels - "boxes": Tensor of dim [num_target_boxes, 4] containing the target box coordinates - - Returns: - A list of size batch_size, containing tuples of (index_i, index_j) where: - - index_i is the indices of the selected predictions (in order) - - index_j is the indices of the corresponding selected targets (in order) - For each batch element, it holds: - len(index_i) = len(index_j) = min(num_queries, num_target_boxes) - """ - if repeat_batch != 1: - raise NotImplementedError("please use BinaryHungarianMatcherV2 instead") - - bs, num_queries = outputs["pred_logits"].shape[:2] - - # We flatten to compute the cost matrices in a batch - out_prob = self.norm(outputs["pred_logits"].flatten(0, 1)).squeeze( - -1 - ) # [batch_size * num_queries] - out_bbox = outputs["pred_boxes"].flatten(0, 1) # [batch_size * num_queries, 4] - - # Also concat the target labels and boxes - tgt_bbox = batched_targets["boxes"] - - # Compute the L1 cost between boxes - cost_bbox = torch.cdist(out_bbox, tgt_bbox, p=1) - - cost_class = -out_prob.unsqueeze(-1).expand_as(cost_bbox) - - assert cost_class.shape == cost_bbox.shape - - # Compute the giou cost betwen boxes - cost_giou = -generalized_box_iou( - box_cxcywh_to_xyxy(out_bbox), box_cxcywh_to_xyxy(tgt_bbox) - ) - - # Final cost matrix - C = ( - self.cost_bbox * cost_bbox - + self.cost_class * cost_class - + self.cost_giou * cost_giou - ) - C = C.view(bs, num_queries, -1).cpu().numpy() - - sizes = torch.cumsum(batched_targets["num_boxes"], -1)[:-1] - costs = [c[i] for i, c in enumerate(np.split(C, sizes.cpu().numpy(), axis=-1))] - return_tgt_indices = False - for c in costs: - n_targ = c.shape[1] - if repeats > 1: - n_targ *= repeats - if c.shape[0] < n_targ: - return_tgt_indices = True - break - if return_tgt_indices: - indices, tgt_indices = zip( - *( - _do_matching( - c, repeats=repeats, return_tgt_indices=return_tgt_indices - ) - for c in costs - ) - ) - tgt_indices = list(tgt_indices) - for i in range(1, len(tgt_indices)): - tgt_indices[i] += sizes[i - 1].item() - tgt_idx = torch.from_numpy(np.concatenate(tgt_indices)).long() - else: - indices = [_do_matching(c, repeats=repeats) for c in costs] - tgt_idx = None - - batch_idx = torch.as_tensor( - sum([[i] * len(src) for i, src in enumerate(indices)], []), dtype=torch.long - ) - src_idx = torch.from_numpy(np.concatenate(indices)).long() - return batch_idx, src_idx, tgt_idx - - -class BinaryFocalHungarianMatcher(nn.Module): - """This class computes an assignment between the targets and the predictions of the network - - For efficiency reasons, the targets don't include the no_object. Because of this, in general, - there are more predictions than targets. In this case, we do a 1-to-1 matching of the best predictions, - while the others are un-matched (and thus treated as non-objects). - """ - - def __init__( - self, - cost_class: float = 1, - cost_bbox: float = 1, - cost_giou: float = 1, - alpha: float = 0.25, - gamma: float = 2.0, - stable: bool = False, - ): - """Creates the matcher - - Params: - cost_class: This is the relative weight of the classification error in the matching cost - cost_bbox: This is the relative weight of the L1 error of the bounding box coordinates in the matching cost - cost_giou: This is the relative weight of the giou loss of the bounding box in the matching cost - """ - super().__init__() - self.cost_class = cost_class - self.cost_bbox = cost_bbox - self.cost_giou = cost_giou - self.norm = nn.Sigmoid() - self.alpha = alpha - self.gamma = gamma - self.stable = stable - assert ( - cost_class != 0 or cost_bbox != 0 or cost_giou != 0 - ), "all costs cant be 0" - - @torch.no_grad() - def forward(self, outputs, batched_targets, repeats=1, repeat_batch=1): - """Performs the matching - - Params: - outputs: This is a dict that contains at least these entries: - "pred_logits": Tensor of dim [batch_size, num_queries, num_classes] with the classification logits - "pred_boxes": Tensor of dim [batch_size, num_queries, 4] with the predicted box coordinates - - targets: This is a list of targets (len(targets) = batch_size), where each target is a dict containing: - "labels": Tensor of dim [num_target_boxes] (where num_target_boxes is the number of ground-truth - objects in the target) containing the class labels - "boxes": Tensor of dim [num_target_boxes, 4] containing the target box coordinates - - Returns: - A list of size batch_size, containing tuples of (index_i, index_j) where: - - index_i is the indices of the selected predictions (in order) - - index_j is the indices of the corresponding selected targets (in order) - For each batch element, it holds: - len(index_i) = len(index_j) = min(num_queries, num_target_boxes) - """ - if repeat_batch != 1: - raise NotImplementedError("please use BinaryHungarianMatcherV2 instead") - - bs, num_queries = outputs["pred_logits"].shape[:2] - - # We flatten to compute the cost matrices in a batch - out_score = outputs["pred_logits"].flatten(0, 1).squeeze(-1) - out_prob = self.norm(out_score) # [batch_size * num_queries] - out_bbox = outputs["pred_boxes"].flatten(0, 1) # [batch_size * num_queries, 4] - - # Also concat the target labels and boxes - tgt_bbox = batched_targets["boxes"] - - # Compute the L1 cost between boxes - cost_bbox = torch.cdist(out_bbox, tgt_bbox, p=1) - - # Compute the giou cost betwen boxes - cost_giou = -generalized_box_iou( - box_cxcywh_to_xyxy(out_bbox), box_cxcywh_to_xyxy(tgt_bbox) - ) - - # cost_class = -out_prob.unsqueeze(-1).expand_as(cost_bbox) - if self.stable: - rescaled_giou = (-cost_giou + 1) / 2 - out_prob = out_prob.unsqueeze(-1).expand_as(cost_bbox) * rescaled_giou - cost_class = -self.alpha * (1 - out_prob) ** self.gamma * torch.log( - out_prob - ) + (1 - self.alpha) * out_prob**self.gamma * torch.log(1 - out_prob) - else: - # directly computing log sigmoid (more numerically stable) - log_out_prob = torch.nn.functional.logsigmoid(out_score) - log_one_minus_out_prob = torch.nn.functional.logsigmoid(-out_score) - cost_class = ( - -self.alpha * (1 - out_prob) ** self.gamma * log_out_prob - + (1 - self.alpha) * out_prob**self.gamma * log_one_minus_out_prob - ) - if not self.stable: - cost_class = cost_class.unsqueeze(-1).expand_as(cost_bbox) - - assert cost_class.shape == cost_bbox.shape - - # Final cost matrix - C = ( - self.cost_bbox * cost_bbox - + self.cost_class * cost_class - + self.cost_giou * cost_giou - ) - C = C.view(bs, num_queries, -1).cpu().numpy() - - sizes = torch.cumsum(batched_targets["num_boxes"], -1)[:-1] - costs = [c[i] for i, c in enumerate(np.split(C, sizes.cpu().numpy(), axis=-1))] - return_tgt_indices = False - for c in costs: - n_targ = c.shape[1] - if repeats > 1: - n_targ *= repeats - if c.shape[0] < n_targ: - return_tgt_indices = True - break - if return_tgt_indices: - indices, tgt_indices = zip( - *( - _do_matching( - c, repeats=repeats, return_tgt_indices=return_tgt_indices - ) - for c in costs - ) - ) - tgt_indices = list(tgt_indices) - for i in range(1, len(tgt_indices)): - tgt_indices[i] += sizes[i - 1].item() - tgt_idx = torch.from_numpy(np.concatenate(tgt_indices)).long() - else: - indices = [_do_matching(c, repeats=repeats) for c in costs] - tgt_idx = None - - batch_idx = torch.as_tensor( - sum([[i] * len(src) for i, src in enumerate(indices)], []), dtype=torch.long - ) - src_idx = torch.from_numpy(np.concatenate(indices)).long() - return batch_idx, src_idx, tgt_idx - - -class BinaryHungarianMatcherV2(nn.Module): - """ - This class computes an assignment between the targets and the predictions - of the network - - For efficiency reasons, the targets don't include the no_object. Because of - this, in general, there are more predictions than targets. In this case, we - do a 1-to-1 matching of the best predictions, while the others are - un-matched (and thus treated as non-objects). - - This is a more efficient implementation of BinaryHungarianMatcher. - """ - - def __init__( - self, - cost_class: float = 1, - cost_bbox: float = 1, - cost_giou: float = 1, - focal: bool = False, - alpha: float = 0.25, - gamma: float = 2.0, - stable: bool = False, - remove_samples_with_0_gt: bool = True, - ): - """ - Creates the matcher - - Params: - - cost_class: Relative weight of the classification error in the - matching cost - - cost_bbox: Relative weight of the L1 error of the bounding box - coordinates in the matching cost - - cost_giou: This is the relative weight of the giou loss of the - bounding box in the matching cost - """ - super().__init__() - self.cost_class = cost_class - self.cost_bbox = cost_bbox - self.cost_giou = cost_giou - self.norm = nn.Sigmoid() - assert ( - cost_class != 0 or cost_bbox != 0 or cost_giou != 0 - ), "all costs cant be 0" - self.focal = focal - if focal: - self.alpha = alpha - self.gamma = gamma - self.stable = stable - self.remove_samples_with_0_gt = remove_samples_with_0_gt - - @torch.no_grad() - def forward( - self, - outputs, - batched_targets, - repeats=1, - repeat_batch=1, - out_is_valid=None, - target_is_valid_padded=None, - ): - """ - Performs the matching. The inputs and outputs are the same as - BinaryHungarianMatcher.forward, except for the optional cached_padded - flag and the optional "_boxes_padded" entry of batched_targets. - - Inputs: - - outputs: A dict with the following keys: - - "pred_logits": Tensor of shape (batch_size, num_queries, 1) with - classification logits - - "pred_boxes": Tensor of shape (batch_size, num_queries, 4) with - predicted box coordinates in cxcywh format. - - batched_targets: A dict of targets. There may be a variable number of - targets per batch entry; suppose that there are T_b targets for batch - entry 0 <= b < batch_size. It should have the following keys: - - "boxes": Tensor of shape (sum_b T_b, 4) giving ground-truth boxes - in cxcywh format for all batch entries packed into a single tensor - - "num_boxes": int64 Tensor of shape (batch_size,) giving the number - of ground-truth boxes per batch entry: num_boxes[b] = T_b - - "_boxes_padded": Tensor of shape (batch_size, max_b T_b, 4) giving - a padded version of ground-truth boxes. If this is not present then - it will be computed from batched_targets["boxes"] instead, but - caching it here can improve performance for repeated calls with the - same targets. - - out_is_valid: If not None, it should be a boolean tensor of shape - (batch_size, num_queries) indicating which predictions are valid. - Invalid predictions are ignored during matching and won't appear in - the output indices. - - target_is_valid_padded: If not None, it should be a boolean tensor of - shape (batch_size, max_num_gt_boxes) in padded format indicating - which GT boxes are valid. Invalid GT boxes are ignored during matching - and won't appear in the output indices. - - Returns: - A list of size batch_size, containing tuples of (idx_i, idx_j): - - idx_i is the indices of the selected predictions (in order) - - idx_j is the indices of the corresponding selected targets - (in order) - For each batch element, it holds: - len(index_i) = len(index_j) - = min(num_queries, num_target_boxes) - """ - _, num_queries = outputs["pred_logits"].shape[:2] - - out_score = outputs["pred_logits"].squeeze(-1) # (B, Q) - out_bbox = outputs["pred_boxes"] # (B, Q, 4)) - - device = out_score.device - - num_boxes = batched_targets["num_boxes"].cpu() - # Get a padded version of target boxes (as precomputed in the collator). - # It should work for both repeat==1 (o2o) and repeat>1 (o2m) matching. - tgt_bbox = batched_targets["boxes_padded"] - if self.remove_samples_with_0_gt: - # keep only samples w/ at least 1 GT box in targets (num_boxes and tgt_bbox) - batch_keep = num_boxes > 0 - num_boxes = num_boxes[batch_keep] - tgt_bbox = tgt_bbox[batch_keep] - if target_is_valid_padded is not None: - target_is_valid_padded = target_is_valid_padded[batch_keep] - # Repeat the targets (for the case of batched aux outputs in the matcher) - if repeat_batch > 1: - # In this case, out_prob and out_bbox will be a concatenation of - # both final and auxiliary outputs, so we also repeat the targets - num_boxes = num_boxes.repeat(repeat_batch) - tgt_bbox = tgt_bbox.repeat(repeat_batch, 1, 1) - if target_is_valid_padded is not None: - target_is_valid_padded = target_is_valid_padded.repeat(repeat_batch, 1) - - # keep only samples w/ at least 1 GT box in outputs - if self.remove_samples_with_0_gt: - if repeat_batch > 1: - batch_keep = batch_keep.repeat(repeat_batch) - out_score = out_score[batch_keep] - out_bbox = out_bbox[batch_keep] - if out_is_valid is not None: - out_is_valid = out_is_valid[batch_keep] - assert out_bbox.shape[0] == tgt_bbox.shape[0] - assert out_bbox.shape[0] == num_boxes.shape[0] - - # Compute the L1 cost between boxes - cost_bbox = torch.cdist(out_bbox, tgt_bbox, p=1) - - # Compute the giou cost betwen boxes - cost_giou = -generalized_box_iou( - box_cxcywh_to_xyxy(out_bbox), box_cxcywh_to_xyxy(tgt_bbox) - ) - - out_prob = self.norm(out_score) - if not self.focal: - cost_class = -out_prob.unsqueeze(-1).expand_as(cost_bbox) - else: - if self.stable: - rescaled_giou = (-cost_giou + 1) / 2 - out_prob = out_prob.unsqueeze(-1).expand_as(cost_bbox) * rescaled_giou - cost_class = -self.alpha * (1 - out_prob) ** self.gamma * torch.log( - out_prob - ) + (1 - self.alpha) * out_prob**self.gamma * torch.log(1 - out_prob) - else: - # directly computing log sigmoid (more numerically stable) - log_out_prob = torch.nn.functional.logsigmoid(out_score) - log_one_minus_out_prob = torch.nn.functional.logsigmoid(-out_score) - cost_class = ( - -self.alpha * (1 - out_prob) ** self.gamma * log_out_prob - + (1 - self.alpha) * out_prob**self.gamma * log_one_minus_out_prob - ) - if not self.stable: - cost_class = cost_class.unsqueeze(-1).expand_as(cost_bbox) - - assert cost_class.shape == cost_bbox.shape - - # Final cost matrix - C = ( - self.cost_bbox * cost_bbox - + self.cost_class * cost_class - + self.cost_giou * cost_giou - ) - # assign a very high cost (1e9) to invalid outputs and targets, so that we can - # filter them out (in `_do_matching`) from bipartite matching results - do_filtering = out_is_valid is not None or target_is_valid_padded is not None - if out_is_valid is not None: - C = torch.where(out_is_valid[:, :, None], C, 1e9) - if target_is_valid_padded is not None: - C = torch.where(target_is_valid_padded[:, None, :], C, 1e9) - C = C.cpu().numpy() - costs = [C[i, :, :s] for i, s in enumerate(num_boxes.tolist())] - return_tgt_indices = ( - do_filtering or torch.any(num_queries < num_boxes * max(repeats, 1)).item() - ) - if len(costs) == 0: - # We have size 0 in the batch dimension, so we return empty matching indices - # (note that this can happen due to `remove_samples_with_0_gt=True` even if - # the original input batch size is not 0, when all queries have empty GTs). - indices = [] - tgt_idx = torch.zeros(0).long().to(device) if return_tgt_indices else None - elif return_tgt_indices: - indices, tgt_indices = zip( - *( - _do_matching( - c, - repeats=repeats, - return_tgt_indices=return_tgt_indices, - do_filtering=do_filtering, - ) - for c in costs - ) - ) - tgt_indices = list(tgt_indices) - sizes = torch.cumsum(num_boxes, -1)[:-1] - for i in range(1, len(tgt_indices)): - tgt_indices[i] += sizes[i - 1].item() - tgt_idx = torch.from_numpy(np.concatenate(tgt_indices)).long().to(device) - else: - indices = [ - _do_matching(c, repeats=repeats, do_filtering=do_filtering) - for c in costs - ] - tgt_idx = None - - if self.remove_samples_with_0_gt: - kept_inds = batch_keep.nonzero().squeeze(1) - batch_idx = torch.as_tensor( - sum([[kept_inds[i]] * len(src) for i, src in enumerate(indices)], []), - dtype=torch.long, - device=device, - ) - else: - batch_idx = torch.as_tensor( - sum([[i] * len(src) for i, src in enumerate(indices)], []), - dtype=torch.long, - device=device, - ) - - # indices could be an empty list (since we remove samples w/ 0 GT boxes) - if len(indices) > 0: - src_idx = torch.from_numpy(np.concatenate(indices)).long().to(device) - else: - src_idx = torch.empty(0, dtype=torch.long, device=device) - return batch_idx, src_idx, tgt_idx - - -class BinaryOneToManyMatcher(nn.Module): - """ - This class computes a greedy assignment between the targets and the predictions of the network. - In this formulation, several predictions can be assigned to each target, but each prediction can be assigned to - at most one target. - - See DAC-Detr for details - """ - - def __init__( - self, - alpha: float = 0.3, - threshold: float = 0.4, - topk: int = 6, - ): - """ - Creates the matcher - - Params: - alpha: relative balancing between classification and localization - threshold: threshold used to select positive predictions - topk: number of top scoring predictions to consider - """ - super().__init__() - self.norm = nn.Sigmoid() - self.alpha = alpha - self.threshold = threshold - self.topk = topk - - @torch.no_grad() - def forward( - self, - outputs, - batched_targets, - repeats=1, - repeat_batch=1, - out_is_valid=None, - target_is_valid_padded=None, - ): - """ - Performs the matching. The inputs and outputs are the same as - BinaryHungarianMatcher.forward - - Inputs: - - outputs: A dict with the following keys: - - "pred_logits": Tensor of shape (batch_size, num_queries, 1) with - classification logits - - "pred_boxes": Tensor of shape (batch_size, num_queries, 4) with - predicted box coordinates in cxcywh format. - - batched_targets: A dict of targets. There may be a variable number of - targets per batch entry; suppose that there are T_b targets for batch - entry 0 <= b < batch_size. It should have the following keys: - - "num_boxes": int64 Tensor of shape (batch_size,) giving the number - of ground-truth boxes per batch entry: num_boxes[b] = T_b - - "_boxes_padded": Tensor of shape (batch_size, max_b T_b, 4) giving - a padded version of ground-truth boxes. If this is not present then - it will be computed from batched_targets["boxes"] instead, but - caching it here can improve performance for repeated calls with the - same targets. - - out_is_valid: If not None, it should be a boolean tensor of shape - (batch_size, num_queries) indicating which predictions are valid. - Invalid predictions are ignored during matching and won't appear in - the output indices. - - target_is_valid_padded: If not None, it should be a boolean tensor of - shape (batch_size, max_num_gt_boxes) in padded format indicating - which GT boxes are valid. Invalid GT boxes are ignored during matching - and won't appear in the output indices. - Returns: - A list of size batch_size, containing tuples of (idx_i, idx_j): - - idx_i is the indices of the selected predictions (in order) - - idx_j is the indices of the corresponding selected targets - (in order) - For each batch element, it holds: - len(index_i) = len(index_j) - = min(num_queries, num_target_boxes) - """ - assert repeats <= 1 and repeat_batch <= 1 - bs, num_queries = outputs["pred_logits"].shape[:2] - - out_prob = self.norm(outputs["pred_logits"]).squeeze(-1) # (B, Q) - out_bbox = outputs["pred_boxes"] # (B, Q, 4)) - - num_boxes = batched_targets["num_boxes"] - - # Get a padded version of target boxes (as precomputed in the collator). - tgt_bbox = batched_targets["boxes_padded"] - assert len(tgt_bbox) == bs - num_targets = tgt_bbox.shape[1] - if num_targets == 0: - return ( - torch.empty(0, dtype=torch.long, device=out_prob.device), - torch.empty(0, dtype=torch.long, device=out_prob.device), - torch.empty(0, dtype=torch.long, device=out_prob.device), - ) - - iou, _ = box_iou(box_cxcywh_to_xyxy(out_bbox), box_cxcywh_to_xyxy(tgt_bbox)) - - assert iou.shape == (bs, num_queries, num_targets) - - # Final cost matrix (higher is better in `C`; this is unlike the case - # of BinaryHungarianMatcherV2 above where lower is better in its `C`) - C = self.alpha * out_prob.unsqueeze(-1) + (1 - self.alpha) * iou - if out_is_valid is not None: - C = torch.where(out_is_valid[:, :, None], C, -1e9) - if target_is_valid_padded is not None: - C = torch.where(target_is_valid_padded[:, None, :], C, -1e9) - - # Selecting topk predictions - matches = C > torch.quantile( - C, 1 - self.topk / num_queries, dim=1, keepdim=True - ) - - # Selecting predictions above threshold - matches = matches & (C > self.threshold) - if out_is_valid is not None: - matches = matches & out_is_valid[:, :, None] - if target_is_valid_padded is not None: - matches = matches & target_is_valid_padded[:, None, :] - - # Removing padding - matches = matches & ( - torch.arange(0, num_targets, device=num_boxes.device)[None] - < num_boxes[:, None] - ).unsqueeze(1) - - batch_idx, src_idx, tgt_idx = torch.nonzero(matches, as_tuple=True) - - cum_num_boxes = torch.cat( - [ - torch.zeros(1, dtype=num_boxes.dtype, device=num_boxes.device), - num_boxes.cumsum(-1)[:-1], - ] - ) - tgt_idx += cum_num_boxes[batch_idx] - - return batch_idx, src_idx, tgt_idx diff --git a/train/nms_helper.py b/train/nms_helper.py deleted file mode 100644 index cd5b6dc..0000000 --- a/train/nms_helper.py +++ /dev/null @@ -1,306 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved -import warnings -from typing import Dict, List - -import numpy as np - -# Check if Numba is available -HAS_NUMBA = False -try: - import numba as nb - - HAS_NUMBA = True -except ImportError: - warnings.warn( - "Numba not found. Using slower pure Python implementations.", UserWarning - ) - - -# -------------------- Helper Functions -------------------- -def is_zero_box(bbox: list) -> bool: - """Check if bounding box is invalid""" - if bbox is None: - return True - return all(x <= 0 for x in bbox[:4]) or len(bbox) < 4 - - -def convert_bbox_format(bbox: list) -> List[float]: - """Convert bbox from (x,y,w,h) to (x1,y1,x2,y2)""" - x, y, w, h = bbox - return [x, y, x + w, y + h] - - -# -------------------- Track-level NMS -------------------- -def process_track_level_nms(video_groups: Dict, nms_threshold: float) -> Dict: - """Apply track-level NMS to all videos""" - for video_id, tracks in video_groups.items(): - track_detections = [] - - # Process tracks - for track_idx, track in enumerate(tracks): - if not track["bboxes"]: - continue - - converted_bboxes = [] - valid_frames = [] - for bbox in track["bboxes"]: - if bbox and not is_zero_box(bbox): - converted_bboxes.append(convert_bbox_format(bbox)) - valid_frames.append(True) - else: - converted_bboxes.append([np.nan] * 4) - valid_frames.append(False) - - if any(valid_frames): - track_detections.append( - { - "track_idx": track_idx, - "bboxes": np.array(converted_bboxes, dtype=np.float32), - "score": track["score"], - } - ) - - # Apply NMS - if track_detections: - scores = np.array([d["score"] for d in track_detections], dtype=np.float32) - keep = apply_track_nms(track_detections, scores, nms_threshold) - - # Suppress non-kept tracks - for idx, track in enumerate(track_detections): - if idx not in keep: - tracks[track["track_idx"]]["bboxes"] = [None] * len(track["bboxes"]) - - return video_groups - - -# -------------------- Frame-level NMS -------------------- -def process_frame_level_nms(video_groups: Dict, nms_threshold: float) -> Dict: - """Apply frame-level NMS to all videos""" - for video_id, tracks in video_groups.items(): - if not tracks: - continue - - num_frames = len(tracks[0]["bboxes"]) - - for frame_idx in range(num_frames): - frame_detections = [] - - # Collect valid detections - for track_idx, track in enumerate(tracks): - bbox = track["bboxes"][frame_idx] - if bbox and not is_zero_box(bbox): - frame_detections.append( - { - "track_idx": track_idx, - "bbox": np.array( - convert_bbox_format(bbox), dtype=np.float32 - ), - "score": track["score"], - } - ) - - # Apply NMS - if frame_detections: - bboxes = np.stack([d["bbox"] for d in frame_detections]) - scores = np.array( - [d["score"] for d in frame_detections], dtype=np.float32 - ) - keep = apply_frame_nms(bboxes, scores, nms_threshold) - - # Suppress non-kept detections - for i, d in enumerate(frame_detections): - if i not in keep: - tracks[d["track_idx"]]["bboxes"][frame_idx] = None - - return video_groups - - -# Track-level NMS helpers ------------------------------------------------------ -def compute_track_iou_matrix( - bboxes_stacked: np.ndarray, valid_masks: np.ndarray, areas: np.ndarray -) -> np.ndarray: - """IoU matrix computation for track-level NMS with fallback to pure Python""" - num_tracks = bboxes_stacked.shape[0] - iou_matrix = np.zeros((num_tracks, num_tracks), dtype=np.float32) - if HAS_NUMBA: - iou_matrix = _compute_track_iou_matrix_numba(bboxes_stacked, valid_masks, areas) - else: - # Pure Python implementation - for i in range(num_tracks): - for j in range(i + 1, num_tracks): - valid_ij = valid_masks[i] & valid_masks[j] - if not valid_ij.any(): - continue - bboxes_i = bboxes_stacked[i, valid_ij] - bboxes_j = bboxes_stacked[j, valid_ij] - area_i = areas[i, valid_ij] - area_j = areas[j, valid_ij] - inter_total = 0.0 - union_total = 0.0 - for k in range(bboxes_i.shape[0]): - x1 = max(bboxes_i[k, 0], bboxes_j[k, 0]) - y1 = max(bboxes_i[k, 1], bboxes_j[k, 1]) - x2 = min(bboxes_i[k, 2], bboxes_j[k, 2]) - y2 = min(bboxes_i[k, 3], bboxes_j[k, 3]) - inter = max(0, x2 - x1) * max(0, y2 - y1) - union = area_i[k] + area_j[k] - inter - inter_total += inter - union_total += union - if union_total > 0: - iou_matrix[i, j] = inter_total / union_total - iou_matrix[j, i] = iou_matrix[i, j] - return iou_matrix - - -if HAS_NUMBA: - - @nb.jit(nopython=True, parallel=True) - def _compute_track_iou_matrix_numba(bboxes_stacked, valid_masks, areas): - """Numba-optimized IoU matrix computation for track-level NMS""" - num_tracks = bboxes_stacked.shape[0] - iou_matrix = np.zeros((num_tracks, num_tracks), dtype=np.float32) - for i in nb.prange(num_tracks): - for j in range(i + 1, num_tracks): - valid_ij = valid_masks[i] & valid_masks[j] - if not valid_ij.any(): - continue - bboxes_i = bboxes_stacked[i, valid_ij] - bboxes_j = bboxes_stacked[j, valid_ij] - area_i = areas[i, valid_ij] - area_j = areas[j, valid_ij] - inter_total = 0.0 - union_total = 0.0 - for k in range(bboxes_i.shape[0]): - x1 = max(bboxes_i[k, 0], bboxes_j[k, 0]) - y1 = max(bboxes_i[k, 1], bboxes_j[k, 1]) - x2 = min(bboxes_i[k, 2], bboxes_j[k, 2]) - y2 = min(bboxes_i[k, 3], bboxes_j[k, 3]) - inter = max(0, x2 - x1) * max(0, y2 - y1) - union = area_i[k] + area_j[k] - inter - inter_total += inter - union_total += union - if union_total > 0: - iou_matrix[i, j] = inter_total / union_total - iou_matrix[j, i] = iou_matrix[i, j] - return iou_matrix - - -def apply_track_nms( - track_detections: List[dict], scores: np.ndarray, nms_threshold: float -) -> List[int]: - """Vectorized track-level NMS implementation""" - if not track_detections: - return [] - bboxes_stacked = np.stack([d["bboxes"] for d in track_detections], axis=0) - valid_masks = ~np.isnan(bboxes_stacked).any(axis=2) - areas = (bboxes_stacked[:, :, 2] - bboxes_stacked[:, :, 0]) * ( - bboxes_stacked[:, :, 3] - bboxes_stacked[:, :, 1] - ) - areas[~valid_masks] = 0 - iou_matrix = compute_track_iou_matrix(bboxes_stacked, valid_masks, areas) - keep = [] - order = np.argsort(-scores) - suppress = np.zeros(len(track_detections), dtype=bool) - for i in range(len(order)): - if not suppress[order[i]]: - keep.append(order[i]) - suppress[order[i:]] = suppress[order[i:]] | ( - iou_matrix[order[i], order[i:]] >= nms_threshold - ) - return keep - - -# Frame-level NMS helpers ------------------------------------------------------ -def compute_frame_ious(bbox: np.ndarray, bboxes: np.ndarray) -> np.ndarray: - """IoU computation for frame-level NMS with fallback to pure Python""" - if HAS_NUMBA: - return _compute_frame_ious_numba(bbox, bboxes) - else: - # Pure Python implementation - ious = np.zeros(len(bboxes), dtype=np.float32) - for i in range(len(bboxes)): - x1 = max(bbox[0], bboxes[i, 0]) - y1 = max(bbox[1], bboxes[i, 1]) - x2 = min(bbox[2], bboxes[i, 2]) - y2 = min(bbox[3], bboxes[i, 3]) - - inter = max(0, x2 - x1) * max(0, y2 - y1) - area1 = (bbox[2] - bbox[0]) * (bbox[3] - bbox[1]) - area2 = (bboxes[i, 2] - bboxes[i, 0]) * (bboxes[i, 3] - bboxes[i, 1]) - union = area1 + area2 - inter - - ious[i] = inter / union if union > 0 else 0.0 - return ious - - -if HAS_NUMBA: - - @nb.jit(nopython=True, parallel=True) - def _compute_frame_ious_numba(bbox, bboxes): - """Numba-optimized IoU computation""" - ious = np.zeros(len(bboxes), dtype=np.float32) - for i in nb.prange(len(bboxes)): - x1 = max(bbox[0], bboxes[i, 0]) - y1 = max(bbox[1], bboxes[i, 1]) - x2 = min(bbox[2], bboxes[i, 2]) - y2 = min(bbox[3], bboxes[i, 3]) - - inter = max(0, x2 - x1) * max(0, y2 - y1) - area1 = (bbox[2] - bbox[0]) * (bbox[3] - bbox[1]) - area2 = (bboxes[i, 2] - bboxes[i, 0]) * (bboxes[i, 3] - bboxes[i, 1]) - union = area1 + area2 - inter - - ious[i] = inter / union if union > 0 else 0.0 - return ious - - -def apply_frame_nms( - bboxes: np.ndarray, scores: np.ndarray, nms_threshold: float -) -> List[int]: - """Frame-level NMS implementation with fallback to pure Python""" - if HAS_NUMBA: - return _apply_frame_nms_numba(bboxes, scores, nms_threshold) - else: - # Pure Python implementation - order = np.argsort(-scores) - keep = [] - suppress = np.zeros(len(bboxes), dtype=bool) - - for i in range(len(order)): - if not suppress[order[i]]: - keep.append(order[i]) - current_bbox = bboxes[order[i]] - - remaining_bboxes = bboxes[order[i + 1 :]] - if len(remaining_bboxes) > 0: # Check if there are any remaining boxes - ious = compute_frame_ious(current_bbox, remaining_bboxes) - suppress[order[i + 1 :]] = suppress[order[i + 1 :]] | ( - ious >= nms_threshold - ) - - return keep - - -if HAS_NUMBA: - - @nb.jit(nopython=True) - def _apply_frame_nms_numba(bboxes, scores, nms_threshold): - """Numba-optimized NMS implementation""" - order = np.argsort(-scores) - keep = [] - suppress = np.zeros(len(bboxes), dtype=nb.boolean) - - for i in range(len(order)): - if not suppress[order[i]]: - keep.append(order[i]) - current_bbox = bboxes[order[i]] - - if i + 1 < len(order): # Check bounds - ious = _compute_frame_ious_numba( - current_bbox, bboxes[order[i + 1 :]] - ) - suppress[order[i + 1 :]] = suppress[order[i + 1 :]] | ( - ious >= nms_threshold - ) - - return keep diff --git a/train/optim/__init__.py b/train/optim/__init__.py deleted file mode 100644 index 46d37d2..0000000 --- a/train/optim/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved diff --git a/train/optim/optimizer.py b/train/optim/optimizer.py deleted file mode 100644 index d401b98..0000000 --- a/train/optim/optimizer.py +++ /dev/null @@ -1,498 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved - -import fnmatch -import inspect -import itertools -import logging -import types -from typing import ( - Any, - Callable, - Dict, - Iterable, - List, - Mapping, - Optional, - Set, - Tuple, - Type, - Union, -) - -import hydra - -import torch -import torch.nn as nn -from omegaconf import DictConfig -from torch import Tensor - - -class Optimizer: - def __init__(self, optimizer, schedulers=None) -> None: - self.optimizer = optimizer - self.schedulers = schedulers - self._validate_optimizer_schedulers() - self.step_schedulers(0.0, 0) - - def _validate_optimizer_schedulers(self): - if self.schedulers is None: - return - for _, set_of_schedulers in enumerate(self.schedulers): - for option, _ in set_of_schedulers.items(): - assert option in self.optimizer.defaults, ( - "Optimizer option " - f"{option} not found in {self.optimizer}. Valid options are " - f"{self.optimizer.defaults.keys()}" - ) - - def step_schedulers(self, where: float, step: int) -> None: - if self.schedulers is None: - return - for i, param_group in enumerate(self.optimizer.param_groups): - for option, scheduler in self.schedulers[i].items(): - if "step" in inspect.signature(scheduler.__call__).parameters: - new_value = scheduler(step=step, where=where) - elif ( - hasattr(scheduler, "scheduler") - and "step" - in inspect.signature(scheduler.scheduler.__call__).parameters - ): - # To handle ValueScaler wrappers - new_value = scheduler(step=step, where=where) - else: - new_value = scheduler(where) - param_group[option] = new_value - - def step(self, where, step, closure=None): - self.step_schedulers(where, step) - return self.optimizer.step(closure) - - def zero_grad(self, *args, **kwargs): - return self.optimizer.zero_grad(*args, **kwargs) - - -def set_default_parameters( - scheduler_cfgs: List[DictConfig], all_parameter_names: Set[str] -) -> None: - """Set up the "default" scheduler with the right parameters. - - Args: - scheduler_cgfs: A list of scheduler configs, where each scheduler also - specifies which parameters it applies to, based on the names of parameters - or the class of the modules. At most one scheduler is allowed to skip this - specification, which is used as a "default" specification for any remaining - parameters. - all_parameter_names: Names of all the parameters to consider. - """ - constraints = [ - scheduler_cfg.parameter_names - for scheduler_cfg in scheduler_cfgs - if scheduler_cfg.parameter_names is not None - ] - if len(constraints) == 0: - default_params = set(all_parameter_names) - else: - default_params = all_parameter_names - set.union(*constraints) - default_count = 0 - for scheduler_cfg in scheduler_cfgs: - if scheduler_cfg.parameter_names is None: - scheduler_cfg.parameter_names = default_params - default_count += 1 - assert default_count <= 1, "Only one scheduler per option can be default" - if default_count == 0: - # No default scheduler specified, add a default, but without any scheduler - # for that option - scheduler_cfgs.append({"parameter_names": default_params}) - - -def name_constraints_to_parameters( - param_constraints: List[Set[str]], named_parameters: Dict[str, Tensor] -) -> List[torch.nn.Parameter]: - """Return parameters which match the intersection of parameter constraints. - - Note that this returns the parameters themselves, not their names. - - Args: - param_constraints: A list, with each element being a set of allowed parameters. - named_parameters: Mapping from a parameter name to the parameter itself. - - Returns: - A list containing the parameters which overlap with _each_ constraint set from - param_constraints. - """ - matching_names = set.intersection(*param_constraints) - return [value for name, value in named_parameters.items() if name in matching_names] - - -def map_scheduler_cfgs_to_param_groups( - all_scheduler_cfgs: Iterable[List[Dict]], - named_parameters: Dict[str, Tensor], -) -> Tuple[List[Dict[Any, Any]], List[Dict[str, List[torch.nn.Parameter]]]]: - """Produce parameter groups corresponding to all the scheduler configs. - - Takes all the scheduler configs, each of which applies to a specific optimizer - option (like "lr" or "weight_decay") and has a set of parameter names which it - applies to, and produces a final set of param groups where each param group - covers all the options which apply to a particular set of parameters. - - Args: - all_scheduler_cfgs: All the scheduler configs covering every option. - named_parameters: Mapping from a parameter name to the parameter itself. - Returns: - Tuple of lists of schedulers and param_groups, where schedulers[i] - applies to param_groups[i]. - """ - - scheduler_cfgs_per_param_group = itertools.product(*all_scheduler_cfgs) - schedulers = [] - param_groups = [] - for scheduler_cfgs in scheduler_cfgs_per_param_group: - param_constraints = [ - scheduler_cfg["parameter_names"] for scheduler_cfg in scheduler_cfgs - ] - matching_parameters = name_constraints_to_parameters( - param_constraints, named_parameters - ) - if len(matching_parameters) == 0: # If no overlap of parameters, skip - continue - schedulers_for_group = { - scheduler_cfg["option"]: scheduler_cfg["scheduler"] - for scheduler_cfg in scheduler_cfgs - if "option" in scheduler_cfg - } - schedulers.append(schedulers_for_group) - param_groups.append({"params": matching_parameters}) - return schedulers, param_groups - - -def validate_param_group_params(param_groups: List[Dict], model: nn.Module): - """Check that the param groups are non-overlapping and cover all the parameters. - - Args: - param_groups: List of all param groups - model: Model to validate against. The check ensures that all the model - parameters are part of param_groups - """ - for pg in param_groups: - # no param should be repeated within a group - assert len(pg["params"]) == len(set(pg["params"])) - parameters = [set(param_group["params"]) for param_group in param_groups] - model_parameters = {parameter for _, parameter in model.named_parameters()} - for p1, p2 in itertools.permutations(parameters, 2): - assert p1.isdisjoint(p2), "Scheduler generated param_groups should be disjoint" - assert set.union(*parameters) == model_parameters, ( - "Scheduler generated param_groups must include all parameters of the model." - f" Found {len(set.union(*parameters))} params whereas model has" - f" {len(model_parameters)} params" - ) - - -def unix_module_cls_pattern_to_parameter_names( - filter_module_cls_names: List[str], - module_cls_to_param_names: Dict[Type, str], -) -> Union[None, Set[str]]: - """Returns param names which pass the filters specified in filter_module_cls_names. - - Args: - filter_module_cls_names: A list of filter strings containing class names, like - ["torch.nn.LayerNorm", "torch.nn.BatchNorm2d"] - module_cls_to_param_names: Mapping from module classes to the parameter names - they contain. See `get_module_cls_to_param_names`. - """ - if filter_module_cls_names is None: - return set() - allowed_parameter_names = [] - for module_cls_name in filter_module_cls_names: - module_cls = hydra.utils.get_class(module_cls_name) - if module_cls not in module_cls_to_param_names: - raise AssertionError( - f"module_cls_name {module_cls_name} does not " - "match any classes in the model" - ) - matching_parameters = module_cls_to_param_names[module_cls] - assert ( - len(matching_parameters) > 0 - ), f"module_cls_name {module_cls_name} does not contain any parameters in the model" - logging.info( - f"Matches for module_cls_name [{module_cls_name}]: {matching_parameters} " - ) - allowed_parameter_names.append(matching_parameters) - return set.union(*allowed_parameter_names) - - -def unix_param_pattern_to_parameter_names( - filter_param_names: Optional[List[str]], - parameter_names: Dict[str, torch.Tensor], -) -> Union[None, Set[str]]: - """Returns param names which pass the filters specified in filter_param_names. - - Args: - filter_param_names: A list of unix-style filter strings with optional - wildcards, like ["block.2.*", "block.2.linear.weight"] - module_cls_to_param_names: Mapping from module classes to the parameter names - they contain. See `get_module_cls_to_param_names`. - """ - - if filter_param_names is None: - return set() - allowed_parameter_names = [] - for param_name in filter_param_names: - matching_parameters = set(fnmatch.filter(parameter_names, param_name)) - assert ( - len(matching_parameters) >= 1 - ), f"param_name {param_name} does not match any parameters in the model" - logging.info(f"Matches for param_name [{param_name}]: {matching_parameters}") - allowed_parameter_names.append(matching_parameters) - return set.union(*allowed_parameter_names) - - -def _unix_pattern_to_parameter_names( - scheduler_cfg: DictConfig, - parameter_names: Set[str], - module_cls_to_param_names: Dict[Type, str], -) -> Union[None, Set[str]]: - """Returns param names which pass the filters specified in scheduler_cfg. - - Args: - scheduler_cfg: The config for the scheduler - parameter_names: The set of all parameter names which will be filtered - """ - if "param_names" not in scheduler_cfg and "module_cls_names" not in scheduler_cfg: - return None - return unix_param_pattern_to_parameter_names( - scheduler_cfg.get("param_names"), parameter_names - ).union( - unix_module_cls_pattern_to_parameter_names( - scheduler_cfg.get("module_cls_names"), module_cls_to_param_names - ) - ) - - -def get_module_cls_to_param_names( - model: nn.Module, param_allowlist: Set[str] = None -) -> Dict[Type, str]: - """Produce a mapping from all the modules classes to the names of parames they own. - - Only counts a parameter as part of the immediate parent module, i.e. recursive - parents do not count. - - Args: - model: Model to iterate over - param_allowlist: If specified, only these param names will be processed - """ - - module_cls_to_params = {} - for module_name, module in model.named_modules(): - module_cls = type(module) - module_cls_to_params.setdefault(module_cls, set()) - for param_name, _ in module.named_parameters(recurse=False): - full_param_name = get_full_parameter_name(module_name, param_name) - if param_allowlist is None or full_param_name in param_allowlist: - module_cls_to_params[module_cls].add(full_param_name) - return module_cls_to_params - - -def construct_optimizer( - model: torch.nn.Module, - optimizer_conf: Any, - options_conf: Mapping[str, List] = None, - param_group_modifiers_conf: List[Callable] = None, - param_allowlist: Optional[Set[str]] = None, - validate_param_groups=True, -) -> Optimizer: - """ - Constructs a stochastic gradient descent or ADAM (or ADAMw) optimizer - with momentum. i.e, constructs a torch.optim.Optimizer with zero-weight decay - Batchnorm and/or no-update 1-D parameters support, based on the config. - - Supports wrapping the optimizer with Layer-wise Adaptive Rate Scaling - (LARS): https://arxiv.org/abs/1708.03888 - - Args: - model: model to perform stochastic gradient descent - optimization or ADAM optimization. - optimizer_conf: Hydra config consisting a partial torch optimizer like SGD or - ADAM, still missing the params argument which this function provides to - produce the final optimizer - param_group_modifiers_conf: Optional user specified functions which can modify - the final scheduler configs before the optimizer's param groups are built - param_allowlist: The parameters to optimize. Parameters which are not part of - this allowlist will be skipped. - validate_param_groups: If enabled, valides that the produced param_groups don't - overlap and cover all the model parameters. - """ - if param_allowlist is None: - param_allowlist = {name for name, _ in model.named_parameters()} - - named_parameters = { - name: param - for name, param in model.named_parameters() - if name in param_allowlist - } - - if not options_conf: - optimizer = hydra.utils.instantiate(optimizer_conf, named_parameters.values()) - return Optimizer(optimizer) - - all_parameter_names = { - name for name, _ in model.named_parameters() if name in param_allowlist - } - module_cls_to_all_param_names = get_module_cls_to_param_names( - model, param_allowlist - ) - - scheduler_cfgs_per_option = hydra.utils.instantiate(options_conf) - all_scheduler_cfgs = [] - for option, scheduler_cfgs in scheduler_cfgs_per_option.items(): - for config in scheduler_cfgs: - config.option = option - config.parameter_names = _unix_pattern_to_parameter_names( - config, all_parameter_names, module_cls_to_all_param_names - ) - set_default_parameters(scheduler_cfgs, all_parameter_names) - all_scheduler_cfgs.append(scheduler_cfgs) - - if param_group_modifiers_conf: - for custom_param_modifier in param_group_modifiers_conf: - custom_param_modifier = hydra.utils.instantiate(custom_param_modifier) - all_scheduler_cfgs = custom_param_modifier( - scheduler_cfgs=all_scheduler_cfgs, model=model - ) - schedulers, param_groups = map_scheduler_cfgs_to_param_groups( - all_scheduler_cfgs, named_parameters - ) - if validate_param_groups: - validate_param_group_params(param_groups, model) - optimizer = hydra.utils.instantiate(optimizer_conf, param_groups) - return Optimizer(optimizer, schedulers) - - -def get_full_parameter_name(module_name, param_name): - if module_name == "": - return param_name - return f"{module_name}.{param_name}" - - -class GradientClipper: - """ - Gradient clipping utils that works for DDP - """ - - def __init__(self, max_norm: float = 1.0, norm_type: int = 2): - assert isinstance(max_norm, (int, float)) or max_norm is None - self.max_norm = max_norm if max_norm is None else float(max_norm) - self.norm_type = norm_type - - def __call__(self, model: nn.Module): - if self.max_norm is None: - return # no-op - - nn.utils.clip_grad_norm_( - model.parameters(), max_norm=self.max_norm, norm_type=self.norm_type - ) - - -class ValueScaler: - def __init__(self, scheduler, mult_val: float): - self.scheduler = scheduler - self.mult_val = mult_val - - def __call__(self, *args, **kwargs): - val = self.scheduler(*args, **kwargs) - return val * self.mult_val - - -def rgetattr(obj, rattrs: str = None): - """ - Like getattr(), but supports dotted notation for nested objects. - rattrs is a str of form 'attr1.attr2', returns obj.attr1.attr2 - """ - if rattrs is None: - return obj - attrs = rattrs.split(".") - for attr in attrs: - obj = getattr(obj, attr) - return obj - - -def layer_decay_param_modifier( - scheduler_cfgs: List[List[Dict]], - model, - layer_decay_value: float, - layer_decay_min: Optional[float] = None, - apply_to: Optional[str] = None, - overrides: List[Dict] = (), -) -> List[List[Dict]]: - """ - Args - - scheduler_cfgs: a list of omegaconf.ListConfigs. - Each element in the list is a omegaconfg.DictConfig with the following structure - { - "scheduler": - "option": possible options are "lr", "weight_decay" etc. - "parameter_names": Set of str indicating param names that this scheduler applies to - } - - model: a model that implements a method `get_layer_id` that maps layer_name to an integer and - and a method get_num_layers. - Alternatively, use apply_to argument to select a specific component of the model. - - layer_decay_value: float - - layer_decay_min: min val for layer decay - - apply_to: optional arg to select which component of the model to apply the the layer decay modifier to - - overrides: to manually override lr for specific patterns. Is a list of dicts. Each dict, has keys "pattern", "value". - Returns - - scheduler_configs: same structure as the input, elements can be modified - """ - model = rgetattr(model, apply_to) - num_layers = model.get_num_layers() + 1 - layer_decays = [ - layer_decay_value ** (num_layers - i) for i in range(num_layers + 1) - ] - if layer_decay_min is not None: - layer_decays = [max(val, layer_decay_min) for val in layer_decays] - final_scheduler_cfgs = [] - # scheduler_cfgs is a list of lists - for scheduler_cfg_group in scheduler_cfgs: - curr_cfg_group = [] - # scheduler_cfg_group is a list of dictionaries - for scheduler_cfg in scheduler_cfg_group: - if scheduler_cfg["option"] != "lr": - curr_cfg_group.append(scheduler_cfg) - continue - # Need sorted so that the list of parameter names is deterministic and consistent - # across re-runs of this job. Else it was causing issues with loading the optimizer - # state during a job restart - parameter_names = sorted(scheduler_cfg["parameter_names"]) - - # Only want one cfg group per layer - layer_cfg_groups = {} - for param_name in parameter_names: - layer_id = num_layers - this_scale = layer_decays[layer_id] - if param_name.startswith(apply_to): - layer_id = model.get_layer_id(param_name) - this_scale = layer_decays[layer_id] - # Overrides - for override in overrides: - if fnmatch.fnmatchcase(param_name, override["pattern"]): - this_scale = float(override["value"]) - layer_id = override["pattern"] - break - - if layer_id not in layer_cfg_groups: - curr_param = { - "option": scheduler_cfg["option"], - "scheduler": ValueScaler( - scheduler_cfg["scheduler"], this_scale - ), - "parameter_names": {param_name}, - } - else: - curr_param = layer_cfg_groups[layer_id] - curr_param["parameter_names"].add(param_name) - layer_cfg_groups[layer_id] = curr_param - - for layer_cfg in layer_cfg_groups.values(): - curr_cfg_group.append(layer_cfg) - - final_scheduler_cfgs.append(curr_cfg_group) - return final_scheduler_cfgs diff --git a/train/optim/schedulers.py b/train/optim/schedulers.py deleted file mode 100644 index 59da840..0000000 --- a/train/optim/schedulers.py +++ /dev/null @@ -1,41 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved - -import math - - -class InverseSquareRootParamScheduler: - def __init__( - self, - base_lr: float, - warmup_steps: int, - cooldown_steps: int, - timescale: int, - ): - self.base_lr = base_lr - self.warmup_steps = warmup_steps - self.cooldown_steps = cooldown_steps - self.timescale = timescale - - def __call__(self, step: int, where: float): - lr = self.base_lr - - if where > 0: - total_steps = step / where - progress = (step - self.warmup_steps) / float( - total_steps - self.warmup_steps - ) - progress = max(min(progress, 1), 0) - else: - progress = 0 - total_steps = 1 - - shift = self.timescale - self.warmup_steps - if self.warmup_steps < step: - lr = lr / math.sqrt((step + shift) / self.timescale) - - if self.warmup_steps: - lr = lr * min(1.0, step / self.warmup_steps) - if self.cooldown_steps: - lr = lr * min(1.0, (total_steps - step) / self.cooldown_steps) - - return lr diff --git a/train/train.py b/train/train.py deleted file mode 100644 index b3e995c..0000000 --- a/train/train.py +++ /dev/null @@ -1,339 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved - -import logging -import os -import random -import sys -import traceback -from argparse import ArgumentParser -from copy import deepcopy - -import submitit -import torch - -from hydra import compose, initialize_config_module -from hydra.utils import instantiate - -from iopath.common.file_io import g_pathmgr -from omegaconf import OmegaConf - -from sam3.train.utils.train_utils import makedir, register_omegaconf_resolvers -from tqdm import tqdm - - -os.environ["HYDRA_FULL_ERROR"] = "1" - - -class SlurmEvent: - QUEUED = "QUEUED" - START = "START" - FINISH = "FINISH" - JOB_ERROR = "JOB_ERROR" - SLURM_SIGNAL = "SLURM_SIGNAL" - - -def handle_custom_resolving(cfg): - # We'll resolve the config here, so we can catch mistakes early. - # However, we need to pass the un-resolved config to the launcher - # (because DVC resolving needs to be done on the node it will run on) - # First, do a copy without triggering resolving - cfg_resolved = OmegaConf.to_container(cfg, resolve=False) - cfg_resolved = OmegaConf.create(cfg_resolved) - return cfg_resolved - - -def single_proc_run(local_rank, main_port, cfg, world_size): - """Single GPU process""" - os.environ["MASTER_ADDR"] = "localhost" - os.environ["MASTER_PORT"] = str(main_port) - os.environ["RANK"] = str(local_rank) - os.environ["LOCAL_RANK"] = str(local_rank) - os.environ["WORLD_SIZE"] = str(world_size) - try: - register_omegaconf_resolvers() - except Exception as e: - logging.info(e) - - trainer = instantiate(cfg.trainer, _recursive_=False) - trainer.run() - - -def single_node_runner(cfg, main_port: int): - assert cfg.launcher.num_nodes == 1 - # assert cfg.launcher.gpus_per_node == 1 - num_proc = cfg.launcher.gpus_per_node - torch.multiprocessing.set_start_method( - "spawn" - ) # CUDA runtime does not support `fork` - if num_proc == 1: - # directly call single_proc so we can easily set breakpoints - # mp.spawn does not let us set breakpoints - single_proc_run(local_rank=0, main_port=main_port, cfg=cfg, world_size=num_proc) - else: - mp_runner = torch.multiprocessing.start_processes - args = (main_port, cfg, num_proc) - # Note: using "fork" below, "spawn" causes time and error regressions. Using - # spawn changes the default multiprocessing context to spawn, which doesn't - # interact well with the dataloaders (likely due to the use of OpenCV). - mp_runner(single_proc_run, args=args, nprocs=num_proc, start_method="spawn") - - -def format_exception(e: Exception, limit=20): - traceback_str = "".join(traceback.format_tb(e.__traceback__, limit=limit)) - return f"{type(e).__name__}: {e}\nTraceback:\n{traceback_str}" - - -class SubmititRunner(submitit.helpers.Checkpointable): - """A callable which is passed to submitit to launch the jobs.""" - - def __init__(self, port, cfg): - self.cfg = cfg - self.port = port - self.has_setup = False - - def run_trainer(self): - job_env = submitit.JobEnvironment() - # Need to add this again so the hydra.job.set_env PYTHONPATH - # is also set when launching jobs. - add_pythonpath_to_sys_path() - os.environ["MASTER_ADDR"] = job_env.hostnames[0] - os.environ["MASTER_PORT"] = str(self.port) - os.environ["RANK"] = str(job_env.global_rank) - os.environ["LOCAL_RANK"] = str(job_env.local_rank) - os.environ["WORLD_SIZE"] = str(job_env.num_tasks) - - register_omegaconf_resolvers() - cfg_resolved = OmegaConf.to_container(self.cfg, resolve=False) - cfg_resolved = OmegaConf.create(cfg_resolved) - - trainer = instantiate(cfg_resolved.trainer, _recursive_=False) - trainer.run() - - def __call__(self): - job_env = submitit.JobEnvironment() - self.setup_job_info(job_env.job_id, job_env.global_rank) - try: - self.run_trainer() - except Exception as e: - # Log the exception. Then raise it again (as what SubmititRunner currently does). - message = format_exception(e) - logging.error(message) - raise e - - def setup_job_info(self, job_id, rank): - """Set up slurm job info""" - self.job_info = { - "job_id": job_id, - "rank": rank, - "cluster": self.cfg.get("cluster", None), - "experiment_log_dir": self.cfg.launcher.experiment_log_dir, - } - - self.has_setup = True - - -def add_pythonpath_to_sys_path(): - if "PYTHONPATH" not in os.environ or not os.environ["PYTHONPATH"]: - return - sys.path = os.environ["PYTHONPATH"].split(":") + sys.path - - -def main(args) -> None: - cfg = compose(config_name=args.config) - if cfg.launcher.experiment_log_dir is None: - cfg.launcher.experiment_log_dir = os.path.join( - os.getcwd(), "sam3_logs", args.config - ) - print("###################### Train App Config ####################") - print(OmegaConf.to_yaml(cfg)) - print("############################################################") - - add_pythonpath_to_sys_path() - makedir(cfg.launcher.experiment_log_dir) - with g_pathmgr.open( - os.path.join(cfg.launcher.experiment_log_dir, "config.yaml"), "w" - ) as f: - f.write(OmegaConf.to_yaml(cfg)) - - cfg_resolved = OmegaConf.to_container(cfg, resolve=False) - cfg_resolved = OmegaConf.create(cfg_resolved) - - with g_pathmgr.open( - os.path.join(cfg.launcher.experiment_log_dir, "config_resolved.yaml"), "w" - ) as f: - f.write(OmegaConf.to_yaml(cfg_resolved, resolve=True)) - - submitit_conf = cfg.get("submitit", None) - assert submitit_conf is not None, "Missing submitit config" - - experiment_log_dir = cfg.launcher.experiment_log_dir - print(f"Experiment Log Dir:\n{experiment_log_dir}") - submitit_dir = os.path.join(experiment_log_dir, "submitit_logs") - - # Prioritize cmd line args - cfg.launcher.gpus_per_node = ( - args.num_gpus if args.num_gpus is not None else cfg.launcher.gpus_per_node - ) - cfg.launcher.num_nodes = ( - args.num_nodes if args.num_nodes is not None else cfg.launcher.num_nodes - ) - submitit_conf.use_cluster = ( - args.use_cluster if args.use_cluster is not None else submitit_conf.use_cluster - ) - if submitit_conf.use_cluster: - executor = submitit.AutoExecutor(folder=submitit_dir) - submitit_conf.partition = ( - args.partition - if args.partition is not None - else submitit_conf.get("partition", None) - ) - submitit_conf.account = ( - args.account - if args.account is not None - else submitit_conf.get("account", None) - ) - submitit_conf.qos = ( - args.qos if args.qos is not None else submitit_conf.get("qos", None) - ) - job_kwargs = { - "timeout_min": 60 * submitit_conf.timeout_hour, - "name": ( - submitit_conf.name if hasattr(submitit_conf, "name") else args.config - ), - "slurm_partition": submitit_conf.partition, - "gpus_per_node": cfg.launcher.gpus_per_node, - "tasks_per_node": cfg.launcher.gpus_per_node, # one task per GPU - "cpus_per_task": submitit_conf.cpus_per_task, - "nodes": cfg.launcher.num_nodes, - "slurm_additional_parameters": { - "exclude": " ".join(submitit_conf.get("exclude_nodes", [])), - }, - } - if "include_nodes" in submitit_conf: - assert ( - len(submitit_conf["include_nodes"]) >= cfg.launcher.num_nodes - ), "Not enough nodes" - job_kwargs["slurm_additional_parameters"]["nodelist"] = " ".join( - submitit_conf["include_nodes"] - ) - if submitit_conf.account is not None: - job_kwargs["slurm_additional_parameters"]["account"] = submitit_conf.account - if submitit_conf.qos is not None: - job_kwargs["slurm_additional_parameters"]["qos"] = submitit_conf.qos - - if submitit_conf.get("mem_gb", None) is not None: - job_kwargs["mem_gb"] = submitit_conf.mem_gb - elif submitit_conf.get("mem", None) is not None: - job_kwargs["slurm_mem"] = submitit_conf.mem - - if submitit_conf.get("constraints", None) is not None: - job_kwargs["slurm_constraint"] = submitit_conf.constraints - - if submitit_conf.get("comment", None) is not None: - job_kwargs["slurm_comment"] = submitit_conf.comment - - # Supports only cpu-bind option within srun_args. New options can be added here - if submitit_conf.get("srun_args", None) is not None: - job_kwargs["slurm_srun_args"] = [] - if submitit_conf.srun_args.get("cpu_bind", None) is not None: - job_kwargs["slurm_srun_args"].extend( - ["--cpu-bind", submitit_conf.srun_args.cpu_bind] - ) - - print("###################### SLURM Config ####################") - print(job_kwargs) - print("##########################################") - executor.update_parameters(**job_kwargs) - - if ( - "job_array" in submitit_conf - and submitit_conf.job_array.get("num_tasks", -1) > 0 - ): - num_tasks = submitit_conf.job_array.num_tasks - job_array_config_dir = os.path.join( - cfg.launcher.experiment_log_dir, "job_array_configs" - ) - makedir(job_array_config_dir) - - job_indices = range(num_tasks) - ports = random.sample( - range(submitit_conf.port_range[0], submitit_conf.port_range[1] + 1), - k=len(job_indices), - ) - - jobs_runners_configs = [] - with executor.batch(): - task_index = 0 - for indices, main_port in tqdm(zip(job_indices, ports)): - curr_cfg = deepcopy(cfg) - curr_cfg.submitit.job_array["task_index"] = task_index - curr_cfg_resolved = handle_custom_resolving(cfg) - runner = SubmititRunner(main_port, curr_cfg) - job = executor.submit(runner) - jobs_runners_configs.append( - (job, runner, curr_cfg, curr_cfg_resolved) - ) - task_index += 1 - - for job, runner, job_cfg, job_cfg_resolved in jobs_runners_configs: - print("Submitit Job ID:", job.job_id) - - # Save job specific config - job_array_config_file = os.path.join( - job_array_config_dir, "{}.config.yaml".format(job.job_id) - ) - with g_pathmgr.open(job_array_config_file, "w") as f: - f.write(OmegaConf.to_yaml(job_cfg)) - - job_array_config_resolved_file = os.path.join( - job_array_config_dir, "{}.config_resolved.yaml".format(job.job_id) - ) - with g_pathmgr.open(job_array_config_resolved_file, "w") as f: - f.write(OmegaConf.to_yaml(job_cfg_resolved, resolve=True)) - - runner.setup_job_info(job.job_id, rank=0) - # runner.log_event(event_type=SlurmEvent.QUEUED) - else: - main_port = random.randint( - submitit_conf.port_range[0], submitit_conf.port_range[1] - ) - runner = SubmititRunner(main_port, cfg) - job = executor.submit(runner) - print(f"Submitit Job ID: {job.job_id}") - runner.setup_job_info(job.job_id, rank=0) - - else: - cfg.launcher.num_nodes = 1 - main_port = random.randint( - submitit_conf.port_range[0], submitit_conf.port_range[1] - ) - single_node_runner(cfg, main_port) - - -if __name__ == "__main__": - initialize_config_module("sam3.train", version_base="1.2") - parser = ArgumentParser() - parser.add_argument( - "-c", - "--config", - required=True, - type=str, - help="path to config file (e.g. configs/roboflow_v100_full_ft_100_images.yaml)", - ) - parser.add_argument( - "--use-cluster", - type=int, - default=None, - help="whether to launch on a cluster, 0: run locally, 1: run on a cluster", - ) - parser.add_argument("--partition", type=str, default=None, help="SLURM partition") - parser.add_argument("--account", type=str, default=None, help="SLURM account") - parser.add_argument("--qos", type=str, default=None, help="SLURM qos") - parser.add_argument( - "--num-gpus", type=int, default=None, help="number of GPUS per node" - ) - parser.add_argument("--num-nodes", type=int, default=None, help="Number of nodes") - args = parser.parse_args() - args.use_cluster = bool(args.use_cluster) if args.use_cluster is not None else None - register_omegaconf_resolvers() - main(args) diff --git a/train/trainer.py b/train/trainer.py deleted file mode 100644 index ac7c1b5..0000000 --- a/train/trainer.py +++ /dev/null @@ -1,1193 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved - -import contextlib -import fnmatch -import gc -import json -import logging -import math -import os -import time -from collections import OrderedDict -from dataclasses import dataclass, field -from typing import Any, Dict, List, Mapping, Optional - -import numpy as np - -import torch -import torch.distributed as dist -import torch.nn as nn -from hydra.utils import instantiate -from iopath.common.file_io import g_pathmgr - -from sam3.model.data_misc import BatchedDatapoint -from sam3.model.model_misc import SAM3Output -from sam3.model.utils.misc import copy_data_to_device - -from sam3.train.optim.optimizer import construct_optimizer - -from sam3.train.utils.checkpoint_utils import ( - assert_skipped_parameters_are_frozen, - exclude_params_matching_unix_pattern, - load_state_dict_into_model, - with_check_parameter_frozen, -) - -from sam3.train.utils.distributed import all_reduce_max, barrier, get_rank - -from sam3.train.utils.logger import Logger, setup_logging -from sam3.train.utils.train_utils import ( - AverageMeter, - collect_dict_keys, - DurationMeter, - get_amp_type, - get_machine_local_and_dist_rank, - get_resume_checkpoint, - human_readable_time, - is_dist_avail_and_initialized, - log_env_variables, - makedir, - MemMeter, - Phase, - ProgressMeter, - set_seeds, - setup_distributed_backend, -) - - -CORE_LOSS_KEY = "core_loss" - - -def unwrap_ddp_if_wrapped(model): - if isinstance(model, torch.nn.parallel.DistributedDataParallel): - return model.module - return model - - -@dataclass -class OptimAMPConf: - enabled: bool = False - amp_dtype: str = "float16" - - -@dataclass -class OptimConf: - optimizer: torch.optim.Optimizer = None - options: Optional[Dict[str, Any]] = None - param_group_modifiers: Optional[List] = None - amp: Optional[Dict[str, Any]] = None - gradient_clip: Any = None - gradient_logger: Any = None - - def __post_init__(self): - # amp - if not isinstance(self.amp, OptimAMPConf): - if self.amp is None: - self.amp = {} - assert isinstance(self.amp, Mapping) - self.amp = OptimAMPConf(**self.amp) - - -@dataclass -class DistributedConf: - backend: Optional[str] = None # inferred from accelerator type - comms_dtype: Optional[str] = None - find_unused_parameters: bool = False - timeout_mins: int = 30 - gradient_as_bucket_view: bool = False # PyTorch DDP default is False - static_graph: bool = False # PyTorch DDP default is False - - -@dataclass -class CudaConf: - cudnn_deterministic: bool = False - cudnn_benchmark: bool = True - allow_tf32: bool = False - # if not None, `matmul_allow_tf32` key will override `allow_tf32` for matmul - matmul_allow_tf32: Optional[bool] = None - # if not None, `cudnn_allow_tf32` key will override `allow_tf32` for cudnn - cudnn_allow_tf32: Optional[bool] = None - - -@dataclass -class CheckpointConf: - save_dir: str - save_freq: int - save_list: List[int] = field(default_factory=list) - model_weight_initializer: Any = None - save_best_meters: List[str] = None - skip_saving_parameters: List[str] = field(default_factory=list) - initialize_after_preemption: Optional[bool] = None - # if not None, training will be resumed from this checkpoint - resume_from: Optional[str] = None - - def infer_missing(self): - if self.initialize_after_preemption is None: - with_skip_saving = len(self.skip_saving_parameters) > 0 - self.initialize_after_preemption = with_skip_saving - return self - - -@dataclass -class LoggingConf: - log_dir: str - log_freq: int # In iterations - tensorboard_writer: Any - log_level_primary: str = "INFO" - log_level_secondary: str = "ERROR" - log_scalar_frequency: int = 100 - log_visual_frequency: int = 100 - scalar_keys_to_log: Optional[Dict[str, Any]] = None - log_batch_stats: bool = False - wandb_writer: Optional[Any] = None - - -class Trainer: - """ - Trainer supporting the DDP training strategies. - """ - - EPSILON = 1e-8 - - def __init__( - self, - *, # the order of these args can change at any time, so they are keyword-only - data: Dict[str, Any], - model: Dict[str, Any], - logging: Dict[str, Any], - checkpoint: Dict[str, Any], - max_epochs: int, - mode: str = "train", - accelerator: str = "cuda", - seed_value: int = 123, - val_epoch_freq: int = 1, - distributed: Dict[str, bool] = None, - cuda: Dict[str, bool] = None, - env_variables: Optional[Dict[str, Any]] = None, - optim: Optional[Dict[str, Any]] = None, - optim_overrides: Optional[List[Dict[str, Any]]] = None, - meters: Optional[Dict[str, Any]] = None, - loss: Optional[Dict[str, Any]] = None, - skip_first_val: bool = False, - skip_saving_ckpts: bool = False, - empty_gpu_mem_cache_after_eval: bool = True, - gradient_accumulation_steps: int = 1, - ): - self._setup_env_variables(env_variables) - self._setup_timers() - - self.data_conf = data - self.model_conf = model - self.logging_conf = LoggingConf(**logging) - self.checkpoint_conf = CheckpointConf(**checkpoint).infer_missing() - self.max_epochs = max_epochs - self.mode = mode - self.val_epoch_freq = val_epoch_freq - self.optim_conf = OptimConf(**optim) if optim is not None else OptimConf() - self.meters_conf = meters - self.loss_conf = loss - self.gradient_accumulation_steps = gradient_accumulation_steps - distributed = DistributedConf(**distributed or {}) - cuda = CudaConf(**cuda or {}) - self.where = 0.0 - - self.skip_first_val = skip_first_val - self.skip_saving_ckpts = skip_saving_ckpts - self.empty_gpu_mem_cache_after_eval = empty_gpu_mem_cache_after_eval - - self._infer_distributed_backend_if_none(distributed, accelerator) - - self._setup_device(accelerator) - - self._setup_torch_dist_and_backend(cuda, distributed) - - makedir(self.logging_conf.log_dir) - setup_logging( - __name__, - output_dir=self.logging_conf.log_dir, - rank=self.rank, - log_level_primary=self.logging_conf.log_level_primary, - log_level_secondary=self.logging_conf.log_level_secondary, - ) - - set_seeds(seed_value, self.max_epochs, self.distributed_rank) - log_env_variables() - - assert ( - is_dist_avail_and_initialized() - ), "Torch distributed needs to be initialized before calling the trainer." - - self._setup_components() # Except Optimizer everything is setup here. - self._move_to_device() - self._construct_optimizers() - self._setup_dataloaders() - - self.time_elapsed_meter = DurationMeter("Time Elapsed", self.device, ":.2f") - - if self.checkpoint_conf.resume_from is not None: - assert os.path.exists( - self.checkpoint_conf.resume_from - ), f"The 'resume_from' checkpoint {self.checkpoint_conf.resume_from} does not exist!" - dst = os.path.join(self.checkpoint_conf.save_dir, "checkpoint.pt") - if self.distributed_rank == 0 and not os.path.exists(dst): - # Copy the "resume_from" checkpoint to the checkpoint folder - # if there is not a checkpoint to resume from already there - makedir(self.checkpoint_conf.save_dir) - g_pathmgr.copy(self.checkpoint_conf.resume_from, dst) - barrier() - - self.load_checkpoint() - self._setup_ddp_distributed_training(distributed, accelerator) - barrier() - - def _setup_timers(self): - """ - Initializes counters for elapsed time and eta. - """ - self.start_time = time.time() - self.ckpt_time_elapsed = 0 - self.est_epoch_time = dict.fromkeys([Phase.TRAIN, Phase.VAL], 0) - - def _get_meters(self, phase_filters=None): - if self.meters is None: - return {} - meters = {} - for phase, phase_meters in self.meters.items(): - if phase_filters is not None and phase not in phase_filters: - continue - for key, key_meters in phase_meters.items(): - if key_meters is None: - continue - for name, meter in key_meters.items(): - meters[f"{phase}_{key}/{name}"] = meter - return meters - - def _infer_distributed_backend_if_none(self, distributed_conf, accelerator): - if distributed_conf.backend is None: - distributed_conf.backend = "nccl" if accelerator == "cuda" else "gloo" - - def _setup_env_variables(self, env_variables_conf) -> None: - if env_variables_conf is not None: - for variable_name, value in env_variables_conf.items(): - os.environ[variable_name] = value - - def _setup_torch_dist_and_backend(self, cuda_conf, distributed_conf) -> None: - if torch.cuda.is_available(): - torch.backends.cudnn.deterministic = cuda_conf.cudnn_deterministic - torch.backends.cudnn.benchmark = cuda_conf.cudnn_benchmark - torch.backends.cuda.matmul.allow_tf32 = ( - cuda_conf.matmul_allow_tf32 - if cuda_conf.matmul_allow_tf32 is not None - else cuda_conf.allow_tf32 - ) - torch.backends.cudnn.allow_tf32 = ( - cuda_conf.cudnn_allow_tf32 - if cuda_conf.cudnn_allow_tf32 is not None - else cuda_conf.allow_tf32 - ) - - self.rank = setup_distributed_backend( - distributed_conf.backend, distributed_conf.timeout_mins - ) - - def _setup_device(self, accelerator): - self.local_rank, self.distributed_rank = get_machine_local_and_dist_rank() - if accelerator == "cuda": - self.device = torch.device("cuda", self.local_rank) - torch.cuda.set_device(self.local_rank) - elif accelerator == "cpu": - self.device = torch.device("cpu") - else: - raise ValueError(f"Unsupported accelerator: {accelerator}") - - def _setup_ddp_distributed_training(self, distributed_conf, accelerator): - assert isinstance(self.model, torch.nn.Module) - - self.model = nn.parallel.DistributedDataParallel( - self.model, - device_ids=[self.local_rank] if accelerator == "cuda" else [], - find_unused_parameters=distributed_conf.find_unused_parameters, - gradient_as_bucket_view=distributed_conf.gradient_as_bucket_view, - static_graph=distributed_conf.static_graph, - ) - if distributed_conf.comms_dtype is not None: # noqa - from torch.distributed.algorithms import ddp_comm_hooks - - amp_type = get_amp_type(distributed_conf.comms_dtype) - if amp_type == torch.bfloat16: - hook = ddp_comm_hooks.default_hooks.bf16_compress_hook - logging.info("Enabling bfloat16 grad communication") - else: - hook = ddp_comm_hooks.default_hooks.fp16_compress_hook - logging.info("Enabling fp16 grad communication") - process_group = None - self.model.register_comm_hook(process_group, hook) - - def _move_to_device(self): - logging.info( - f"Moving components to device {self.device} and local rank {self.local_rank}." - ) - - self.model.to(self.device) - - logging.info( - f"Done moving components to device {self.device} and local rank {self.local_rank}." - ) - - def save_checkpoint(self, epoch, checkpoint_names=None): - if self.skip_saving_ckpts: - logging.info( - "skip_saving_ckpts is set to True. So, no checkpoints have been saved." - ) - return - checkpoint_folder = self.checkpoint_conf.save_dir - makedir(checkpoint_folder) - if checkpoint_names is None: - checkpoint_names = ["checkpoint"] - if ( - self.checkpoint_conf.save_freq > 0 - and (int(epoch) % self.checkpoint_conf.save_freq == 0) - ) or int(epoch) in self.checkpoint_conf.save_list: - checkpoint_names.append(f"checkpoint_{int(epoch)}") - - checkpoint_paths = [] - for ckpt_name in checkpoint_names: - checkpoint_paths.append(os.path.join(checkpoint_folder, f"{ckpt_name}.pt")) - - state_dict = unwrap_ddp_if_wrapped(self.model).state_dict() - state_dict = exclude_params_matching_unix_pattern( - patterns=self.checkpoint_conf.skip_saving_parameters, state_dict=state_dict - ) - - checkpoint = { - "model": state_dict, - "optimizer": self.optim.optimizer.state_dict(), - "epoch": epoch, - "loss": self.loss.state_dict(), - "steps": self.steps, - "time_elapsed": self.time_elapsed_meter.val, - "best_meter_values": self.best_meter_values, - } - if self.optim_conf.amp.enabled: - checkpoint["scaler"] = self.scaler.state_dict() - - # DDP checkpoints are only saved on rank 0 (all workers are identical) - if self.distributed_rank != 0: - return - - for checkpoint_path in checkpoint_paths: - self._save_checkpoint(checkpoint, checkpoint_path) - - def _save_checkpoint(self, checkpoint, checkpoint_path): - """ - Save a checkpoint while guarding against the job being killed in the middle - of checkpoint saving (which corrupts the checkpoint file and ruins the - entire training since usually only the last checkpoint is kept per run). - - We first save the new checkpoint to a temp file (with a '.tmp' suffix), and - and move it to overwrite the old checkpoint_path. - """ - checkpoint_path_tmp = f"{checkpoint_path}.tmp" - with g_pathmgr.open(checkpoint_path_tmp, "wb") as f: - torch.save(checkpoint, f) - # after torch.save is completed, replace the old checkpoint with the new one - if g_pathmgr.exists(checkpoint_path): - # remove the old checkpoint_path file first (otherwise g_pathmgr.mv fails) - g_pathmgr.rm(checkpoint_path) - success = g_pathmgr.mv(checkpoint_path_tmp, checkpoint_path) - assert success - - def load_checkpoint(self): - ckpt_path = get_resume_checkpoint(self.checkpoint_conf.save_dir) - if ckpt_path is None: - self._init_model_state() - else: - if self.checkpoint_conf.initialize_after_preemption: - self._call_model_initializer() - self._load_resuming_checkpoint(ckpt_path) - - def _init_model_state(self): - # Checking that parameters that won't be saved are indeed frozen - # We do this check here before even saving the model to catch errors - # are early as possible and not at the end of the first epoch - assert_skipped_parameters_are_frozen( - patterns=self.checkpoint_conf.skip_saving_parameters, - model=self.model, - ) - - # Checking that parameters that won't be saved are initialized from - # within the model definition, unless `initialize_after_preemption` - # is explicitly set to `True`. If not, this is a bug, and after - # preemption, the `skip_saving_parameters` will have random values - allow_init_skip_parameters = self.checkpoint_conf.initialize_after_preemption - with with_check_parameter_frozen( - patterns=self.checkpoint_conf.skip_saving_parameters, - model=self.model, - disabled=allow_init_skip_parameters, - ): - self._call_model_initializer() - - def _call_model_initializer(self): - model_weight_initializer = instantiate( - self.checkpoint_conf.model_weight_initializer - ) - if model_weight_initializer is not None: - logging.info( - f"Loading pretrained checkpoint from {self.checkpoint_conf.model_weight_initializer}" - ) - self.model = model_weight_initializer(model=self.model) - - def _load_resuming_checkpoint(self, ckpt_path: str): - logging.info(f"Resuming training from {ckpt_path}") - - with g_pathmgr.open(ckpt_path, "rb") as f: - checkpoint = torch.load(f, map_location="cpu") - load_state_dict_into_model( - model=self.model, - state_dict=checkpoint["model"], - ignore_missing_keys=self.checkpoint_conf.skip_saving_parameters, - ) - - self.optim.optimizer.load_state_dict(checkpoint["optimizer"]) - self.loss.load_state_dict(checkpoint["loss"], strict=True) - self.epoch = checkpoint["epoch"] - self.steps = checkpoint["steps"] - self.ckpt_time_elapsed = checkpoint.get("time_elapsed") - - if self.optim_conf.amp.enabled and "scaler" in checkpoint: - self.scaler.load_state_dict(checkpoint["scaler"]) - - self.best_meter_values = checkpoint.get("best_meter_values", {}) - - if "train_dataset" in checkpoint and self.train_dataset is not None: - self.train_dataset.load_checkpoint_state(checkpoint["train_dataset"]) - - def is_intermediate_val_epoch(self, epoch): - skip_epoch = self.skip_first_val and epoch == 0 - return ( - epoch % self.val_epoch_freq == 0 - and epoch < self.max_epochs - 1 - and not skip_epoch - ) - - def _find_loss(self, key: str): - if key in self.loss: - return self.loss[key] - - assert key != "all", "Loss must be specified for key='all'" - assert ( - "default" in self.loss - ), f"Key {key} not found in losss, and no default provided" - return self.loss["default"] - - def _find_meter(self, phase: str, key: str): - if key in self.meters[phase]: - return self.meters[phase][key] - - for cand_key, meter in self.meters[phase].items(): - if fnmatch.fnmatch(key, cand_key): - return meter - return None - - def _step( - self, - batch: BatchedDatapoint, - model: nn.Module, - phase: str, - ): - key, batch = batch.popitem() - batch = copy_data_to_device(batch, self.device, non_blocking=True) - - find_stages = model(batch) - find_targets = [ - unwrap_ddp_if_wrapped(model).back_convert(x) for x in batch.find_targets - ] - batch_size = len(batch.img_batch) - loss = self._find_loss(key)(find_stages, find_targets) - - loss_str = f"Losses/{phase}_{key}_loss" - - loss_log_str = os.path.join("Step_Losses", loss_str) - - # loss contains multiple sub-components we wish to log - step_losses = {} - if isinstance(loss, dict): - step_losses.update( - {f"Losses/{phase}_{key}_{k}": v for k, v in loss.items()} - ) - loss = self._log_loss_detailed_and_return_core_loss( - loss, loss_log_str, self.steps[phase] - ) - - if self.steps[phase] % self.logging_conf.log_scalar_frequency == 0: - self.logger.log( - loss_log_str, - loss, - self.steps[phase], - ) - - self.steps[phase] += 1 - - ret_tuple = {loss_str: loss}, batch_size, step_losses - - if phase not in self.meters: - return ret_tuple - - meters_dict = self._find_meter(phase, key) - if meters_dict is None: - return ret_tuple - if meters_dict is not None: - for _, meter in meters_dict.items(): - meter.update( - find_stages=find_stages, - find_metadatas=batch.find_metadatas, - model=model, - batch=batch, - key=key, - ) - # Cleanup memory - if isinstance(find_stages, SAM3Output): - for fs in find_stages: - for k in list(fs.keys()): - del fs[k] - - return ret_tuple - - def run(self): - assert self.mode in ["train", "train_only", "val"] - if self.mode == "train": - if self.epoch > 0: - logging.info(f"Resuming training from epoch: {self.epoch}") - # resuming from a checkpoint - if self.is_intermediate_val_epoch(self.epoch - 1): - logging.info("Running previous val epoch") - self.epoch -= 1 - self.run_val() - self.epoch += 1 - self.run_train() - self.run_val() - elif self.mode == "val": - self.run_val() - elif self.mode == "train_only": - self.run_train() - - def _setup_dataloaders(self): - self.train_dataset = None - self.val_dataset = None - - if self.mode in ["train", "val"]: - self.val_dataset = instantiate(self.data_conf.get(Phase.VAL, None)) - - if self.mode in ["train", "train_only"]: - self.train_dataset = instantiate(self.data_conf.train) - - def run_train(self): - while self.epoch < self.max_epochs: - dataloader = self.train_dataset.get_loader(epoch=int(self.epoch)) - barrier() - outs = self.train_epoch(dataloader) - self.logger.log_dict(outs, self.epoch) # Logged only on rank 0 - - # log train to text file. - if self.distributed_rank == 0: - with g_pathmgr.open( - os.path.join(self.logging_conf.log_dir, "train_stats.json"), - "a", - ) as f: - f.write(json.dumps(outs) + "\n") - - # Save checkpoint before validating - self.save_checkpoint(self.epoch + 1) - - del dataloader - gc.collect() - - # Run val, not running on last epoch since will run after the - # loop anyway - if self.is_intermediate_val_epoch(self.epoch): - self.run_val() - if torch.cuda.is_available() and self.empty_gpu_mem_cache_after_eval: - # release memory buffers held by the model during eval (which typically - # involves a lot more frames in video grounding that during training) - torch.cuda.empty_cache() - - if self.distributed_rank == 0: - self.best_meter_values.update(self._get_trainer_state("train")) - with g_pathmgr.open( - os.path.join(self.logging_conf.log_dir, "best_stats.json"), - "a", - ) as f: - f.write(json.dumps(self.best_meter_values) + "\n") - - self.epoch += 1 - # epoch was incremented in the loop but the val step runs out of the loop - self.epoch -= 1 - - def run_val(self): - if not self.val_dataset: - return - - dataloader = self.val_dataset.get_loader(epoch=int(self.epoch)) - outs = self.val_epoch(dataloader, phase=Phase.VAL) - del dataloader - gc.collect() - self.logger.log_dict(outs, self.epoch) # Logged only on rank 0 - - if self.distributed_rank == 0: - with g_pathmgr.open( - os.path.join(self.logging_conf.log_dir, "val_stats.json"), - "a", - ) as f: - f.write(json.dumps(outs) + "\n") - - def val_epoch(self, val_loader, phase): - batch_time = AverageMeter("Batch Time", self.device, ":.2f") - data_time = AverageMeter("Data Time", self.device, ":.2f") - mem = MemMeter("Mem (GB)", self.device, ":.2f") - - iters_per_epoch = len(val_loader) - - curr_phases = [phase] - curr_models = [self.model] - - loss_names = [] - for p in curr_phases: - for key in self.loss.keys(): - loss_names.append(f"Losses/{p}_{key}_loss") - - loss_mts = OrderedDict( - [(name, AverageMeter(name, self.device, ":.2e")) for name in loss_names] - ) - extra_loss_mts = {} - - for model in curr_models: - model.eval() - if hasattr(unwrap_ddp_if_wrapped(model), "on_validation_epoch_start"): - unwrap_ddp_if_wrapped(model).on_validation_epoch_start() - - progress = ProgressMeter( - iters_per_epoch, - [batch_time, data_time, mem, self.time_elapsed_meter, *loss_mts.values()], - self._get_meters(curr_phases), - prefix="Val Epoch: [{}]".format(self.epoch), - ) - - end = time.time() - - for data_iter, batch in enumerate(val_loader): - # measure data loading time - data_time.update(time.time() - end) - - # batch = batch.to(self.device, non_blocking=True) - - # compute output - with torch.no_grad(): - with torch.amp.autocast( - device_type="cuda", - enabled=(self.optim_conf.amp.enabled if self.optim_conf else False), - dtype=( - get_amp_type(self.optim_conf.amp.amp_dtype) - if self.optim_conf - else None - ), - ): - for phase, model in zip(curr_phases, curr_models): - loss_dict, batch_size, extra_losses = self._step( - batch, - model, - phase, - ) - - assert len(loss_dict) == 1 - loss_key, loss = loss_dict.popitem() - - if loss_key in loss_mts: - loss_mts[loss_key].update(loss.item(), batch_size) - - for k, v in extra_losses.items(): - if k not in extra_loss_mts: - extra_loss_mts[k] = AverageMeter(k, self.device, ":.2e") - extra_loss_mts[k].update(v.item(), batch_size) - - # measure elapsed time - batch_time.update(time.time() - end) - end = time.time() - - self.time_elapsed_meter.update( - time.time() - self.start_time + self.ckpt_time_elapsed - ) - - if torch.cuda.is_available(): - mem.update(reset_peak_usage=True) - - if data_iter % self.logging_conf.log_freq == 0: - progress.display(data_iter) - - if data_iter % self.logging_conf.log_scalar_frequency == 0: - # Log progress meters. - for progress_meter in progress.meters: - self.logger.log( - os.path.join("Step_Stats", phase, progress_meter.name), - progress_meter.val, - self.steps[Phase.VAL], - ) - - if data_iter % 10 == 0: - dist.barrier() - - self.est_epoch_time[phase] = batch_time.avg * iters_per_epoch - self._log_timers(phase) - for model in curr_models: - if hasattr(unwrap_ddp_if_wrapped(model), "on_validation_epoch_end"): - unwrap_ddp_if_wrapped(model).on_validation_epoch_end() - - out_dict = self._log_meters_and_save_best_ckpts(curr_phases) - - for k, v in loss_mts.items(): - out_dict[k] = v.avg - for k, v in extra_loss_mts.items(): - out_dict[k] = v.avg - - for phase in curr_phases: - out_dict.update(self._get_trainer_state(phase)) - self._reset_meters(curr_phases) - logging.info(f"Meters: {out_dict}") - return out_dict - - def _get_trainer_state(self, phase): - return { - "Trainer/where": self.where, - "Trainer/epoch": self.epoch, - f"Trainer/steps_{phase}": self.steps[phase], - } - - def train_epoch(self, train_loader): - # Init stat meters - batch_time_meter = AverageMeter("Batch Time", self.device, ":.2f") - data_time_meter = AverageMeter("Data Time", self.device, ":.2f") - mem_meter = MemMeter("Mem (GB)", self.device, ":.2f") - data_times = [] - phase = Phase.TRAIN - - iters_per_epoch = len(train_loader) - - loss_names = [] - for batch_key in self.loss.keys(): - loss_names.append(f"Losses/{phase}_{batch_key}_loss") - - loss_mts = OrderedDict( - [(name, AverageMeter(name, self.device, ":.2e")) for name in loss_names] - ) - extra_loss_mts = {} - - progress = ProgressMeter( - iters_per_epoch, - [ - batch_time_meter, - data_time_meter, - mem_meter, - self.time_elapsed_meter, - *loss_mts.values(), - ], - self._get_meters([phase]), - prefix="Train Epoch: [{}]".format(self.epoch), - ) - - # Model training loop - self.model.train() - end = time.time() - - for data_iter, batch in enumerate(train_loader): - # measure data loading time - data_time_meter.update(time.time() - end) - data_times.append(data_time_meter.val) - # batch = batch.to( - # self.device, non_blocking=True - # ) # move tensors in a tensorclass - - try: - self._run_step(batch, phase, loss_mts, extra_loss_mts) - - # compute gradient and do optim step - exact_epoch = self.epoch + float(data_iter) / iters_per_epoch - self.where = float(exact_epoch) / self.max_epochs - assert self.where <= 1 + self.EPSILON - if self.where < 1.0: - self.optim.step_schedulers( - self.where, step=int(exact_epoch * iters_per_epoch) - ) - else: - logging.warning( - f"Skipping scheduler update since the training is at the end, i.e, {self.where} of [0,1]." - ) - - # Log schedulers - if data_iter % self.logging_conf.log_scalar_frequency == 0: - for j, param_group in enumerate(self.optim.optimizer.param_groups): - for option in self.optim.schedulers[j]: - optim_prefix = ( - "" + f"{j}_" - if len(self.optim.optimizer.param_groups) > 1 - else "" - ) - self.logger.log( - os.path.join("Optim", f"{optim_prefix}", option), - param_group[option], - self.steps[phase], - ) - - # Clipping gradients and detecting diverging gradients - if self.gradient_clipper is not None: - self.scaler.unscale_(self.optim.optimizer) - self.gradient_clipper(model=self.model) - - if self.gradient_logger is not None: - self.gradient_logger( - self.model, rank=self.distributed_rank, where=self.where - ) - - # Optimizer step: the scaler will make sure gradients are not - # applied if the gradients are infinite - self.scaler.step(self.optim.optimizer) - self.scaler.update() - - # measure elapsed time - batch_time_meter.update(time.time() - end) - end = time.time() - - self.time_elapsed_meter.update( - time.time() - self.start_time + self.ckpt_time_elapsed - ) - - mem_meter.update(reset_peak_usage=True) - if data_iter % self.logging_conf.log_freq == 0: - progress.display(data_iter) - - if data_iter % self.logging_conf.log_scalar_frequency == 0: - # Log progress meters. - for progress_meter in progress.meters: - self.logger.log( - os.path.join("Step_Stats", phase, progress_meter.name), - progress_meter.val, - self.steps[phase], - ) - - # Catching NaN/Inf errors in the loss - except FloatingPointError as e: - raise e - - self.est_epoch_time[Phase.TRAIN] = batch_time_meter.avg * iters_per_epoch - self._log_timers(Phase.TRAIN) - self._log_sync_data_times(Phase.TRAIN, data_times) - - out_dict = self._log_meters_and_save_best_ckpts([Phase.TRAIN]) - - for k, v in loss_mts.items(): - out_dict[k] = v.avg - for k, v in extra_loss_mts.items(): - out_dict[k] = v.avg - out_dict.update(self._get_trainer_state(phase)) - logging.info(f"Losses and meters: {out_dict}") - self._reset_meters([phase]) - return out_dict - - def _log_sync_data_times(self, phase, data_times): - data_times = all_reduce_max(torch.tensor(data_times)).tolist() - steps = range(self.steps[phase] - len(data_times), self.steps[phase]) - for step, data_time in zip(steps, data_times): - if step % self.logging_conf.log_scalar_frequency == 0: - self.logger.log( - os.path.join("Step_Stats", phase, "Data Time Synced"), - data_time, - step, - ) - - def _run_step( - self, - batch: BatchedDatapoint, - phase: str, - loss_mts: Dict[str, AverageMeter], - extra_loss_mts: Dict[str, AverageMeter], - raise_on_error: bool = True, - ): - """ - Run the forward / backward - """ - - # it's important to set grads to None, especially with Adam since 0 - # grads will also update a model even if the step doesn't produce - # gradients - self.optim.zero_grad(set_to_none=True) - - if self.gradient_accumulation_steps > 1: - assert isinstance( - batch, list - ), f"Expected a list of batches, got {type(batch)}" - assert ( - len(batch) == self.gradient_accumulation_steps - ), f"Expected {self.gradient_accumulation_steps} batches, got {len(batch)}" - accum_steps = len(batch) - else: - accum_steps = 1 - batch = [batch] - - for i, chunked_batch in enumerate(batch): - ddp_context = ( - self.model.no_sync() - if i < accum_steps - 1 - else contextlib.nullcontext() - ) - with ddp_context: - with torch.amp.autocast( - device_type="cuda", - enabled=self.optim_conf.amp.enabled, - dtype=get_amp_type(self.optim_conf.amp.amp_dtype), - ): - loss_dict, batch_size, extra_losses = self._step( - chunked_batch, - self.model, - phase, - ) - - assert len(loss_dict) == 1 - loss_key, loss = loss_dict.popitem() - - if not math.isfinite(loss.item()): - error_msg = f"Loss is {loss.item()}, attempting to stop training" - logging.error(error_msg) - if raise_on_error: - raise FloatingPointError(error_msg) - else: - return - - self.scaler.scale(loss).backward() - loss_mts[loss_key].update(loss.item(), batch_size) - for extra_loss_key, extra_loss in extra_losses.items(): - if extra_loss_key not in extra_loss_mts: - extra_loss_mts[extra_loss_key] = AverageMeter( - extra_loss_key, self.device, ":.2e" - ) - extra_loss_mts[extra_loss_key].update(extra_loss.item(), batch_size) - - def _log_meters_and_save_best_ckpts(self, phases: List[str]): - logging.info("Synchronizing meters") - out_dict = {} - checkpoint_save_keys = [] - for key, meter in self._get_meters(phases).items(): - meter_output = meter.compute_synced() - is_better_check = getattr(meter, "is_better", None) - - for meter_subkey, meter_value in meter_output.items(): - out_dict[os.path.join("Meters_train", key, meter_subkey)] = meter_value - - if is_better_check is None: - continue - - tracked_meter_key = os.path.join(key, meter_subkey) - if tracked_meter_key not in self.best_meter_values or is_better_check( - meter_value, - self.best_meter_values[tracked_meter_key], - ): - self.best_meter_values[tracked_meter_key] = meter_value - - if ( - self.checkpoint_conf.save_best_meters is not None - and key in self.checkpoint_conf.save_best_meters - ): - checkpoint_save_keys.append(tracked_meter_key.replace("/", "_")) - - if len(checkpoint_save_keys) > 0: - self.save_checkpoint(self.epoch + 1, checkpoint_save_keys) - - return out_dict - - def _log_timers(self, phase): - time_remaining = 0 - epochs_remaining = self.max_epochs - self.epoch - 1 - val_epochs_remaining = sum( - n % self.val_epoch_freq == 0 for n in range(self.epoch, self.max_epochs) - ) - - # Adding the guaranteed val run at the end if val_epoch_freq doesn't coincide with - # the end epoch. - if (self.max_epochs - 1) % self.val_epoch_freq != 0: - val_epochs_remaining += 1 - - # Remove the current val run from estimate - if phase == Phase.VAL: - val_epochs_remaining -= 1 - - time_remaining += ( - epochs_remaining * self.est_epoch_time[Phase.TRAIN] - + val_epochs_remaining * self.est_epoch_time[Phase.VAL] - ) - - self.logger.log( - os.path.join("Step_Stats", phase, self.time_elapsed_meter.name), - self.time_elapsed_meter.val, - self.steps[phase], - ) - - logging.info(f"Estimated time remaining: {human_readable_time(time_remaining)}") - - def _reset_meters(self, phases: str) -> None: - for meter in self._get_meters(phases).values(): - meter.reset() - - def _check_val_key_match(self, val_keys, phase): - if val_keys is not None: - # Check if there are any duplicates - assert len(val_keys) == len( - set(val_keys) - ), f"Duplicate keys in val datasets, keys: {val_keys}" - - # Check that the keys match the meter keys - if self.meters_conf is not None and phase in self.meters_conf: - assert set(val_keys) == set(self.meters_conf[phase].keys()), ( - f"Keys in val datasets do not match the keys in meters." - f"\nMissing in meters: {set(val_keys) - set(self.meters_conf[phase].keys())}" - f"\nMissing in val datasets: {set(self.meters_conf[phase].keys()) - set(val_keys)}" - ) - - if self.loss_conf is not None: - loss_keys = set(self.loss_conf.keys()) - set(["all"]) - if "default" not in loss_keys: - for k in val_keys: - assert ( - k in loss_keys - ), f"Error: key {k} is not defined in the losses, and no default is set" - - def _setup_components(self): - # Get the keys for all the val datasets, if any - val_phase = Phase.VAL - val_keys = None - if self.data_conf.get(val_phase, None) is not None: - val_keys = collect_dict_keys(self.data_conf[val_phase]) - # Additional checks on the sanity of the config for val datasets - self._check_val_key_match(val_keys, phase=val_phase) - - logging.info("Setting up components: Model, loss, optim, meters etc.") - self.epoch = 0 - self.steps = {Phase.TRAIN: 0, Phase.VAL: 0} - - self.logger = Logger(self.logging_conf) - - self.model = instantiate(self.model_conf, _convert_="all") - print_model_summary(self.model) - - self.loss = None - if self.loss_conf: - self.loss = { - key: el # wrap_base_loss(el) - for (key, el) in instantiate(self.loss_conf, _convert_="all").items() - } - self.loss = nn.ModuleDict(self.loss) - - self.meters = {} - self.best_meter_values = {} - if self.meters_conf: - self.meters = instantiate(self.meters_conf, _convert_="all") - - self.scaler = torch.amp.GradScaler( - self.device, - enabled=self.optim_conf.amp.enabled if self.optim_conf else False, - ) - - self.gradient_clipper = ( - instantiate(self.optim_conf.gradient_clip) if self.optim_conf else None - ) - self.gradient_logger = ( - instantiate(self.optim_conf.gradient_logger) if self.optim_conf else None - ) - - logging.info("Finished setting up components: Model, loss, optim, meters etc.") - - def _construct_optimizers(self): - self.optim = construct_optimizer( - self.model, - self.optim_conf.optimizer, - self.optim_conf.options, - self.optim_conf.param_group_modifiers, - ) - - def _log_loss_detailed_and_return_core_loss(self, loss, loss_str, step): - core_loss = loss.pop(CORE_LOSS_KEY) - if step % self.logging_conf.log_scalar_frequency == 0: - for k in loss: - log_str = os.path.join(loss_str, k) - self.logger.log(log_str, loss[k], step) - return core_loss - - -def print_model_summary(model: torch.nn.Module, log_dir: str = ""): - """ - Prints the model and the number of parameters in the model. - # Multiple packages provide this info in a nice table format - # However, they need us to provide an `input` (as they also write down the output sizes) - # Our models are complex, and a single input is restrictive. - # https://github.com/sksq96/pytorch-summary - # https://github.com/nmhkahn/torchsummaryX - """ - if get_rank() != 0: - return - param_kwargs = {} - trainable_parameters = sum( - p.numel() for p in model.parameters(**param_kwargs) if p.requires_grad - ) - total_parameters = sum(p.numel() for p in model.parameters(**param_kwargs)) - non_trainable_parameters = total_parameters - trainable_parameters - logging.info("==" * 10) - logging.info(f"Summary for model {type(model)}") - logging.info(f"Model is {model}") - logging.info(f"\tTotal parameters {get_human_readable_count(total_parameters)}") - logging.info( - f"\tTrainable parameters {get_human_readable_count(trainable_parameters)}" - ) - logging.info( - f"\tNon-Trainable parameters {get_human_readable_count(non_trainable_parameters)}" - ) - logging.info("==" * 10) - - if log_dir: - output_fpath = os.path.join(log_dir, "model.txt") - with g_pathmgr.open(output_fpath, "w") as f: - print(model, file=f) - - -PARAMETER_NUM_UNITS = [" ", "K", "M", "B", "T"] - - -def get_human_readable_count(number: int) -> str: - """ - Abbreviates an integer number with K, M, B, T for thousands, millions, - billions and trillions, respectively. - Examples: - >>> get_human_readable_count(123) - '123 ' - >>> get_human_readable_count(1234) # (one thousand) - '1.2 K' - >>> get_human_readable_count(2e6) # (two million) - '2.0 M' - >>> get_human_readable_count(3e9) # (three billion) - '3.0 B' - >>> get_human_readable_count(4e14) # (four hundred trillion) - '400 T' - >>> get_human_readable_count(5e15) # (more than trillion) - '5,000 T' - Args: - number: a positive integer number - Return: - A string formatted according to the pattern described above. - """ - assert number >= 0 - labels = PARAMETER_NUM_UNITS - num_digits = int(np.floor(np.log10(number)) + 1 if number > 0 else 1) - num_groups = int(np.ceil(num_digits / 3)) - num_groups = min(num_groups, len(labels)) # don't abbreviate beyond trillions - shift = -3 * (num_groups - 1) - number = number * (10**shift) - index = num_groups - 1 - if index < 1 or number >= 100: - return f"{int(number):,d} {labels[index]}" - else: - return f"{number:,.1f} {labels[index]}" diff --git a/train/transforms/__init__.py b/train/transforms/__init__.py deleted file mode 100644 index 46d37d2..0000000 --- a/train/transforms/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved diff --git a/train/transforms/basic.py b/train/transforms/basic.py deleted file mode 100644 index 3e8cf1f..0000000 --- a/train/transforms/basic.py +++ /dev/null @@ -1,455 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved - -""" -Transforms and data augmentation for both image + bbox. -""" - -import math -import random -from typing import Iterable - -import PIL -import torch -import torchvision.transforms as T -import torchvision.transforms.functional as F - -from sam3.model.box_ops import box_xyxy_to_cxcywh -from sam3.model.data_misc import interpolate - - -def crop(image, target, region): - cropped_image = F.crop(image, *region) - - target = target.copy() - i, j, h, w = region - - # should we do something wrt the original size? - target["size"] = torch.tensor([h, w]) - - fields = ["labels", "area", "iscrowd", "positive_map"] - - if "boxes" in target: - boxes = target["boxes"] - max_size = torch.as_tensor([w, h], dtype=torch.float32) - cropped_boxes = boxes - torch.as_tensor([j, i, j, i], dtype=torch.float32) - cropped_boxes = torch.min(cropped_boxes.reshape(-1, 2, 2), max_size) - cropped_boxes = cropped_boxes.clamp(min=0) - area = (cropped_boxes[:, 1, :] - cropped_boxes[:, 0, :]).prod(dim=1) - target["boxes"] = cropped_boxes.reshape(-1, 4) - target["area"] = area - fields.append("boxes") - - if "input_boxes" in target: - boxes = target["input_boxes"] - max_size = torch.as_tensor([w, h], dtype=torch.float32) - cropped_boxes = boxes - torch.as_tensor([j, i, j, i], dtype=torch.float32) - cropped_boxes = torch.min(cropped_boxes.reshape(-1, 2, 2), max_size) - cropped_boxes = cropped_boxes.clamp(min=0) - target["input_boxes"] = cropped_boxes.reshape(-1, 4) - - if "masks" in target: - # FIXME should we update the area here if there are no boxes? - target["masks"] = target["masks"][:, i : i + h, j : j + w] - fields.append("masks") - - # remove elements for which the boxes or masks that have zero area - if "boxes" in target or "masks" in target: - # favor boxes selection when defining which elements to keep - # this is compatible with previous implementation - if "boxes" in target: - cropped_boxes = target["boxes"].reshape(-1, 2, 2) - keep = torch.all(cropped_boxes[:, 1, :] > cropped_boxes[:, 0, :], dim=1) - else: - keep = target["masks"].flatten(1).any(1) - - for field in fields: - if field in target: - target[field] = target[field][keep] - - return cropped_image, target - - -def hflip(image, target): - flipped_image = F.hflip(image) - - w, h = image.size - - target = target.copy() - if "boxes" in target: - boxes = target["boxes"] - boxes = boxes[:, [2, 1, 0, 3]] * torch.as_tensor( - [-1, 1, -1, 1] - ) + torch.as_tensor([w, 0, w, 0]) - target["boxes"] = boxes - - if "input_boxes" in target: - boxes = target["input_boxes"] - boxes = boxes[:, [2, 1, 0, 3]] * torch.as_tensor( - [-1, 1, -1, 1] - ) + torch.as_tensor([w, 0, w, 0]) - target["input_boxes"] = boxes - - if "masks" in target: - target["masks"] = target["masks"].flip(-1) - - if "text_input" in target: - text_input = ( - target["text_input"] - .replace("left", "[TMP]") - .replace("right", "left") - .replace("[TMP]", "right") - ) - target["text_input"] = text_input - - return flipped_image, target - - -def resize(image, target, size, max_size=None, square=False): - # size can be min_size (scalar) or (w, h) tuple - - def get_size_with_aspect_ratio(image_size, size, max_size=None): - w, h = image_size - if max_size is not None: - min_original_size = float(min((w, h))) - max_original_size = float(max((w, h))) - if max_original_size / min_original_size * size > max_size: - size = int(round(max_size * min_original_size / max_original_size)) - - if (w <= h and w == size) or (h <= w and h == size): - return (h, w) - - if w < h: - ow = size - oh = int(size * h / w) - else: - oh = size - ow = int(size * w / h) - - return (oh, ow) - - def get_size(image_size, size, max_size=None): - if isinstance(size, (list, tuple)): - return size[::-1] - else: - return get_size_with_aspect_ratio(image_size, size, max_size) - - if square: - size = size, size - else: - size = get_size(image.size, size, max_size) - rescaled_image = F.resize(image, size) - - if target is None: - return rescaled_image, None - - ratios = tuple( - float(s) / float(s_orig) for s, s_orig in zip(rescaled_image.size, image.size) - ) - ratio_width, ratio_height = ratios - - target = target.copy() - if "boxes" in target: - boxes = target["boxes"] - scaled_boxes = boxes * torch.as_tensor( - [ratio_width, ratio_height, ratio_width, ratio_height], dtype=torch.float32 - ) - target["boxes"] = scaled_boxes - if "input_boxes" in target: - boxes = target["input_boxes"] - scaled_boxes = boxes * torch.as_tensor( - [ratio_width, ratio_height, ratio_width, ratio_height], dtype=torch.float32 - ) - target["input_boxes"] = scaled_boxes - - if "area" in target: - area = target["area"] - scaled_area = area * (ratio_width * ratio_height) - target["area"] = scaled_area - - h, w = size - target["size"] = torch.tensor([h, w]) - - if "masks" in target: - target["masks"] = ( - interpolate(target["masks"][:, None].float(), size, mode="nearest")[:, 0] - > 0.5 - ) - - return rescaled_image, target - - -def pad(image, target, padding): - if len(padding) == 2: - # assumes that we only pad on the bottom right corners - padded_image = F.pad(image, (0, 0, padding[0], padding[1])) - else: - # left, top, right, bottom - padded_image = F.pad(image, (padding[0], padding[1], padding[2], padding[3])) - if target is None: - return padded_image, None - target = target.copy() - - w, h = padded_image.size - - # should we do something wrt the original size? - target["size"] = torch.tensor([h, w]) - if "boxes" in target and len(padding) == 4: - boxes = target["boxes"] - boxes = boxes + torch.as_tensor( - [padding[0], padding[1], padding[0], padding[1]], dtype=torch.float32 - ) - target["boxes"] = boxes - - if "input_boxes" in target and len(padding) == 4: - boxes = target["input_boxes"] - boxes = boxes + torch.as_tensor( - [padding[0], padding[1], padding[0], padding[1]], dtype=torch.float32 - ) - target["input_boxes"] = boxes - - if "masks" in target: - if len(padding) == 2: - target["masks"] = torch.nn.functional.pad( - target["masks"], (0, padding[0], 0, padding[1]) - ) - else: - target["masks"] = torch.nn.functional.pad( - target["masks"], (padding[0], padding[2], padding[1], padding[3]) - ) - return padded_image, target - - -class RandomCrop: - def __init__(self, size): - self.size = size - - def __call__(self, img, target): - region = T.RandomCrop.get_params(img, self.size) - return crop(img, target, region) - - -class RandomSizeCrop: - def __init__(self, min_size: int, max_size: int, respect_boxes: bool = False): - self.min_size = min_size - self.max_size = max_size - self.respect_boxes = respect_boxes # if True we can't crop a box out - - def __call__(self, img: PIL.Image.Image, target: dict): - init_boxes = len(target["boxes"]) - init_boxes_tensor = target["boxes"].clone() - if self.respect_boxes and init_boxes > 0: - minW, minH, maxW, maxH = ( - min(img.width, self.min_size), - min(img.width, self.min_size), - min(img.width, self.max_size), - min(img.height, self.max_size), - ) - minX, minY = ( - target["boxes"][:, 0].max().item() + 10.0, - target["boxes"][:, 1].max().item() + 10.0, - ) - minX = min(img.width, minX) - minY = min(img.height, minY) - maxX, maxY = ( - target["boxes"][:, 2].min().item() - 10, - target["boxes"][:, 3].min().item() - 10, - ) - maxX = max(0.0, maxX) - maxY = max(0.0, maxY) - minW = max(minW, minX - maxX) - minH = max(minH, minY - maxY) - w = random.uniform(minW, max(minW, maxW)) - h = random.uniform(minH, max(minH, maxH)) - if minX > maxX: - # i = random.uniform(max(0, minX - w + 1), max(maxX, max(0, minX - w + 1))) - i = random.uniform(max(0, minX - w), max(maxX, max(0, minX - w))) - else: - i = random.uniform( - max(0, minX - w + 1), max(maxX - 1, max(0, minX - w + 1)) - ) - if minY > maxY: - # j = random.uniform(max(0, minY - h + 1), max(maxY, max(0, minY - h + 1))) - j = random.uniform(max(0, minY - h), max(maxY, max(0, minY - h))) - else: - j = random.uniform( - max(0, minY - h + 1), max(maxY - 1, max(0, minY - h + 1)) - ) - result_img, result_target = crop(img, target, [j, i, h, w]) - assert ( - len(result_target["boxes"]) == init_boxes - ), f"img_w={img.width}\timg_h={img.height}\tminX={minX}\tminY={minY}\tmaxX={maxX}\tmaxY={maxY}\tminW={minW}\tminH={minH}\tmaxW={maxW}\tmaxH={maxH}\tw={w}\th={h}\ti={i}\tj={j}\tinit_boxes={init_boxes_tensor}\tresults={result_target['boxes']}" - - return result_img, result_target - else: - w = random.randint(self.min_size, min(img.width, self.max_size)) - h = random.randint(self.min_size, min(img.height, self.max_size)) - region = T.RandomCrop.get_params(img, (h, w)) - result_img, result_target = crop(img, target, region) - return result_img, result_target - - -class CenterCrop: - def __init__(self, size): - self.size = size - - def __call__(self, img, target): - image_width, image_height = img.size - crop_height, crop_width = self.size - crop_top = int(round((image_height - crop_height) / 2.0)) - crop_left = int(round((image_width - crop_width) / 2.0)) - return crop(img, target, (crop_top, crop_left, crop_height, crop_width)) - - -class RandomHorizontalFlip: - def __init__(self, p=0.5): - self.p = p - - def __call__(self, img, target): - if random.random() < self.p: - return hflip(img, target) - return img, target - - -class RandomResize: - def __init__(self, sizes, max_size=None, square=False): - if isinstance(sizes, int): - sizes = (sizes,) - assert isinstance(sizes, Iterable) - self.sizes = list(sizes) - self.max_size = max_size - self.square = square - - def __call__(self, img, target=None): - size = random.choice(self.sizes) - return resize(img, target, size, self.max_size, square=self.square) - - -class RandomPad: - def __init__(self, max_pad): - self.max_pad = max_pad - - def __call__(self, img, target): - pad_x = random.randint(0, self.max_pad) - pad_y = random.randint(0, self.max_pad) - return pad(img, target, (pad_x, pad_y)) - - -class PadToSize: - def __init__(self, size): - self.size = size - - def __call__(self, img, target): - w, h = img.size - pad_x = self.size - w - pad_y = self.size - h - assert pad_x >= 0 and pad_y >= 0 - pad_left = random.randint(0, pad_x) - pad_right = pad_x - pad_left - pad_top = random.randint(0, pad_y) - pad_bottom = pad_y - pad_top - return pad(img, target, (pad_left, pad_top, pad_right, pad_bottom)) - - -class Identity: - def __call__(self, img, target): - return img, target - - -class RandomSelect: - """ - Randomly selects between transforms1 and transforms2, - with probability p for transforms1 and (1 - p) for transforms2 - """ - - def __init__(self, transforms1=None, transforms2=None, p=0.5): - self.transforms1 = transforms1 or Identity() - self.transforms2 = transforms2 or Identity() - self.p = p - - def __call__(self, img, target): - if random.random() < self.p: - return self.transforms1(img, target) - return self.transforms2(img, target) - - -class ToTensor: - def __call__(self, img, target): - return F.to_tensor(img), target - - -class RandomErasing: - def __init__(self, *args, **kwargs): - self.eraser = T.RandomErasing(*args, **kwargs) - - def __call__(self, img, target): - return self.eraser(img), target - - -class Normalize: - def __init__(self, mean, std): - self.mean = mean - self.std = std - - def __call__(self, image, target=None): - image = F.normalize(image, mean=self.mean, std=self.std) - if target is None: - return image, None - target = target.copy() - h, w = image.shape[-2:] - if "boxes" in target: - boxes = target["boxes"] - boxes = box_xyxy_to_cxcywh(boxes) - boxes = boxes / torch.tensor([w, h, w, h], dtype=torch.float32) - target["boxes"] = boxes - if "input_boxes" in target: - boxes = target["input_boxes"] - boxes = box_xyxy_to_cxcywh(boxes) - boxes = boxes / torch.tensor([w, h, w, h], dtype=torch.float32) - target["input_boxes"] = boxes - return image, target - - -class RemoveDifficult: - def __init__(self, enabled=False): - self.remove_difficult = enabled - - def __call__(self, image, target=None): - if target is None: - return image, None - target = target.copy() - keep = ~target["iscrowd"].to(torch.bool) | (not self.remove_difficult) - if "boxes" in target: - target["boxes"] = target["boxes"][keep] - target["labels"] = target["labels"][keep] - target["iscrowd"] = target["iscrowd"][keep] - return image, target - - -class Compose: - def __init__(self, transforms): - self.transforms = transforms - - def __call__(self, image, target): - for t in self.transforms: - image, target = t(image, target) - return image, target - - def __repr__(self): - format_string = self.__class__.__name__ + "(" - for t in self.transforms: - format_string += "\n" - format_string += " {0}".format(t) - format_string += "\n)" - return format_string - - -def get_random_resize_scales(size, min_size, rounded): - stride = 128 if rounded else 32 - min_size = int(stride * math.ceil(min_size / stride)) - scales = list(range(min_size, size + 1, stride)) - return scales - - -def get_random_resize_max_size(size, ratio=5 / 3): - max_size = round(ratio * size) - return max_size diff --git a/train/transforms/basic_for_api.py b/train/transforms/basic_for_api.py deleted file mode 100644 index e0ec2af..0000000 --- a/train/transforms/basic_for_api.py +++ /dev/null @@ -1,1396 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved - -""" -Transforms and data augmentation for both image + bbox. -""" - -import logging - -import numbers -import random -from collections.abc import Sequence -from typing import Iterable - -import torch -import torchvision.transforms as T -import torchvision.transforms.functional as F -import torchvision.transforms.v2.functional as Fv2 - -from PIL import Image as PILImage - -from sam3.model.box_ops import box_xyxy_to_cxcywh, masks_to_boxes -from sam3.train.data.sam3_image_dataset import Datapoint -from torchvision.transforms import InterpolationMode - - -def crop( - datapoint, - index, - region, - v2=False, - check_validity=True, - check_input_validity=True, - recompute_box_from_mask=False, -): - if v2: - rtop, rleft, rheight, rwidth = (int(round(r)) for r in region) - datapoint.images[index].data = Fv2.crop( - datapoint.images[index].data, - top=rtop, - left=rleft, - height=rheight, - width=rwidth, - ) - else: - datapoint.images[index].data = F.crop(datapoint.images[index].data, *region) - - i, j, h, w = region - - # should we do something wrt the original size? - datapoint.images[index].size = (h, w) - - for obj in datapoint.images[index].objects: - # crop the mask - if obj.segment is not None: - obj.segment = F.crop(obj.segment, int(i), int(j), int(h), int(w)) - - # crop the bounding box - if recompute_box_from_mask and obj.segment is not None: - # here the boxes are still in XYXY format with absolute coordinates (they are - # converted to CxCyWH with relative coordinates in basic_for_api.NormalizeAPI) - obj.bbox, obj.area = get_bbox_xyxy_abs_coords_from_mask(obj.segment) - else: - if recompute_box_from_mask and obj.segment is None and obj.area > 0: - logging.warning( - "Cannot recompute bounding box from mask since `obj.segment` is None. " - "Falling back to directly cropping from the input bounding box." - ) - boxes = obj.bbox.view(1, 4) - max_size = torch.as_tensor([w, h], dtype=torch.float32) - cropped_boxes = boxes - torch.as_tensor([j, i, j, i], dtype=torch.float32) - cropped_boxes = torch.min(cropped_boxes.reshape(-1, 2, 2), max_size) - cropped_boxes = cropped_boxes.clamp(min=0) - obj.area = (cropped_boxes[:, 1, :] - cropped_boxes[:, 0, :]).prod(dim=1) - obj.bbox = cropped_boxes.reshape(-1, 4) - - for query in datapoint.find_queries: - if query.semantic_target is not None: - query.semantic_target = F.crop( - query.semantic_target, int(i), int(j), int(h), int(w) - ) - if query.image_id == index and query.input_bbox is not None: - boxes = query.input_bbox - max_size = torch.as_tensor([w, h], dtype=torch.float32) - cropped_boxes = boxes - torch.as_tensor([j, i, j, i], dtype=torch.float32) - cropped_boxes = torch.min(cropped_boxes.reshape(-1, 2, 2), max_size) - cropped_boxes = cropped_boxes.clamp(min=0) - - # cur_area = (cropped_boxes[:, 1, :] - cropped_boxes[:, 0, :]).prod(dim=1) - # if check_input_validity: - # assert ( - # (cur_area > 0).all().item() - # ), "Some input box got cropped out by the crop transform" - - query.input_bbox = cropped_boxes.reshape(-1, 4) - if query.image_id == index and query.input_points is not None: - print( - "Warning! Point cropping with this function may lead to unexpected results" - ) - points = query.input_points - # Unlike right-lower box edges, which are exclusive, the - # point must be in [0, length-1], hence the -1 - max_size = torch.as_tensor([w, h], dtype=torch.float32) - 1 - cropped_points = points - torch.as_tensor([j, i, 0], dtype=torch.float32) - cropped_points[:, :, :2] = torch.min(cropped_points[:, :, :2], max_size) - cropped_points[:, :, :2] = cropped_points[:, :, :2].clamp(min=0) - query.input_points = cropped_points - - if check_validity: - # Check that all boxes are still valid - for obj in datapoint.images[index].objects: - assert obj.area > 0, "Box {} has no area".format(obj.bbox) - - return datapoint - - -def hflip(datapoint, index): - datapoint.images[index].data = F.hflip(datapoint.images[index].data) - - w, h = datapoint.images[index].data.size - for obj in datapoint.images[index].objects: - boxes = obj.bbox.view(1, 4) - boxes = boxes[:, [2, 1, 0, 3]] * torch.as_tensor( - [-1, 1, -1, 1] - ) + torch.as_tensor([w, 0, w, 0]) - obj.bbox = boxes - if obj.segment is not None: - obj.segment = F.hflip(obj.segment) - - for query in datapoint.find_queries: - if query.semantic_target is not None: - query.semantic_target = F.hflip(query.semantic_target) - if query.image_id == index and query.input_bbox is not None: - boxes = query.input_bbox - boxes = boxes[:, [2, 1, 0, 3]] * torch.as_tensor( - [-1, 1, -1, 1] - ) + torch.as_tensor([w, 0, w, 0]) - query.input_bbox = boxes - if query.image_id == index and query.input_points is not None: - points = query.input_points - points = points * torch.as_tensor([-1, 1, 1]) + torch.as_tensor([w, 0, 0]) - query.input_points = points - return datapoint - - -def get_size_with_aspect_ratio(image_size, size, max_size=None): - w, h = image_size - if max_size is not None: - min_original_size = float(min((w, h))) - max_original_size = float(max((w, h))) - if max_original_size / min_original_size * size > max_size: - size = max_size * min_original_size / max_original_size - - if (w <= h and w == size) or (h <= w and h == size): - return (h, w) - - if w < h: - ow = int(round(size)) - oh = int(round(size * h / w)) - else: - oh = int(round(size)) - ow = int(round(size * w / h)) - - return (oh, ow) - - -def resize(datapoint, index, size, max_size=None, square=False, v2=False): - # size can be min_size (scalar) or (w, h) tuple - - def get_size(image_size, size, max_size=None): - if isinstance(size, (list, tuple)): - return size[::-1] - else: - return get_size_with_aspect_ratio(image_size, size, max_size) - - if square: - size = size, size - else: - cur_size = ( - datapoint.images[index].data.size()[-2:][::-1] - if v2 - else datapoint.images[index].data.size - ) - size = get_size(cur_size, size, max_size) - - old_size = ( - datapoint.images[index].data.size()[-2:][::-1] - if v2 - else datapoint.images[index].data.size - ) - if v2: - datapoint.images[index].data = Fv2.resize( - datapoint.images[index].data, size, antialias=True - ) - else: - datapoint.images[index].data = F.resize(datapoint.images[index].data, size) - - new_size = ( - datapoint.images[index].data.size()[-2:][::-1] - if v2 - else datapoint.images[index].data.size - ) - ratios = tuple(float(s) / float(s_orig) for s, s_orig in zip(new_size, old_size)) - ratio_width, ratio_height = ratios - - for obj in datapoint.images[index].objects: - boxes = obj.bbox.view(1, 4) - scaled_boxes = boxes * torch.as_tensor( - [ratio_width, ratio_height, ratio_width, ratio_height], dtype=torch.float32 - ) - obj.bbox = scaled_boxes - obj.area *= ratio_width * ratio_height - if obj.segment is not None: - obj.segment = F.resize(obj.segment[None, None], size).squeeze() - - for query in datapoint.find_queries: - if query.semantic_target is not None: - query.semantic_target = F.resize( - query.semantic_target[None, None], size - ).squeeze() - if query.image_id == index and query.input_bbox is not None: - boxes = query.input_bbox - scaled_boxes = boxes * torch.as_tensor( - [ratio_width, ratio_height, ratio_width, ratio_height], - dtype=torch.float32, - ) - query.input_bbox = scaled_boxes - if query.image_id == index and query.input_points is not None: - points = query.input_points - scaled_points = points * torch.as_tensor( - [ratio_width, ratio_height, 1], - dtype=torch.float32, - ) - query.input_points = scaled_points - - h, w = size - datapoint.images[index].size = (h, w) - return datapoint - - -def pad(datapoint, index, padding, v2=False): - old_h, old_w = datapoint.images[index].size - h, w = old_h, old_w - if len(padding) == 2: - # assumes that we only pad on the bottom right corners - if v2: - datapoint.images[index].data = Fv2.pad( - datapoint.images[index].data, (0, 0, padding[0], padding[1]) - ) - else: - datapoint.images[index].data = F.pad( - datapoint.images[index].data, (0, 0, padding[0], padding[1]) - ) - h += padding[1] - w += padding[0] - else: - if v2: - # left, top, right, bottom - datapoint.images[index].data = Fv2.pad( - datapoint.images[index].data, - (padding[0], padding[1], padding[2], padding[3]), - ) - else: - # left, top, right, bottom - datapoint.images[index].data = F.pad( - datapoint.images[index].data, - (padding[0], padding[1], padding[2], padding[3]), - ) - h += padding[1] + padding[3] - w += padding[0] + padding[2] - - datapoint.images[index].size = (h, w) - - for obj in datapoint.images[index].objects: - if len(padding) != 2: - obj.bbox += torch.as_tensor( - [padding[0], padding[1], padding[0], padding[1]], dtype=torch.float32 - ) - if obj.segment is not None: - if v2: - if len(padding) == 2: - obj.segment = Fv2.pad( - obj.segment[None], (0, 0, padding[0], padding[1]) - ).squeeze(0) - else: - obj.segment = Fv2.pad(obj.segment[None], tuple(padding)).squeeze(0) - else: - if len(padding) == 2: - obj.segment = F.pad(obj.segment, (0, 0, padding[0], padding[1])) - else: - obj.segment = F.pad(obj.segment, tuple(padding)) - - for query in datapoint.find_queries: - if query.semantic_target is not None: - if v2: - if len(padding) == 2: - query.semantic_target = Fv2.pad( - query.semantic_target[None, None], - (0, 0, padding[0], padding[1]), - ).squeeze() - else: - query.semantic_target = Fv2.pad( - query.semantic_target[None, None], tuple(padding) - ).squeeze() - else: - if len(padding) == 2: - query.semantic_target = F.pad( - query.semantic_target[None, None], - (0, 0, padding[0], padding[1]), - ).squeeze() - else: - query.semantic_target = F.pad( - query.semantic_target[None, None], tuple(padding) - ).squeeze() - if query.image_id == index and query.input_bbox is not None: - if len(padding) != 2: - query.input_bbox += torch.as_tensor( - [padding[0], padding[1], padding[0], padding[1]], - dtype=torch.float32, - ) - if query.image_id == index and query.input_points is not None: - if len(padding) != 2: - query.input_points += torch.as_tensor( - [padding[0], padding[1], 0], dtype=torch.float32 - ) - - return datapoint - - -class RandomSizeCropAPI: - def __init__( - self, - min_size: int, - max_size: int, - respect_boxes: bool, - consistent_transform: bool, - respect_input_boxes: bool = True, - v2: bool = False, - recompute_box_from_mask: bool = False, - ): - self.min_size = min_size - self.max_size = max_size - self.respect_boxes = respect_boxes # if True we can't crop a box out - self.respect_input_boxes = respect_input_boxes - self.consistent_transform = consistent_transform - self.v2 = v2 - self.recompute_box_from_mask = recompute_box_from_mask - - def _sample_no_respect_boxes(self, img): - w = random.randint(self.min_size, min(img.width, self.max_size)) - h = random.randint(self.min_size, min(img.height, self.max_size)) - return T.RandomCrop.get_params(img, (h, w)) - - def _sample_respect_boxes(self, img, boxes, points, min_box_size=10.0): - """ - Assure that no box or point is dropped via cropping, though portions - of boxes may be removed. - """ - if len(boxes) == 0 and len(points) == 0: - return self._sample_no_respect_boxes(img) - - if self.v2: - img_height, img_width = img.size()[-2:] - else: - img_width, img_height = img.size - - minW, minH, maxW, maxH = ( - min(img_width, self.min_size), - min(img_height, self.min_size), - min(img_width, self.max_size), - min(img_height, self.max_size), - ) - - # The crop box must extend one pixel beyond points to the bottom/right - # to assure the exclusive box contains the points. - minX = ( - torch.cat([boxes[:, 0] + min_box_size, points[:, 0] + 1], dim=0) - .max() - .item() - ) - minY = ( - torch.cat([boxes[:, 1] + min_box_size, points[:, 1] + 1], dim=0) - .max() - .item() - ) - minX = min(img_width, minX) - minY = min(img_height, minY) - maxX = torch.cat([boxes[:, 2] - min_box_size, points[:, 0]], dim=0).min().item() - maxY = torch.cat([boxes[:, 3] - min_box_size, points[:, 1]], dim=0).min().item() - maxX = max(0.0, maxX) - maxY = max(0.0, maxY) - minW = max(minW, minX - maxX) - minH = max(minH, minY - maxY) - w = random.uniform(minW, max(minW, maxW)) - h = random.uniform(minH, max(minH, maxH)) - if minX > maxX: - # i = random.uniform(max(0, minX - w + 1), max(maxX, max(0, minX - w + 1))) - i = random.uniform(max(0, minX - w), max(maxX, max(0, minX - w))) - else: - i = random.uniform( - max(0, minX - w + 1), max(maxX - 1, max(0, minX - w + 1)) - ) - if minY > maxY: - # j = random.uniform(max(0, minY - h + 1), max(maxY, max(0, minY - h + 1))) - j = random.uniform(max(0, minY - h), max(maxY, max(0, minY - h))) - else: - j = random.uniform( - max(0, minY - h + 1), max(maxY - 1, max(0, minY - h + 1)) - ) - - return [j, i, h, w] - - def __call__(self, datapoint, **kwargs): - if self.respect_boxes or self.respect_input_boxes: - if self.consistent_transform: - # Check that all the images are the same size - w, h = datapoint.images[0].data.size - for img in datapoint.images: - assert img.data.size == (w, h) - - all_boxes = [] - # Getting all boxes in all the images - if self.respect_boxes: - all_boxes += [ - obj.bbox.view(-1, 4) - for img in datapoint.images - for obj in img.objects - ] - # Get all the boxes in the find queries - if self.respect_input_boxes: - all_boxes += [ - q.input_bbox.view(-1, 4) - for q in datapoint.find_queries - if q.input_bbox is not None - ] - if all_boxes: - all_boxes = torch.cat(all_boxes, 0) - else: - all_boxes = torch.empty(0, 4) - - all_points = [ - q.input_points.view(-1, 3)[:, :2] - for q in datapoint.find_queries - if q.input_points is not None - ] - if all_points: - all_points = torch.cat(all_points, 0) - else: - all_points = torch.empty(0, 2) - - crop_param = self._sample_respect_boxes( - datapoint.images[0].data, all_boxes, all_points - ) - for i in range(len(datapoint.images)): - datapoint = crop( - datapoint, - i, - crop_param, - v2=self.v2, - check_validity=self.respect_boxes, - check_input_validity=self.respect_input_boxes, - recompute_box_from_mask=self.recompute_box_from_mask, - ) - return datapoint - else: - for i in range(len(datapoint.images)): - all_boxes = [] - # Get all boxes in the current image - if self.respect_boxes: - all_boxes += [ - obj.bbox.view(-1, 4) for obj in datapoint.images[i].objects - ] - # Get all the boxes in the find queries that correspond to this image - if self.respect_input_boxes: - all_boxes += [ - q.input_bbox.view(-1, 4) - for q in datapoint.find_queries - if q.image_id == i and q.input_bbox is not None - ] - if all_boxes: - all_boxes = torch.cat(all_boxes, 0) - else: - all_boxes = torch.empty(0, 4) - - all_points = [ - q.input_points.view(-1, 3)[:, :2] - for q in datapoint.find_queries - if q.input_points is not None - ] - if all_points: - all_points = torch.cat(all_points, 0) - else: - all_points = torch.empty(0, 2) - - crop_param = self._sample_respect_boxes( - datapoint.images[i].data, all_boxes, all_points - ) - datapoint = crop( - datapoint, - i, - crop_param, - v2=self.v2, - check_validity=self.respect_boxes, - check_input_validity=self.respect_input_boxes, - recompute_box_from_mask=self.recompute_box_from_mask, - ) - return datapoint - else: - if self.consistent_transform: - # Check that all the images are the same size - w, h = datapoint.images[0].data.size - for img in datapoint.images: - assert img.data.size == (w, h) - - crop_param = self._sample_no_respect_boxes(datapoint.images[0].data) - for i in range(len(datapoint.images)): - datapoint = crop( - datapoint, - i, - crop_param, - v2=self.v2, - check_validity=self.respect_boxes, - check_input_validity=self.respect_input_boxes, - recompute_box_from_mask=self.recompute_box_from_mask, - ) - return datapoint - else: - for i in range(len(datapoint.images)): - crop_param = self._sample_no_respect_boxes(datapoint.images[i].data) - datapoint = crop( - datapoint, - i, - crop_param, - v2=self.v2, - check_validity=self.respect_boxes, - check_input_validity=self.respect_input_boxes, - recompute_box_from_mask=self.recompute_box_from_mask, - ) - return datapoint - - -class CenterCropAPI: - def __init__(self, size, consistent_transform, recompute_box_from_mask=False): - self.size = size - self.consistent_transform = consistent_transform - self.recompute_box_from_mask = recompute_box_from_mask - - def _sample_crop(self, image_width, image_height): - crop_height, crop_width = self.size - crop_top = int(round((image_height - crop_height) / 2.0)) - crop_left = int(round((image_width - crop_width) / 2.0)) - return crop_top, crop_left, crop_height, crop_width - - def __call__(self, datapoint, **kwargs): - if self.consistent_transform: - # Check that all the images are the same size - w, h = datapoint.images[0].data.size - for img in datapoint.images: - assert img.size == (w, h) - - crop_top, crop_left, crop_height, crop_width = self._sample_crop(w, h) - for i in range(len(datapoint.images)): - datapoint = crop( - datapoint, - i, - (crop_top, crop_left, crop_height, crop_width), - recompute_box_from_mask=self.recompute_box_from_mask, - ) - return datapoint - - for i in range(len(datapoint.images)): - w, h = datapoint.images[i].data.size - crop_top, crop_left, crop_height, crop_width = self._sample_crop(w, h) - datapoint = crop( - datapoint, - i, - (crop_top, crop_left, crop_height, crop_width), - recompute_box_from_mask=self.recompute_box_from_mask, - ) - - return datapoint - - -class RandomHorizontalFlip: - def __init__(self, consistent_transform, p=0.5): - self.p = p - self.consistent_transform = consistent_transform - - def __call__(self, datapoint, **kwargs): - if self.consistent_transform: - if random.random() < self.p: - for i in range(len(datapoint.images)): - datapoint = hflip(datapoint, i) - return datapoint - for i in range(len(datapoint.images)): - if random.random() < self.p: - datapoint = hflip(datapoint, i) - return datapoint - - -class RandomResizeAPI: - def __init__( - self, sizes, consistent_transform, max_size=None, square=False, v2=False - ): - if isinstance(sizes, int): - sizes = (sizes,) - assert isinstance(sizes, Iterable) - self.sizes = list(sizes) - self.max_size = max_size - self.square = square - self.consistent_transform = consistent_transform - self.v2 = v2 - - def __call__(self, datapoint, **kwargs): - if self.consistent_transform: - size = random.choice(self.sizes) - for i in range(len(datapoint.images)): - datapoint = resize( - datapoint, i, size, self.max_size, square=self.square, v2=self.v2 - ) - return datapoint - for i in range(len(datapoint.images)): - size = random.choice(self.sizes) - datapoint = resize( - datapoint, i, size, self.max_size, square=self.square, v2=self.v2 - ) - return datapoint - - -class ScheduledRandomResizeAPI(RandomResizeAPI): - def __init__(self, size_scheduler, consistent_transform, square=False): - self.size_scheduler = size_scheduler - # Just a meaningful init value for super - params = self.size_scheduler(epoch_num=0) - sizes, max_size = params["sizes"], params["max_size"] - super().__init__(sizes, consistent_transform, max_size=max_size, square=square) - - def __call__(self, datapoint, **kwargs): - assert "epoch" in kwargs, "Param scheduler needs to know the current epoch" - params = self.size_scheduler(kwargs["epoch"]) - sizes, max_size = params["sizes"], params["max_size"] - self.sizes = sizes - self.max_size = max_size - datapoint = super(ScheduledRandomResizeAPI, self).__call__(datapoint, **kwargs) - return datapoint - - -class RandomPadAPI: - def __init__(self, max_pad, consistent_transform): - self.max_pad = max_pad - self.consistent_transform = consistent_transform - - def _sample_pad(self): - pad_x = random.randint(0, self.max_pad) - pad_y = random.randint(0, self.max_pad) - return pad_x, pad_y - - def __call__(self, datapoint, **kwargs): - if self.consistent_transform: - pad_x, pad_y = self._sample_pad() - for i in range(len(datapoint.images)): - datapoint = pad(datapoint, i, (pad_x, pad_y)) - return datapoint - - for i in range(len(datapoint.images)): - pad_x, pad_y = self._sample_pad() - datapoint = pad(datapoint, i, (pad_x, pad_y)) - return datapoint - - -class PadToSizeAPI: - def __init__(self, size, consistent_transform, bottom_right=False, v2=False): - self.size = size - self.consistent_transform = consistent_transform - self.v2 = v2 - self.bottom_right = bottom_right - - def _sample_pad(self, w, h): - pad_x = self.size - w - pad_y = self.size - h - assert pad_x >= 0 and pad_y >= 0 - pad_left = random.randint(0, pad_x) - pad_right = pad_x - pad_left - pad_top = random.randint(0, pad_y) - pad_bottom = pad_y - pad_top - return pad_left, pad_top, pad_right, pad_bottom - - def __call__(self, datapoint, **kwargs): - if self.consistent_transform: - # Check that all the images are the same size - w, h = datapoint.images[0].data.size - for img in datapoint.images: - assert img.size == (w, h) - if self.bottom_right: - pad_right = self.size - w - pad_bottom = self.size - h - padding = (pad_right, pad_bottom) - else: - padding = self._sample_pad(w, h) - for i in range(len(datapoint.images)): - datapoint = pad(datapoint, i, padding, v2=self.v2) - return datapoint - - for i, img in enumerate(datapoint.images): - w, h = img.data.size - if self.bottom_right: - pad_right = self.size - w - pad_bottom = self.size - h - padding = (pad_right, pad_bottom) - else: - padding = self._sample_pad(w, h) - datapoint = pad(datapoint, i, padding, v2=self.v2) - return datapoint - - -class RandomMosaicVideoAPI: - def __init__(self, prob=0.15, grid_h=2, grid_w=2, use_random_hflip=False): - self.prob = prob - self.grid_h = grid_h - self.grid_w = grid_w - self.use_random_hflip = use_random_hflip - - def __call__(self, datapoint, **kwargs): - if random.random() > self.prob: - return datapoint - - # select a random location to place the target mask in the mosaic - target_grid_y = random.randint(0, self.grid_h - 1) - target_grid_x = random.randint(0, self.grid_w - 1) - # whether to flip each grid in the mosaic horizontally - if self.use_random_hflip: - should_hflip = torch.rand(self.grid_h, self.grid_w) < 0.5 - else: - should_hflip = torch.zeros(self.grid_h, self.grid_w, dtype=torch.bool) - for i in range(len(datapoint.images)): - datapoint = random_mosaic_frame( - datapoint, - i, - grid_h=self.grid_h, - grid_w=self.grid_w, - target_grid_y=target_grid_y, - target_grid_x=target_grid_x, - should_hflip=should_hflip, - ) - - return datapoint - - -def random_mosaic_frame( - datapoint, - index, - grid_h, - grid_w, - target_grid_y, - target_grid_x, - should_hflip, -): - # Step 1: downsize the images and paste them into a mosaic - image_data = datapoint.images[index].data - is_pil = isinstance(image_data, PILImage.Image) - if is_pil: - H_im = image_data.height - W_im = image_data.width - image_data_output = PILImage.new("RGB", (W_im, H_im)) - else: - H_im = image_data.size(-2) - W_im = image_data.size(-1) - image_data_output = torch.zeros_like(image_data) - - downsize_cache = {} - for grid_y in range(grid_h): - for grid_x in range(grid_w): - y_offset_b = grid_y * H_im // grid_h - x_offset_b = grid_x * W_im // grid_w - y_offset_e = (grid_y + 1) * H_im // grid_h - x_offset_e = (grid_x + 1) * W_im // grid_w - H_im_downsize = y_offset_e - y_offset_b - W_im_downsize = x_offset_e - x_offset_b - - if (H_im_downsize, W_im_downsize) in downsize_cache: - image_data_downsize = downsize_cache[(H_im_downsize, W_im_downsize)] - else: - image_data_downsize = F.resize( - image_data, - size=(H_im_downsize, W_im_downsize), - interpolation=InterpolationMode.BILINEAR, - antialias=True, # antialiasing for downsizing - ) - downsize_cache[(H_im_downsize, W_im_downsize)] = image_data_downsize - if should_hflip[grid_y, grid_x].item(): - image_data_downsize = F.hflip(image_data_downsize) - - if is_pil: - image_data_output.paste(image_data_downsize, (x_offset_b, y_offset_b)) - else: - image_data_output[:, y_offset_b:y_offset_e, x_offset_b:x_offset_e] = ( - image_data_downsize - ) - - datapoint.images[index].data = image_data_output - - # Step 2: downsize the masks and paste them into the target grid of the mosaic - # (note that we don't scale input/target boxes since they are not used in TA) - for obj in datapoint.images[index].objects: - if obj.segment is None: - continue - assert obj.segment.shape == (H_im, W_im) and obj.segment.dtype == torch.uint8 - segment_output = torch.zeros_like(obj.segment) - - target_y_offset_b = target_grid_y * H_im // grid_h - target_x_offset_b = target_grid_x * W_im // grid_w - target_y_offset_e = (target_grid_y + 1) * H_im // grid_h - target_x_offset_e = (target_grid_x + 1) * W_im // grid_w - target_H_im_downsize = target_y_offset_e - target_y_offset_b - target_W_im_downsize = target_x_offset_e - target_x_offset_b - - segment_downsize = F.resize( - obj.segment[None, None], - size=(target_H_im_downsize, target_W_im_downsize), - interpolation=InterpolationMode.BILINEAR, - antialias=True, # antialiasing for downsizing - )[0, 0] - if should_hflip[target_grid_y, target_grid_x].item(): - segment_downsize = F.hflip(segment_downsize[None, None])[0, 0] - - segment_output[ - target_y_offset_b:target_y_offset_e, target_x_offset_b:target_x_offset_e - ] = segment_downsize - obj.segment = segment_output - - return datapoint - - -class ScheduledPadToSizeAPI(PadToSizeAPI): - def __init__(self, size_scheduler, consistent_transform): - self.size_scheduler = size_scheduler - size = self.size_scheduler(epoch_num=0)["sizes"] - super().__init__(size, consistent_transform) - - def __call__(self, datapoint, **kwargs): - assert "epoch" in kwargs, "Param scheduler needs to know the current epoch" - params = self.size_scheduler(kwargs["epoch"]) - self.size = params["resolution"] - return super(ScheduledPadToSizeAPI, self).__call__(datapoint, **kwargs) - - -class IdentityAPI: - def __call__(self, datapoint, **kwargs): - return datapoint - - -class RandomSelectAPI: - """ - Randomly selects between transforms1 and transforms2, - with probability p for transforms1 and (1 - p) for transforms2 - """ - - def __init__(self, transforms1=None, transforms2=None, p=0.5): - self.transforms1 = transforms1 or IdentityAPI() - self.transforms2 = transforms2 or IdentityAPI() - self.p = p - - def __call__(self, datapoint, **kwargs): - if random.random() < self.p: - return self.transforms1(datapoint, **kwargs) - return self.transforms2(datapoint, **kwargs) - - -class ToTensorAPI: - def __init__(self, v2=False): - self.v2 = v2 - - def __call__(self, datapoint: Datapoint, **kwargs): - for img in datapoint.images: - if self.v2: - img.data = Fv2.to_image_tensor(img.data) - # img.data = Fv2.to_dtype(img.data, torch.uint8, scale=True) - # img.data = Fv2.convert_image_dtype(img.data, torch.uint8) - else: - img.data = F.to_tensor(img.data) - return datapoint - - -class NormalizeAPI: - def __init__(self, mean, std, v2=False): - self.mean = mean - self.std = std - self.v2 = v2 - - def __call__(self, datapoint: Datapoint, **kwargs): - for img in datapoint.images: - if self.v2: - img.data = Fv2.convert_image_dtype(img.data, torch.float32) - img.data = Fv2.normalize(img.data, mean=self.mean, std=self.std) - else: - img.data = F.normalize(img.data, mean=self.mean, std=self.std) - for obj in img.objects: - boxes = obj.bbox - cur_h, cur_w = img.data.shape[-2:] - boxes = box_xyxy_to_cxcywh(boxes) - boxes = boxes / torch.tensor( - [cur_w, cur_h, cur_w, cur_h], dtype=torch.float32 - ) - obj.bbox = boxes - - for query in datapoint.find_queries: - if query.input_bbox is not None: - boxes = query.input_bbox - cur_h, cur_w = datapoint.images[query.image_id].data.shape[-2:] - boxes = box_xyxy_to_cxcywh(boxes) - boxes = boxes / torch.tensor( - [cur_w, cur_h, cur_w, cur_h], dtype=torch.float32 - ) - query.input_bbox = boxes - if query.input_points is not None: - points = query.input_points - cur_h, cur_w = datapoint.images[query.image_id].data.shape[-2:] - points = points / torch.tensor([cur_w, cur_h, 1.0], dtype=torch.float32) - query.input_points = points - - return datapoint - - -class ComposeAPI: - def __init__(self, transforms): - self.transforms = transforms - - def __call__(self, datapoint, **kwargs): - for t in self.transforms: - datapoint = t(datapoint, **kwargs) - return datapoint - - def __repr__(self): - format_string = self.__class__.__name__ + "(" - for t in self.transforms: - format_string += "\n" - format_string += " {0}".format(t) - format_string += "\n)" - return format_string - - -class RandomGrayscale: - def __init__(self, consistent_transform, p=0.5): - self.p = p - self.consistent_transform = consistent_transform - self.Grayscale = T.Grayscale(num_output_channels=3) - - def __call__(self, datapoint: Datapoint, **kwargs): - if self.consistent_transform: - if random.random() < self.p: - for img in datapoint.images: - img.data = self.Grayscale(img.data) - return datapoint - for img in datapoint.images: - if random.random() < self.p: - img.data = self.Grayscale(img.data) - return datapoint - - -class ColorJitter: - def __init__(self, consistent_transform, brightness, contrast, saturation, hue): - self.consistent_transform = consistent_transform - self.brightness = ( - brightness - if isinstance(brightness, list) - else [max(0, 1 - brightness), 1 + brightness] - ) - self.contrast = ( - contrast - if isinstance(contrast, list) - else [max(0, 1 - contrast), 1 + contrast] - ) - self.saturation = ( - saturation - if isinstance(saturation, list) - else [max(0, 1 - saturation), 1 + saturation] - ) - self.hue = hue if isinstance(hue, list) or hue is None else ([-hue, hue]) - - def __call__(self, datapoint: Datapoint, **kwargs): - if self.consistent_transform: - # Create a color jitter transformation params - ( - fn_idx, - brightness_factor, - contrast_factor, - saturation_factor, - hue_factor, - ) = T.ColorJitter.get_params( - self.brightness, self.contrast, self.saturation, self.hue - ) - for img in datapoint.images: - if not self.consistent_transform: - ( - fn_idx, - brightness_factor, - contrast_factor, - saturation_factor, - hue_factor, - ) = T.ColorJitter.get_params( - self.brightness, self.contrast, self.saturation, self.hue - ) - for fn_id in fn_idx: - if fn_id == 0 and brightness_factor is not None: - img.data = F.adjust_brightness(img.data, brightness_factor) - elif fn_id == 1 and contrast_factor is not None: - img.data = F.adjust_contrast(img.data, contrast_factor) - elif fn_id == 2 and saturation_factor is not None: - img.data = F.adjust_saturation(img.data, saturation_factor) - elif fn_id == 3 and hue_factor is not None: - img.data = F.adjust_hue(img.data, hue_factor) - return datapoint - - -class RandomAffine: - def __init__( - self, - degrees, - consistent_transform, - scale=None, - translate=None, - shear=None, - image_mean=(123, 116, 103), - log_warning=True, - num_tentatives=1, - image_interpolation="bicubic", - ): - """ - The mask is required for this transform. - if consistent_transform if True, then the same random affine is applied to all frames and masks. - """ - self.degrees = degrees if isinstance(degrees, list) else ([-degrees, degrees]) - self.scale = scale - self.shear = ( - shear if isinstance(shear, list) else ([-shear, shear] if shear else None) - ) - self.translate = translate - self.fill_img = image_mean - self.consistent_transform = consistent_transform - self.log_warning = log_warning - self.num_tentatives = num_tentatives - - if image_interpolation == "bicubic": - self.image_interpolation = InterpolationMode.BICUBIC - elif image_interpolation == "bilinear": - self.image_interpolation = InterpolationMode.BILINEAR - else: - raise NotImplementedError - - def __call__(self, datapoint: Datapoint, **kwargs): - for _tentative in range(self.num_tentatives): - res = self.transform_datapoint(datapoint) - if res is not None: - return res - - if self.log_warning: - logging.warning( - f"Skip RandomAffine for zero-area mask in first frame after {self.num_tentatives} tentatives" - ) - return datapoint - - def transform_datapoint(self, datapoint: Datapoint): - _, height, width = F.get_dimensions(datapoint.images[0].data) - img_size = [width, height] - - if self.consistent_transform: - # Create a random affine transformation - affine_params = T.RandomAffine.get_params( - degrees=self.degrees, - translate=self.translate, - scale_ranges=self.scale, - shears=self.shear, - img_size=img_size, - ) - - for img_idx, img in enumerate(datapoint.images): - this_masks = [ - obj.segment.unsqueeze(0) if obj.segment is not None else None - for obj in img.objects - ] - if not self.consistent_transform: - # if not consistent we create a new affine params for every frame&mask pair Create a random affine transformation - affine_params = T.RandomAffine.get_params( - degrees=self.degrees, - translate=self.translate, - scale_ranges=self.scale, - shears=self.shear, - img_size=img_size, - ) - - transformed_bboxes, transformed_masks = [], [] - for i in range(len(img.objects)): - if this_masks[i] is None: - transformed_masks.append(None) - # Dummy bbox for a dummy target - transformed_bboxes.append(torch.tensor([[0, 0, 0, 0]])) - else: - transformed_mask = F.affine( - this_masks[i], - *affine_params, - interpolation=InterpolationMode.NEAREST, - fill=0.0, - ) - if img_idx == 0 and transformed_mask.max() == 0: - # We are dealing with a video and the object is not visible in the first frame - # Return the datapoint without transformation - return None - transformed_bbox = masks_to_boxes(transformed_mask) - transformed_bboxes.append(transformed_bbox) - transformed_masks.append(transformed_mask.squeeze()) - - for i in range(len(img.objects)): - img.objects[i].bbox = transformed_bboxes[i] - img.objects[i].segment = transformed_masks[i] - - img.data = F.affine( - img.data, - *affine_params, - interpolation=self.image_interpolation, - fill=self.fill_img, - ) - return datapoint - - -class RandomResizedCrop: - def __init__( - self, - consistent_transform, - size, - scale=None, - ratio=None, - log_warning=True, - num_tentatives=4, - keep_aspect_ratio=False, - ): - """ - The mask is required for this transform. - if consistent_transform if True, then the same random resized crop is applied to all frames and masks. - """ - if isinstance(size, numbers.Number): - self.size = (int(size), int(size)) - elif isinstance(size, Sequence) and len(size) == 1: - self.size = (size[0], size[0]) - elif len(size) != 2: - raise ValueError("Please provide only two dimensions (h, w) for size.") - else: - self.size = size - - self.scale = scale if scale is not None else (0.08, 1.0) - self.ratio = ratio if ratio is not None else (3.0 / 4.0, 4.0 / 3.0) - self.consistent_transform = consistent_transform - self.log_warning = log_warning - self.num_tentatives = num_tentatives - self.keep_aspect_ratio = keep_aspect_ratio - - def __call__(self, datapoint: Datapoint, **kwargs): - for _tentative in range(self.num_tentatives): - res = self.transform_datapoint(datapoint) - if res is not None: - return res - - if self.log_warning: - logging.warning( - f"Skip RandomResizeCrop for zero-area mask in first frame after {self.num_tentatives} tentatives" - ) - return datapoint - - def transform_datapoint(self, datapoint: Datapoint): - if self.keep_aspect_ratio: - original_size = datapoint.images[0].size - original_ratio = original_size[1] / original_size[0] - ratio = [r * original_ratio for r in self.ratio] - else: - ratio = self.ratio - - if self.consistent_transform: - # Create a random crop transformation - crop_params = T.RandomResizedCrop.get_params( - img=datapoint.images[0].data, - scale=self.scale, - ratio=ratio, - ) - - for img_idx, img in enumerate(datapoint.images): - if not self.consistent_transform: - # Create a random crop transformation - crop_params = T.RandomResizedCrop.get_params( - img=img.data, - scale=self.scale, - ratio=ratio, - ) - - this_masks = [ - obj.segment.unsqueeze(0) if obj.segment is not None else None - for obj in img.objects - ] - - transformed_bboxes, transformed_masks = [], [] - for i in range(len(img.objects)): - if this_masks[i] is None: - transformed_masks.append(None) - # Dummy bbox for a dummy target - transformed_bboxes.append(torch.tensor([[0, 0, 0, 0]])) - else: - transformed_mask = F.resized_crop( - this_masks[i], - *crop_params, - size=self.size, - interpolation=InterpolationMode.NEAREST, - ) - if img_idx == 0 and transformed_mask.max() == 0: - # We are dealing with a video and the object is not visible in the first frame - # Return the datapoint without transformation - return None - transformed_masks.append(transformed_mask.squeeze()) - transformed_bbox = masks_to_boxes(transformed_mask) - transformed_bboxes.append(transformed_bbox) - - # Set the new boxes and masks if all transformed masks and boxes are good. - for i in range(len(img.objects)): - img.objects[i].bbox = transformed_bboxes[i] - img.objects[i].segment = transformed_masks[i] - - img.data = F.resized_crop( - img.data, - *crop_params, - size=self.size, - interpolation=InterpolationMode.BILINEAR, - ) - return datapoint - - -class ResizeToMaxIfAbove: - # Resize datapoint image if one of its sides is larger that max_size - def __init__( - self, - max_size=None, - ): - self.max_size = max_size - - def __call__(self, datapoint: Datapoint, **kwargs): - _, height, width = F.get_dimensions(datapoint.images[0].data) - - if height <= self.max_size and width <= self.max_size: - # The original frames are small enough - return datapoint - elif height >= width: - new_height = self.max_size - new_width = int(round(self.max_size * width / height)) - else: - new_height = int(round(self.max_size * height / width)) - new_width = self.max_size - - size = new_height, new_width - - for index in range(len(datapoint.images)): - datapoint.images[index].data = F.resize(datapoint.images[index].data, size) - - for obj in datapoint.images[index].objects: - obj.segment = F.resize( - obj.segment[None, None], - size, - interpolation=InterpolationMode.NEAREST, - ).squeeze() - - h, w = size - datapoint.images[index].size = (h, w) - return datapoint - - -def get_bbox_xyxy_abs_coords_from_mask(mask): - """Get the bounding box (XYXY format w/ absolute coordinates) of a binary mask.""" - assert mask.dim() == 2 - rows = torch.any(mask, dim=1) - cols = torch.any(mask, dim=0) - row_inds = rows.nonzero().view(-1) - col_inds = cols.nonzero().view(-1) - if row_inds.numel() == 0: - # mask is empty - bbox = torch.zeros(1, 4, dtype=torch.float32) - bbox_area = 0.0 - else: - ymin, ymax = row_inds.min(), row_inds.max() - xmin, xmax = col_inds.min(), col_inds.max() - bbox = torch.tensor([xmin, ymin, xmax, ymax], dtype=torch.float32).view(1, 4) - bbox_area = float((ymax - ymin) * (xmax - xmin)) - return bbox, bbox_area - - -class MotionBlur: - def __init__(self, kernel_size=5, consistent_transform=True, p=0.5): - assert kernel_size % 2 == 1, "Kernel size must be odd." - self.kernel_size = kernel_size - self.consistent_transform = consistent_transform - self.p = p - - def __call__(self, datapoint: Datapoint, **kwargs): - if random.random() >= self.p: - return datapoint - if self.consistent_transform: - # Generate a single motion blur kernel for all images - kernel = self._generate_motion_blur_kernel() - for img in datapoint.images: - if not self.consistent_transform: - # Generate a new motion blur kernel for each image - kernel = self._generate_motion_blur_kernel() - img.data = self._apply_motion_blur(img.data, kernel) - - return datapoint - - def _generate_motion_blur_kernel(self): - kernel = torch.zeros((self.kernel_size, self.kernel_size)) - direction = random.choice(["horizontal", "vertical", "diagonal"]) - if direction == "horizontal": - kernel[self.kernel_size // 2, :] = 1.0 - elif direction == "vertical": - kernel[:, self.kernel_size // 2] = 1.0 - elif direction == "diagonal": - for i in range(self.kernel_size): - kernel[i, i] = 1.0 - kernel /= kernel.sum() - return kernel - - def _apply_motion_blur(self, image, kernel): - if isinstance(image, PILImage.Image): - image = F.to_tensor(image) - channels = image.shape[0] - kernel = kernel.to(image.device).unsqueeze(0).unsqueeze(0) - blurred_image = torch.nn.functional.conv2d( - image.unsqueeze(0), - kernel.repeat(channels, 1, 1, 1), - padding=self.kernel_size // 2, - groups=channels, - ) - return F.to_pil_image(blurred_image.squeeze(0)) - - -class LargeScaleJitter: - def __init__( - self, - scale_range=(0.1, 2.0), - aspect_ratio_range=(0.75, 1.33), - crop_size=(640, 640), - consistent_transform=True, - p=0.5, - ): - """ - Args:rack - scale_range (tuple): Range of scaling factors (min_scale, max_scale). - aspect_ratio_range (tuple): Range of aspect ratios (min_aspect_ratio, max_aspect_ratio). - crop_size (tuple): Target size of the cropped region (width, height). - consistent_transform (bool): Whether to apply the same transformation across all frames. - p (float): Probability of applying the transformation. - """ - self.scale_range = scale_range - self.aspect_ratio_range = aspect_ratio_range - self.crop_size = crop_size - self.consistent_transform = consistent_transform - self.p = p - - def __call__(self, datapoint: Datapoint, **kwargs): - if random.random() >= self.p: - return datapoint - - # Sample a single scale factor and aspect ratio for all frames - log_ratio = torch.log(torch.tensor(self.aspect_ratio_range)) - scale_factor = torch.empty(1).uniform_(*self.scale_range).item() - aspect_ratio = torch.exp( - torch.empty(1).uniform_(log_ratio[0], log_ratio[1]) - ).item() - - for idx, img in enumerate(datapoint.images): - if not self.consistent_transform: - # Sample a new scale factor and aspect ratio for each frame - log_ratio = torch.log(torch.tensor(self.aspect_ratio_range)) - scale_factor = torch.empty(1).uniform_(*self.scale_range).item() - aspect_ratio = torch.exp( - torch.empty(1).uniform_(log_ratio[0], log_ratio[1]) - ).item() - - # Compute the dimensions of the jittered crop - original_width, original_height = img.data.size - target_area = original_width * original_height * scale_factor - crop_width = int(round((target_area * aspect_ratio) ** 0.5)) - crop_height = int(round((target_area / aspect_ratio) ** 0.5)) - - # Randomly select the top-left corner of the crop - crop_x = random.randint(0, max(0, original_width - crop_width)) - crop_y = random.randint(0, max(0, original_height - crop_height)) - - # Extract the cropped region - datapoint = crop(datapoint, idx, (crop_x, crop_y, crop_width, crop_height)) - - # Resize the cropped region to the target crop size - datapoint = resize(datapoint, idx, self.crop_size) - - return datapoint diff --git a/train/transforms/filter_query_transforms.py b/train/transforms/filter_query_transforms.py deleted file mode 100644 index 2d6708f..0000000 --- a/train/transforms/filter_query_transforms.py +++ /dev/null @@ -1,607 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved - -import logging -import random - -from collections import defaultdict -from typing import List, Optional, Union - -import torch - -from sam3.train.data.sam3_image_dataset import Datapoint, FindQuery, Object - - -class FilterDataPointQueries: - find_ids_to_filter: set = None - get_ids_to_filter: set = None - obj_ids_to_filter: set = None # stored as pairs (img_id, obj_id) - - def identify_queries_to_filter(self, datapoint: Datapoint) -> None: - """ - Compute set of query ids to keep, for both find and get queries - """ - raise NotImplementedError - - def _do_filter_query(self, query: Union[FindQuery], query_id: int): - assert self.find_ids_to_filter is not None - - return query_id in self.find_ids_to_filter - - -class FilterQueryWithText(FilterDataPointQueries): - """ - Filter all datapoints which have query text in a specified list of exluded terms - """ - - def __init__( - self, exclude_find_keys: List[str] = None, exclude_get_keys: List[str] = None - ): - self.find_filter_keys = exclude_find_keys if exclude_find_keys else [] - self.get_filter_keys = exclude_get_keys if exclude_get_keys else [] - - def identify_queries_to_filter(self, datapoint): - self.obj_ids_to_filter = set() - del_find_ids = [] - del_get_ids = [] - for i, f_q in enumerate(datapoint.find_queries): - if f_q.query_text in self.find_filter_keys: - del_find_ids.append(i) - - self.find_ids_to_filter = set(del_find_ids) - - -class KeepMaxNumFindQueries(FilterDataPointQueries): - def __init__( - self, max_num_find_queries: int, retain_positive_queries: bool = False - ): - self.max_num_find_queries = max_num_find_queries - self.retain_positive_queries = retain_positive_queries - - def identify_queries_to_filter(self, datapoint: Datapoint) -> None: - self.obj_ids_to_filter = set() - num_find_queries = len(datapoint.find_queries) - if num_find_queries <= self.max_num_find_queries: - self.find_ids_to_filter = set() # keep all find queries - return - - if not self.retain_positive_queries: - all_find_query_ids = list(range(num_find_queries)) - num_queries_to_filter = max(0, num_find_queries - self.max_num_find_queries) - query_ids_to_filter = random.sample( - all_find_query_ids, k=num_queries_to_filter - ) - else: - # keep up to max_num_find_queries postive find queries and fill - # the remaining slots (if any) with negative find queries - pos_find_ids, neg_find_ids = [], [] - for i, f_q in enumerate(datapoint.find_queries): - # Negative finds return an empty list of object_ids_output - if len(f_q.object_ids_output) == 0: - neg_find_ids.append(i) - else: - pos_find_ids.append(i) - - if len(pos_find_ids) >= self.max_num_find_queries: - # we have more positive find queries than `max_num_find_queries`, - # so we subsample postive find queries and remove all negative find queries - num_queries_to_filter = len(pos_find_ids) - self.max_num_find_queries - query_ids_to_filter = random.sample( - pos_find_ids, k=num_queries_to_filter - ) - query_ids_to_filter.extend(neg_find_ids) - else: - # we have fewer positive find queries than `max_num_find_queries` - # so we need to fill the remaining with negative find queries - num_queries_to_filter = num_find_queries - self.max_num_find_queries - query_ids_to_filter = random.sample( - neg_find_ids, k=num_queries_to_filter - ) - - assert len(query_ids_to_filter) == num_find_queries - self.max_num_find_queries - self.find_ids_to_filter = set(query_ids_to_filter) - - -class KeepMaxNumFindQueriesVideo(FilterDataPointQueries): - def __init__( - self, - video_mosaic_max_num_find_queries_per_frame: int, - retain_positive_queries: bool = False, - ): - self.video_mosaic_max_num_find_queries_per_frame = ( - video_mosaic_max_num_find_queries_per_frame - ) - self.retain_positive_queries = retain_positive_queries - - def identify_queries_to_filter(self, datapoint: Datapoint) -> None: - self.obj_ids_to_filter = set() - num_find_queries = len(datapoint.find_queries) - - findQueries_to_imageIds = defaultdict(list) - max_queries_per_frame = True - for i, f_q in enumerate(datapoint.find_queries): - findQueries_to_imageIds[f_q.image_id].append(i) - if ( - len(findQueries_to_imageIds[f_q.image_id]) - > self.video_mosaic_max_num_find_queries_per_frame - ): - max_queries_per_frame = False - - if max_queries_per_frame: - self.find_ids_to_filter = set() - return - - num_frames = len(findQueries_to_imageIds) - findQueries_0 = findQueries_to_imageIds[0] - num_find_queries_0 = len(findQueries_0) - max_num_find_queries_per_frame = ( - self.video_mosaic_max_num_find_queries_per_frame - ) - if not self.retain_positive_queries: - find_query_ids_0 = list(range(num_find_queries_0)) - num_queries_to_filter = max( - 0, num_find_queries_0 - max_num_find_queries_per_frame - ) - query_ids_to_filter_0 = random.sample( - find_query_ids_0, k=num_queries_to_filter - ) - else: - # keep up to max_num_find_queries postive find queries and fill - # the remaining slots (if any) with negative find queries - pos_find_ids_0, neg_find_ids_0 = [], [] - for i, f_q_id in enumerate(findQueries_0): - f_q = datapoint.find_queries[f_q_id] - # Negative finds return an empty list of object_ids_output - if len(f_q.object_ids_output) == 0: - neg_find_ids_0.append(i) - else: - pos_find_ids_0.append(i) - - if len(pos_find_ids_0) >= max_num_find_queries_per_frame: - # we have more positive find queries than `max_num_find_queries`, - # so we subsample postive find queries and remove all negative find queries - num_queries_to_filter = ( - len(pos_find_ids_0) - max_num_find_queries_per_frame - ) - query_ids_to_filter_0 = random.sample( - pos_find_ids_0, k=num_queries_to_filter - ) - query_ids_to_filter_0.extend(neg_find_ids_0) - else: - # we have fewer positive find queries than `max_num_find_queries` - # so we need to fill the remaining with negative find queries - num_queries_to_filter = ( - num_find_queries_0 - max_num_find_queries_per_frame - ) - query_ids_to_filter_0 = random.sample( - neg_find_ids_0, k=num_queries_to_filter - ) - - # get based on frame 0 all find queries from all the frames with the same indices as in frame 0 - query_ids_to_filter = [] - for i in range(num_frames): - findQueries_i = findQueries_to_imageIds[i] - query_ids_to_filter.extend( - [findQueries_i[j] for j in query_ids_to_filter_0] - ) - - assert ( - len(query_ids_to_filter) - == num_find_queries - - self.video_mosaic_max_num_find_queries_per_frame * num_frames - ) - self.find_ids_to_filter = set(query_ids_to_filter) - - -class KeepSemanticFindQueriesOnly(FilterDataPointQueries): - def identify_queries_to_filter(self, datapoint: Datapoint) -> None: - self.obj_ids_to_filter = set() - self.find_ids_to_filter = { - i for i, q in enumerate(datapoint.find_queries) if q.input_bbox is not None - } # filter (remove) geometric find queries (whose input_bbox is not None) - - # Keep all get queries which don't depend on filtered finds - - -class KeepUnaryFindQueriesOnly(FilterDataPointQueries): - def identify_queries_to_filter(self, datapoint: Datapoint) -> None: - self.obj_ids_to_filter = set() - self.find_ids_to_filter = set() - - # Keep all get queries which don't depend on filtered finds - - -class FilterZeroBoxQueries(FilterDataPointQueries): - """ - Filters all find queries which predict a box with zero area - """ - - @staticmethod - def _is_zero_area_object(obj: Object): - # Check if height or width of bounding box is zero - bbox = obj.bbox # Assume in XYXY format - height = bbox[..., 3].item() - bbox[..., 1].item() - width = bbox[..., 2].item() - bbox[..., 0].item() - - return height == 0 or width == 0 - - def identify_queries_to_filter(self, datapoint): - self.obj_ids_to_filter = set() - - # Find objects with zero area - # Assume only one image per datapoint - image_objects = datapoint.images[0].objects - exclude_objects = { - obj_id - for obj_id, obj in enumerate(image_objects) - if self._is_zero_area_object(obj) - } - - # If a query predicts an object with zero area, drop the whole find query - del_find_ids = [] - for i, f_q in enumerate(datapoint.find_queries): - f_q_objects = set(f_q.object_ids_output) - if len(exclude_objects.intersection(f_q_objects)) > 0: - del_find_ids.append(i) - - self.find_ids_to_filter = set(del_find_ids) - - -class FilterFindQueriesWithTooManyOut(FilterDataPointQueries): - """ - Filters all find queries which have more than a specified number of objects in the output - """ - - def __init__(self, max_num_objects: int): - self.max_num_objects = max_num_objects - - def identify_queries_to_filter(self, datapoint): - self.obj_ids_to_filter = set() - - # If a query predicts more than max_num_objects, drop the whole find query - del_find_ids = [] - for i, f_q in enumerate(datapoint.find_queries): - if len(f_q.object_ids_output) > self.max_num_objects: - del_find_ids.append(i) - - self.find_ids_to_filter = set(del_find_ids) - - -class FilterEmptyTargets(FilterDataPointQueries): - """ - Filters all targets which have zero area - """ - - def identify_queries_to_filter(self, datapoint): - self.obj_ids_to_filter = set() - - for img_id in range(len(datapoint.images)): - for obj_id, obj in enumerate(datapoint.images[img_id].objects): - if obj.area < 1e-6: - self.obj_ids_to_filter.add((img_id, obj_id)) - self.find_ids_to_filter = set() - - -class FilterNonExhaustiveFindQueries(FilterDataPointQueries): - """ - Filters all find queries which are non-exhaustive - """ - - def __init__(self, exhaustivity_type: str): - """ - Args: - exhaustivity_type: Can be "pixel" or "instance": - -pixel: filter queries where the union of all segments covers every pixel belonging to target class - -instance: filter queries where there are non-separable or non annotated instances - Note that instance exhaustivity implies pixel exhaustivity - """ - assert exhaustivity_type in ["pixel", "instance"] - self.exhaustivity_type = exhaustivity_type - - def identify_queries_to_filter(self, datapoint): - self.obj_ids_to_filter = set() - - # If a query predicts more than max_num_objects, drop the whole find query - del_find_ids = [] - for i, f_q in enumerate(datapoint.find_queries): - if self.exhaustivity_type == "instance": - if not f_q.is_exhaustive: - del_find_ids.append(i) - elif self.exhaustivity_type == "pixel": - if f_q.is_pixel_exhaustive is not None and not f_q.is_pixel_exhaustive: - del_find_ids.append(i) - else: - raise RuntimeError( - f"Unknown exhaustivity type {self.exhaustivity_type}" - ) - - self.find_ids_to_filter = set(del_find_ids) - - -class FilterInvalidGeometricQueries(FilterDataPointQueries): - """ - Filters geometric queries whose output got deleted (eg due to cropping) - """ - - def identify_queries_to_filter(self, datapoint): - self.obj_ids_to_filter = set() - - # If a query predicts more than max_num_objects, drop the whole find query - del_find_ids = [] - for i, f_q in enumerate(datapoint.find_queries): - if f_q.input_bbox is not None and f_q.query_text == "geometric": - if len(f_q.object_ids_output) == 0: - del_find_ids.append(i) - self.find_ids_to_filter = set(del_find_ids) - - -class FlexibleFilterFindGetQueries: - def __init__( - self, query_filter: FilterDataPointQueries, enabled: bool = True - ) -> None: - self.query_filter = query_filter - self.enabled = enabled - - def __call__(self, datapoint, **kwargs): - if not self.enabled: - return datapoint - - # Identify all queries to filter - self.query_filter.identify_queries_to_filter(datapoint=datapoint) - - del_find_ids = [] - del_get_ids = [] - for i, f_q in enumerate(datapoint.find_queries): - if self.query_filter._do_filter_query(f_q, i): - datapoint.find_queries[i] = None - del_find_ids.append(i) - - new_find_queries = [] - new_get_queries = [] - - find_old_to_new_map = {} - get_old_to_new_map = {} - - find_counter = 0 - get_counter = 0 - - for i, f_q in enumerate(datapoint.find_queries): - if f_q is not None: - find_old_to_new_map[i] = find_counter - find_counter += 1 - new_find_queries.append(f_q) - - start_with_zero_check = False - for n_f_q in new_find_queries: - if n_f_q.query_processing_order == 0: - start_with_zero_check = True - break - - if len(new_find_queries) == 0: - start_with_zero_check = True - - assert ( - start_with_zero_check - ), "Invalid Find queries, they need to start at query_processing_order = 0" - - datapoint.find_queries = new_find_queries - - if len(datapoint.find_queries) == 0: - print("Warning: No find queries left in datapoint, this is not allowed") - print("Filtering function:", self.query_filter) - print("Datapoint:", datapoint) - raise ValueError - - # The deletion may have removed intermediate steps, so we need to remap to make them contiguous again - all_stages = sorted( - list(set(q.query_processing_order for q in datapoint.find_queries)) - ) - stage_map = {qpo: i for i, qpo in enumerate(all_stages)} - for i in range(len(datapoint.find_queries)): - qpo = datapoint.find_queries[i].query_processing_order - datapoint.find_queries[i].query_processing_order = stage_map[qpo] - - # Final step, clear up objects that are not used anymore - for img_id in range(len(datapoint.images)): - all_objects_ids = set( - i - for find in datapoint.find_queries - for i in find.object_ids_output - if find.image_id == img_id - ) - unused_ids = ( - set(range(len(datapoint.images[img_id].objects))) - all_objects_ids - ) - for tgt_img_id, tgt_obj_id in self.query_filter.obj_ids_to_filter: - if tgt_img_id == img_id: - unused_ids.add(tgt_obj_id) - - if len(unused_ids) > 0: - old_objects = datapoint.images[img_id].objects - object_old_to_new_map = {} - new_objects = [] - for i, o in enumerate(old_objects): - if i not in unused_ids: - object_old_to_new_map[i] = len(new_objects) - new_objects.append(o) - - datapoint.images[img_id].objects = new_objects - - # Remap the outputs of the find queries - affected_find_queries_ids = set() - object_old_to_new_map_per_query = {} - for fid, find in enumerate(datapoint.find_queries): - if find.image_id == img_id: - old_object_ids_output = find.object_ids_output - object_old_to_new_map_per_query[fid] = {} - find.object_ids_output = [] - for oid, old_obj_id in enumerate(old_object_ids_output): - if old_obj_id not in unused_ids: - new_obj_id = object_old_to_new_map[old_obj_id] - find.object_ids_output.append(new_obj_id) - object_old_to_new_map_per_query[fid][oid] = ( - len(find.object_ids_output) - 1 - ) - affected_find_queries_ids.add(fid) - - # finally remove unused images - all_imgs_to_keep = set() - for f_q in datapoint.find_queries: - all_imgs_to_keep.add(f_q.image_id) - - old_img_id_to_new_img_id = {} - new_images = [] - for img_id, img in enumerate(datapoint.images): - if img_id in all_imgs_to_keep: - old_img_id_to_new_img_id[img_id] = len(new_images) - new_images.append(img) - datapoint.images = new_images - - for f_q in datapoint.find_queries: - f_q.image_id = old_img_id_to_new_img_id[f_q.image_id] - - return datapoint - - -class AddPrefixSuffixToFindText: - """ - Add prefix or suffix strings to find query text on the fly. - - If `condition_on_text` is True, the prefix or suffix strings are only added - to those find query text in `condition_text_list` (case-insensitive). - """ - - def __init__( - self, - prefix: Optional[str] = None, - suffix: Optional[str] = None, - condition_on_text: bool = False, - condition_text_list: Optional[List[str]] = None, - enabled: bool = True, - ) -> None: - self.prefix = prefix - self.suffix = suffix - self.condition_on_text = condition_on_text - if self.condition_on_text: - assert condition_text_list is not None - self.condition_text_set = {s.lower().strip() for s in condition_text_list} - self.enabled = enabled - if self.enabled: - logging.info( - f"AddPrefixSuffixToFindText: prefix={prefix}, suffix={suffix}, " - f"condition_on_text={condition_on_text}, condition_text_list={condition_text_list}" - ) - - def __call__(self, datapoint, **kwargs): - if not self.enabled: - return datapoint - - for find in datapoint.find_queries: - if find.query_text == "geometric": - # skip geometric find queries - continue - if ( - self.condition_on_text - and find.query_text.lower().strip() not in self.condition_text_set - ): - # if condition_on_text is True, skip those queries not in condition_text_set - continue - - # add prefix and/or suffix strings to the find query text - if self.prefix is not None: - find.query_text = self.prefix + find.query_text - if self.suffix is not None: - find.query_text = find.query_text + self.suffix - - return datapoint - - -class FilterCrowds(FilterDataPointQueries): - def identify_queries_to_filter(self, datapoint: Datapoint) -> None: - """ - Compute set of query ids to keep, for both find and get queries - """ - self.obj_ids_to_filter = set() - self.find_ids_to_filter = set() - # self.get_ids_to_filter = set() - for img_id, img in enumerate(datapoint.images): - for obj_id, obj in enumerate(img.objects): - if obj.is_crowd: - self.obj_ids_to_filter.add((img_id, obj_id)) - - -class TextQueryToVisual: - """ - Transform a test query to a visual query (with some proba), using any of the output targets as the prompt - """ - - def __init__(self, probability, keep_text_queries=False) -> None: - self.probability = probability - assert 0 <= probability <= 1 - self.keep_text_queries = keep_text_queries - - def __call__(self, datapoint: Datapoint, **kwargs): - for find in datapoint.find_queries: - if find.input_bbox is not None or find.input_points is not None: - # skip geometric find queries - continue - - if len(find.object_ids_output) == 0: - # Can't create a visual query, skip - continue - - if find.query_processing_order > 0: - # Second stage query, can't use - continue - - if random.random() > self.probability: - continue - - selected_vq_id = random.choice(find.object_ids_output) - img_id = find.image_id - - find.input_bbox = datapoint.images[img_id].objects[selected_vq_id].bbox - find.input_bbox_label = torch.ones(1, dtype=torch.bool) - if not self.keep_text_queries: - find.query_text = "visual" - - return datapoint - - -class RemoveInputBoxes: - """ - Remove input boxes from find queries - """ - - def __init__(self) -> None: - pass - - def __call__(self, datapoint: Datapoint, **kwargs): - for find in datapoint.find_queries: - if find.input_bbox is None: - continue - - if find.query_text == "geometric": - print("Warning: removing input box from geometric find query") - - find.input_bbox = None - return datapoint - - -class OverwriteTextQuery: - """ - With some probability, overwrite the text query with a custom text - """ - - def __init__(self, target_text, probability=1.0) -> None: - self.probability = probability - self.target_text = target_text - assert 0 <= probability <= 1 - - def __call__(self, datapoint: Datapoint, **kwargs): - for find in datapoint.find_queries: - if random.random() > self.probability: - continue - - find.query_text = self.target_text - - return datapoint diff --git a/train/transforms/point_sampling.py b/train/transforms/point_sampling.py deleted file mode 100644 index e083fde..0000000 --- a/train/transforms/point_sampling.py +++ /dev/null @@ -1,345 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved - -import cv2 -import numpy as np -import torch -from PIL import Image as PILImage -from pycocotools import mask as mask_util - -from sam3.train.data.sam3_image_dataset import Datapoint -from torchvision.ops import masks_to_boxes - - -def sample_points_from_rle(rle, n_points, mode, box=None, normalize=True): - """ - Sample random points from a mask provided in COCO RLE format. 'mode' - 'mode' is in ["centered", "random_mask", "random_box"] - "centered": points are sampled farthest from the mask edges and each other - "random_mask": points are sampled uniformly from the mask - "random_box": points are sampled uniformly from the annotation's box - 'box' must be provided if 'mode' is "random_box". - If 'normalize' is true, points are in [0,1], relative to mask h,w. - """ - mask = np.ascontiguousarray(mask_util.decode(rle)) - points = sample_points_from_mask(mask, n_points, mode, box) - - if normalize: - h, w = mask.shape - norm = np.array([w, h, 1.0])[None, :] - points = points / norm - - return points - - -def sample_points_from_mask(mask, n_points, mode, box=None): - if mode == "centered": - points = center_positive_sample(mask, n_points) - elif mode == "random_mask": - points = uniform_positive_sample(mask, n_points) - elif mode == "random_box": - assert box is not None, "'random_box' mode requires a provided box." - points = uniform_sample_from_box(mask, box, n_points) - else: - raise ValueError(f"Unknown point sampling mode {mode}.") - return points - - -def uniform_positive_sample(mask, n_points): - """ - Samples positive points uniformly from the mask. Only integer pixel - values are sampled. - """ - # Sampling directly from the uncompressed RLE would be faster but is - # likely unnecessary. - mask_points = np.stack(np.nonzero(mask), axis=0).transpose(1, 0) - assert len(mask_points) > 0, "Can't sample positive points from an empty mask." - selected_idxs = np.random.randint(low=0, high=len(mask_points), size=n_points) - selected_points = mask_points[selected_idxs] - - selected_points = selected_points[:, ::-1] # (y, x) -> (x, y) - labels = np.ones((len(selected_points), 1)) - selected_points = np.concatenate([selected_points, labels], axis=1) - - return selected_points - - -def center_positive_sample(mask, n_points): - """ - Samples points farthest from mask edges (by distance transform) - and subsequent points also farthest from each other. Each new point - sampled is treated as an edge for future points. Edges of the image are - treated as edges of the mask. - """ - - # Pad mask by one pixel on each end to assure distance transform - # avoids edges - padded_mask = np.pad(mask, 1) - - points = [] - for _ in range(n_points): - assert np.max(mask) > 0, "Can't sample positive points from an empty mask." - dist = cv2.distanceTransform(padded_mask, cv2.DIST_L2, 0) - point = np.unravel_index(dist.argmax(), dist.shape) - # Mark selected point as background so next point avoids it - padded_mask[point[0], point[1]] = 0 - points.append(point[::-1]) # (y, x) -> (x, y) - - points = np.stack(points, axis=0) - points = points - 1 # Subtract left/top padding of 1 - labels = np.ones((len(points), 1)) - points = np.concatenate([points, labels], axis=1) - - return points - - -def uniform_sample_from_box(mask, box, n_points): - """ - Sample points uniformly from the provided box. The points' labels - are determined by the provided mask. Does not guarantee a positive - point is sampled. The box is assumed unnormalized in XYXY format. - Points are sampled at integer values. - """ - - # Since lower/right edges are exclusive, ceil can be applied to all edges - int_box = np.ceil(box) - - x = np.random.randint(low=int_box[0], high=int_box[2], size=n_points) - y = np.random.randint(low=int_box[1], high=int_box[3], size=n_points) - labels = mask[y, x] - points = np.stack([x, y, labels], axis=1) - - return points - - -def rescale_box_xyxy(box, factor, imsize=None): - """ - Rescale a box providing in unnormalized XYXY format, fixing the center. - If imsize is provided, clamp to the image. - """ - cx, cy = (box[0] + box[2]) / 2, (box[1] + box[3]) / 2 - w, h = box[2] - box[0], box[3] - box[1] - - new_w, new_h = factor * w, factor * h - - new_x0, new_y0 = cx - new_w / 2, cy - new_h / 2 - new_x1, new_y1 = cx + new_w / 2, cy + new_h / 2 - - if imsize is not None: - new_x0 = max(min(new_x0, imsize[1]), 0) - new_x1 = max(min(new_x1, imsize[1]), 0) - new_y0 = max(min(new_y0, imsize[0]), 0) - new_y1 = max(min(new_y1, imsize[0]), 0) - - return [new_x0, new_y0, new_x1, new_y1] - - -def noise_box(box, im_size, box_noise_std, box_noise_max, min_box_area): - if box_noise_std <= 0.0: - return box - noise = box_noise_std * torch.randn(size=(4,)) - w, h = box[2] - box[0], box[3] - box[1] - scale_factor = torch.tensor([w, h, w, h]) - noise = noise * scale_factor - if box_noise_max is not None: - noise = torch.clamp(noise, -box_noise_max, box_noise_max) - input_box = box + noise - # Clamp to maximum image size - img_clamp = torch.tensor([im_size[1], im_size[0], im_size[1], im_size[0]]) - input_box = torch.maximum(input_box, torch.zeros_like(input_box)) - input_box = torch.minimum(input_box, img_clamp) - if (input_box[2] - input_box[0]) * (input_box[3] - input_box[1]) <= min_box_area: - return box - - return input_box - - -class RandomGeometricInputsAPI: - """ - For geometric queries, replaces the input box or points with a random - one sampled from the GT mask. Segments must be provided for objects - that are targets of geometric queries, and must be binary masks. Existing - point and box queries in the datapoint will be ignored and completely replaced. - Will sample points and boxes in XYXY format in absolute pixel space. - - Geometry queries are currently determined by taking any query whose - query text is a set value. - - Args: - num_points (int or (int, int)): how many points to sample. If a tuple, - sample a random number of points uniformly over the inclusive range. - box_chance (float): fraction of time a box is sampled. A box will replace - one sampled point. - box_noise_std (float): if greater than 0, add noise to the sampled boxes - with this std. Noise is relative to the length of the box side. - box_noise_max (int): if not none, truncate any box noise larger than this - in terms of absolute pixels. - resample_box_from_mask (bool): if True, any sampled box will be determined - by finding the extrema of the provided mask. If False, the bbox provided - in the target object will be used. - point_sample_mode (str): In ["centered", "random_mask", "random_box"], - controlling how points are sampled: - "centered": points are sampled farthest from the mask edges and each other - "random_mask": points are sampled uniformly from the mask - "random_box": points are sampled uniformly from the annotation's box - Note that "centered" may be too slow for on-line generation. - geometric_query_str (str): what string in query_text indicates a - geometry query. - minimum_box_area (float): sampled boxes with area this size or smaller after - noising will use the original box instead. It is the input's responsibility - to avoid original boxes that violate necessary area bounds. - concat_points (bool): if True, any sampled points will be added to existing - ones instead of replacing them. - - """ - - def __init__( - self, - num_points, - box_chance, - box_noise_std=0.0, - box_noise_max=None, - minimum_box_area=0.0, - resample_box_from_mask=False, - point_sample_mode="random_mask", - sample_box_scale_factor=1.0, - geometric_query_str="geometric", - concat_points=False, - ): - self.num_points = num_points - if not isinstance(self.num_points, int): - # Convert from inclusive range to exclusive range expected by torch - self.num_points[1] += 1 - self.num_points = tuple(self.num_points) - self.box_chance = box_chance - self.box_noise_std = box_noise_std - self.box_noise_max = box_noise_max - self.minimum_box_area = minimum_box_area - self.resample_box_from_mask = resample_box_from_mask - self.point_sample_mode = point_sample_mode - assert point_sample_mode in [ - "centered", - "random_mask", - "random_box", - ], "Unknown point sample mode." - self.geometric_query_str = geometric_query_str - self.concat_points = concat_points - self.sample_box_scale_factor = sample_box_scale_factor - - def _sample_num_points_and_if_box(self): - if isinstance(self.num_points, tuple): - n_points = torch.randint( - low=self.num_points[0], high=self.num_points[1], size=(1,) - ).item() - else: - n_points = self.num_points - if self.box_chance > 0.0: - use_box = torch.rand(size=(1,)).item() < self.box_chance - n_points -= int(use_box) # box stands in for one point - else: - use_box = False - return n_points, use_box - - def _get_original_box(self, target_object): - if not self.resample_box_from_mask: - return target_object.bbox - mask = target_object.segment - return masks_to_boxes(mask[None, :, :])[0] - - def _get_target_object(self, datapoint, query): - img = datapoint.images[query.image_id] - targets = query.object_ids_output - assert ( - len(targets) == 1 - ), "Geometric queries only support a single target object." - target_idx = targets[0] - return img.objects[target_idx] - - def __call__(self, datapoint, **kwargs): - for query in datapoint.find_queries: - if query.query_text != self.geometric_query_str: - continue - - target_object = self._get_target_object(datapoint, query) - n_points, use_box = self._sample_num_points_and_if_box() - box = self._get_original_box(target_object) - - mask = target_object.segment - if n_points > 0: - # FIXME: The conversion to numpy and back to reuse code - # is awkward, but this is all in the dataloader worker anyway - # on CPU and so I don't think it should matter. - if self.sample_box_scale_factor != 1.0: - sample_box = rescale_box_xyxy( - box.numpy(), self.sample_box_scale_factor, mask.shape - ) - else: - sample_box = box.numpy() - input_points = sample_points_from_mask( - mask.numpy(), - n_points, - self.point_sample_mode, - sample_box, - ) - input_points = torch.as_tensor(input_points) - input_points = input_points[None, :, :] - if self.concat_points and query.input_points is not None: - input_points = torch.cat([query.input_points, input_points], dim=1) - else: - input_points = query.input_points if self.concat_points else None - - if use_box: - w, h = datapoint.images[query.image_id].size - input_box = noise_box( - box, - (h, w), - box_noise_std=self.box_noise_std, - box_noise_max=self.box_noise_max, - min_box_area=self.minimum_box_area, - ) - input_box = input_box[None, :] - else: - input_box = query.input_bbox if self.concat_points else None - - query.input_points = input_points - query.input_bbox = input_box - - return datapoint - - -class RandomizeInputBbox: - """ - Simplified version of the geometric transform that only deals with input boxes - """ - - def __init__( - self, - box_noise_std=0.0, - box_noise_max=None, - minimum_box_area=0.0, - ): - self.box_noise_std = box_noise_std - self.box_noise_max = box_noise_max - self.minimum_box_area = minimum_box_area - - def __call__(self, datapoint: Datapoint, **kwargs): - for query in datapoint.find_queries: - if query.input_bbox is None: - continue - - img = datapoint.images[query.image_id].data - if isinstance(img, PILImage.Image): - w, h = img.size - else: - assert isinstance(img, torch.Tensor) - h, w = img.shape[-2:] - - for box_id in range(query.input_bbox.shape[0]): - query.input_bbox[box_id, :] = noise_box( - query.input_bbox[box_id, :].view(4), - (h, w), - box_noise_std=self.box_noise_std, - box_noise_max=self.box_noise_max, - min_box_area=self.minimum_box_area, - ).view(1, 4) - - return datapoint diff --git a/train/transforms/segmentation.py b/train/transforms/segmentation.py deleted file mode 100644 index 3e97db0..0000000 --- a/train/transforms/segmentation.py +++ /dev/null @@ -1,157 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved - -import numpy as np -import pycocotools.mask as mask_utils -import torch - -import torchvision.transforms.functional as F -from PIL import Image as PILImage - -from sam3.model.box_ops import masks_to_boxes - -from sam3.train.data.sam3_image_dataset import Datapoint - - -class InstanceToSemantic(object): - """Convert instance segmentation to semantic segmentation.""" - - def __init__(self, delete_instance=True, use_rle=False): - self.delete_instance = delete_instance - self.use_rle = use_rle - - def __call__(self, datapoint: Datapoint, **kwargs): - for fquery in datapoint.find_queries: - h, w = datapoint.images[fquery.image_id].size - - if self.use_rle: - all_segs = [ - datapoint.images[fquery.image_id].objects[obj_id].segment - for obj_id in fquery.object_ids_output - ] - if len(all_segs) > 0: - # we need to double check that all rles are the correct size - # Otherwise cocotools will fail silently to an empty [0,0] mask - for seg in all_segs: - assert seg["size"] == all_segs[0]["size"], ( - "Instance segments have inconsistent sizes. " - f"Found sizes {seg['size']} and {all_segs[0]['size']}" - ) - fquery.semantic_target = mask_utils.merge(all_segs) - else: - # There is no good way to create an empty RLE of the correct size - # We resort to converting an empty box to RLE - fquery.semantic_target = mask_utils.frPyObjects( - np.array([[0, 0, 0, 0]], dtype=np.float64), h, w - )[0] - - else: - # `semantic_target` is uint8 and remains uint8 throughout the transforms - # (it contains binary 0 and 1 values just like `segment` for each object) - fquery.semantic_target = torch.zeros((h, w), dtype=torch.uint8) - for obj_id in fquery.object_ids_output: - segment = datapoint.images[fquery.image_id].objects[obj_id].segment - if segment is not None: - assert ( - isinstance(segment, torch.Tensor) - and segment.dtype == torch.uint8 - ) - fquery.semantic_target |= segment - - if self.delete_instance: - for img in datapoint.images: - for obj in img.objects: - del obj.segment - obj.segment = None - - return datapoint - - -class RecomputeBoxesFromMasks: - """Recompute bounding boxes from masks.""" - - def __call__(self, datapoint: Datapoint, **kwargs): - for img in datapoint.images: - for obj in img.objects: - # Note: if the mask is empty, the bounding box will be undefined - # The empty targets should be subsequently filtered - obj.bbox = masks_to_boxes(obj.segment) - obj.area = obj.segment.sum().item() - - return datapoint - - -class DecodeRle: - """This transform decodes RLEs into binary segments. - Implementing it as a transforms allows lazy loading. Some transforms (eg query filters) - may be deleting masks, so decoding them from the beginning is wasteful. - - This transforms needs to be called before any kind of geometric manipulation - """ - - def __call__(self, datapoint: Datapoint, **kwargs): - imgId2size = {} - warning_shown = False - for imgId, img in enumerate(datapoint.images): - if isinstance(img.data, PILImage.Image): - img_w, img_h = img.data.size - elif isinstance(img.data, torch.Tensor): - img_w, img_h = img.data.shape[-2:] - else: - raise RuntimeError(f"Unexpected image type {type(img.data)}") - - imgId2size[imgId] = (img_h, img_w) - - for obj in img.objects: - if obj.segment is not None and not isinstance( - obj.segment, torch.Tensor - ): - if mask_utils.area(obj.segment) == 0: - print("Warning, empty mask found, approximating from box") - obj.segment = torch.zeros(img_h, img_w, dtype=torch.uint8) - x1, y1, x2, y2 = obj.bbox.int().tolist() - obj.segment[y1 : max(y2, y1 + 1), x1 : max(x1 + 1, x2)] = 1 - else: - obj.segment = mask_utils.decode(obj.segment) - # segment is uint8 and remains uint8 throughout the transforms - obj.segment = torch.tensor(obj.segment).to(torch.uint8) - - if list(obj.segment.shape) != [img_h, img_w]: - # Should not happen often, but adding for security - if not warning_shown: - print( - f"Warning expected instance segmentation size to be {[img_h, img_w]} but found {list(obj.segment.shape)}" - ) - # Printing only once per datapoint to avoid spam - warning_shown = True - - obj.segment = F.resize( - obj.segment[None], (img_h, img_w) - ).squeeze(0) - - assert list(obj.segment.shape) == [img_h, img_w] - - warning_shown = False - for query in datapoint.find_queries: - if query.semantic_target is not None and not isinstance( - query.semantic_target, torch.Tensor - ): - query.semantic_target = mask_utils.decode(query.semantic_target) - # segment is uint8 and remains uint8 throughout the transforms - query.semantic_target = torch.tensor(query.semantic_target).to( - torch.uint8 - ) - if tuple(query.semantic_target.shape) != imgId2size[query.image_id]: - if not warning_shown: - print( - f"Warning expected semantic segmentation size to be {imgId2size[query.image_id]} but found {tuple(query.semantic_target.shape)}" - ) - # Printing only once per datapoint to avoid spam - warning_shown = True - - query.semantic_target = F.resize( - query.semantic_target[None], imgId2size[query.image_id] - ).squeeze(0) - - assert tuple(query.semantic_target.shape) == imgId2size[query.image_id] - - return datapoint diff --git a/train/utils/__init__.py b/train/utils/__init__.py deleted file mode 100644 index 46d37d2..0000000 --- a/train/utils/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved diff --git a/train/utils/checkpoint_utils.py b/train/utils/checkpoint_utils.py deleted file mode 100644 index 7f2736a..0000000 --- a/train/utils/checkpoint_utils.py +++ /dev/null @@ -1,358 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved - - -import contextlib -import fnmatch -import logging -from typing import ( - Any, - Callable, - Dict, - List, - Mapping, - Optional, - Sequence, - Set, - Tuple, - Union, -) - -import numpy as np -import torch -import torch.nn as nn -from iopath.common.file_io import g_pathmgr -from torch.jit._script import RecursiveScriptModule - - -def unix_pattern_to_parameter_names( - constraints: List[str], all_parameter_names: Sequence[str] -) -> Union[None, Set[str]]: - """ - Go through the list of parameter names and select those that match - any of the provided constraints - """ - parameter_names = [] - for param_name in constraints: - matching_parameters = set(fnmatch.filter(all_parameter_names, param_name)) - assert ( - len(matching_parameters) > 0 - ), f"param_names {param_name} don't match any param in the given names." - parameter_names.append(matching_parameters) - return set.union(*parameter_names) - - -def filter_params_matching_unix_pattern( - patterns: List[str], state_dict: Dict[str, torch.Tensor] -) -> Dict[str, torch.Tensor]: - """ - Remove from the state dictionary the parameters matching the provided unix patterns - - Args: - patterns: the list of unix patterns to exclude - state_dict: the dictionary to filter - - Returns: - A new state dictionary - """ - if len(patterns) == 0: - return {} - - all_keys = list(state_dict.keys()) - included_keys = unix_pattern_to_parameter_names(patterns, all_keys) - return {k: state_dict[k] for k in included_keys} - - -def exclude_params_matching_unix_pattern( - patterns: List[str], state_dict: Dict[str, torch.Tensor] -) -> Dict[str, torch.Tensor]: - """ - Remove from the state dictionary the parameters matching the provided unix patterns - - Args: - patterns: the list of unix patterns to exclude - state_dict: the dictionary to filter - - Returns: - A new state dictionary - """ - if len(patterns) == 0: - return state_dict - - all_keys = list(state_dict.keys()) - excluded_keys = unix_pattern_to_parameter_names(patterns, all_keys) - return {k: v for k, v in state_dict.items() if k not in excluded_keys} - - -def _get_state_dict_summary(state_dict: Dict[str, torch.Tensor]): - keys = [] - trace = [] - for k, v in state_dict.items(): - keys.append(k) - trace.append(v.sum().item()) - trace = np.array(trace)[np.argsort(keys)] - return trace - - -def assert_skipped_parameters_are_frozen(model: nn.Module, patterns: List[str]): - """ - Verifies that all the parameters matching the provided patterns - are frozen - this acts as a safeguard when ignoring parameter - when saving checkpoints - if the parameters are in fact trainable - """ - if not patterns: - return - - frozen_state_dict = filter_params_matching_unix_pattern( - patterns=patterns, state_dict=model.state_dict() - ) - non_frozen_keys = { - n - for n, p in model.named_parameters() - if n in frozen_state_dict and p.requires_grad - } - if non_frozen_keys: - raise ValueError( - f"Parameters excluded with `skip_saving_parameters` should be frozen: {non_frozen_keys}" - ) - - -@contextlib.contextmanager -def with_check_parameter_frozen( - model: nn.Module, patterns: List[str], disabled: bool = True -): - """ - Context manager that inspects a model surrounding a piece of code - and verifies if the model has been updated by this piece of code - - The function will raise an exception if the model has been updated - on at least one of the parameter that matches one of the pattern - - Args: - model: the model that might have been updated - patterns: for the parameters we want to observe - allowed: - """ - if not patterns or disabled: - yield - return - - frozen_state_dict = filter_params_matching_unix_pattern( - patterns=patterns, state_dict=model.state_dict() - ) - summary_before = _get_state_dict_summary(frozen_state_dict) - - yield - - frozen_state_dict = filter_params_matching_unix_pattern( - patterns=patterns, state_dict=model.state_dict() - ) - summary_after = _get_state_dict_summary(frozen_state_dict) - - if not np.allclose(summary_before, summary_after, atol=1e-6): - raise ValueError( - f""" - The `model_weight_initializer` has initialized parameters frozen with `skip_saving_parameters`. - You can resolve this error by either initializing those parameters from within the model definition - or using the flag `trainer.checkpoint.initialize_after_preemption` to True. - """ - ) - - -class CkptExcludeKernel: - """ - Removes the keys from the given model state_dict that match the key_pattern. - - Args: - key_pattern: Patterns used to select the keys in the state_dict - that are eligible for this kernel. - """ - - def __init__(self, key_pattern: List[str]): - self.key_pattern = key_pattern - - def __call__(self, state_dict: Dict): - """ - Args: - state_dict: A dictionary representing the given checkpoint's state dict. - """ - if len(self.key_pattern) == 0: - return state_dict - exclude_keys = unix_pattern_to_parameter_names( - self.key_pattern, state_dict.keys() - ) - return {k: v for k, v in state_dict.items() if k not in exclude_keys} - - -def load_checkpoint( - path_list: List[str], - pick_recursive_keys: Optional[List[str]] = None, - map_location: str = "cpu", -) -> Any: - """ - Loads a checkpoint from the specified path. - - Args: - path_list: A list of paths which contain the checkpoint. Each element - is tried (in order) until a file that exists is found. That file is then - used to read the checkpoint. - pick_recursive_keys: Picks sub dicts from the loaded checkpoint if not None. - For pick_recursive_keys = ["a", "b"], will return checkpoint_dict["a"]["b"] - map_location (str): a function, torch.device, string or a dict specifying how to - remap storage locations - - Returns: Model with the matchin pre-trained weights loaded. - """ - path_exists = False - for path in path_list: - if g_pathmgr.exists(path): - path_exists = True - break - - if not path_exists: - raise ValueError(f"No path exists in {path_list}") - - with g_pathmgr.open(path, "rb") as f: - checkpoint = torch.load(f, map_location=map_location) - - logging.info(f"Loaded checkpoint from {path}") - if pick_recursive_keys is not None: - for key in pick_recursive_keys: - checkpoint = checkpoint[key] - return checkpoint - - -def get_state_dict(checkpoint, ckpt_state_dict_keys): - if isinstance(checkpoint, RecursiveScriptModule): - # This is a torchscript JIT model - return checkpoint.state_dict() - pre_train_dict = checkpoint - for i, key in enumerate(ckpt_state_dict_keys): - if (isinstance(pre_train_dict, Mapping) and key not in pre_train_dict) or ( - isinstance(pre_train_dict, Sequence) and key >= len(pre_train_dict) - ): - key_str = ( - '["' + '"]["'.join(list(map(ckpt_state_dict_keys[:i], str))) + '"]' - ) - raise KeyError( - f"'{key}' not found in checkpoint{key_str} " - f"with keys: {pre_train_dict.keys()}" - ) - pre_train_dict = pre_train_dict[key] - return pre_train_dict - - -def load_checkpoint_and_apply_kernels( - checkpoint_path: str, - checkpoint_kernels: List[Callable] = None, - ckpt_state_dict_keys: Tuple[str] = ("state_dict",), - map_location: str = "cpu", -) -> nn.Module: - """ - Performs checkpoint loading with a variety of pre-processing kernel applied in - sequence. - - Args: - checkpoint_path (str): Path to the checkpoint. - checkpoint_kernels List(Callable): A list of checkpoint processing kernels - to apply in the specified order. Supported kernels include `CkptIncludeKernel`, - `CkptExcludeKernel`, etc. These kernels are applied in the - given order. - ckpt_state_dict_keys (str): Keys containing the model state dict. - map_location (str): a function, torch.device, string or a dict specifying how to - remap storage locations - - Returns: Model with the matchin pre-trained weights loaded. - """ - assert g_pathmgr.exists(checkpoint_path), "Checkpoint '{}' not found".format( - checkpoint_path - ) - - # Load the checkpoint on CPU to avoid GPU mem spike. - with g_pathmgr.open(checkpoint_path, "rb") as f: - checkpoint = torch.load(f, map_location=map_location) - - pre_train_dict = get_state_dict(checkpoint, ckpt_state_dict_keys) - - # Not logging into info etc since it's a huge log - logging.debug( - "Loaded Checkpoint State Dict pre-kernel application: %s" - % str(", ".join(list(pre_train_dict.keys()))) - ) - # Apply kernels - if checkpoint_kernels is not None: - for f in checkpoint_kernels: - pre_train_dict = f(state_dict=pre_train_dict) - - logging.debug( - "Loaded Checkpoint State Dict Post-kernel application %s" - % str(", ".join(list(pre_train_dict.keys()))) - ) - - return pre_train_dict - - -def check_load_state_dict_errors( - missing_keys, - unexpected_keys, - strict: bool, - ignore_missing_keys: List[str] = None, - ignore_unexpected_keys: List[str] = None, -): - if ignore_missing_keys is not None and len(ignore_missing_keys) > 0: - ignored_keys = unix_pattern_to_parameter_names( - ignore_missing_keys, missing_keys - ) - missing_keys = [key for key in missing_keys if key not in ignored_keys] - - if ignore_unexpected_keys is not None and len(ignore_unexpected_keys) > 0: - ignored_unexpected_keys = unix_pattern_to_parameter_names( - ignore_unexpected_keys, unexpected_keys - ) - unexpected_keys = [ - key for key in unexpected_keys if key not in ignored_unexpected_keys - ] - - err = "State key mismatch." - if unexpected_keys: - err += f" Unexpected keys: {unexpected_keys}." - if missing_keys: - err += f" Missing keys: {missing_keys}." - - if unexpected_keys or missing_keys: - logging.warning(err) - if unexpected_keys or strict: - raise KeyError(err) - - -def load_state_dict_into_model( - state_dict: Dict, - model: nn.Module, - strict: bool = True, - ignore_missing_keys: List[str] = None, - ignore_unexpected_keys: List[str] = None, - checkpoint_kernels: List[Callable] = None, -): - """ - Loads a state dict into the given model. - - Args: - state_dict: A dictionary containing the model's - state dict, or a subset if strict is False - model: Model to load the checkpoint weights into - strict: raise if the state_dict has missing state keys - ignore_missing_keys: unix pattern of keys to ignore - """ - # Apply kernels - if checkpoint_kernels is not None: - for f in checkpoint_kernels: - state_dict = f(state_dict=state_dict) - missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False) - - check_load_state_dict_errors( - missing_keys, - unexpected_keys, - strict=strict, - ignore_missing_keys=ignore_missing_keys, - ignore_unexpected_keys=ignore_unexpected_keys, - ) - return model diff --git a/train/utils/distributed.py b/train/utils/distributed.py deleted file mode 100644 index 3c87a91..0000000 --- a/train/utils/distributed.py +++ /dev/null @@ -1,585 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved - -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. - -# This source code is licensed under the license found in the -# LICENSE file in the root directory of this source tree. - -import datetime -import functools -import io -import logging -import os -import random -import tempfile -import time -from typing import Any, Callable, List, Tuple - -import torch -import torch.autograd as autograd -import torch.distributed as dist - - -# Default to GPU 0 -_cuda_device_index: int = 0 - -# Setting _cuda_device_index to -1 internally implies that we should use CPU -_CPU_DEVICE_INDEX = -1 -_PRIMARY_RANK = 0 - - -@functools.lru_cache() -def _get_global_gloo_group(): - """ - Return a process group based on gloo backend, containing all the ranks - The result is cached. - """ - - if dist.get_backend() == "nccl": - # Increase timeout from 1800 sec to 43200 sec (12 hr) to avoid some processes - # being much slower than others causing a timeout (which can happen in relation - # or LVIS class mAP evaluation). - timeout = 43200 - return dist.new_group( - backend="gloo", - timeout=datetime.timedelta(seconds=timeout), - ) - - return dist.group.WORLD - - -def is_main_process(): - """Return true if the current process is the main one""" - return get_rank() == 0 - - -def all_gather_via_filesys(data, filesys_save_dir=None, gather_to_rank_0_only=False): - """ - Run all_gather on arbitrary picklable data (not necessarily tensors), similar to - `all_gather` above, but using filesystem instead of collective ops. - - If gather_to_rank_0_only is True, only rank 0 will load the gathered object list - (and other ranks will have an empty list). - """ - world_size = get_world_size() - if world_size == 1: - return [data] - - print("gathering via files") - cpu_group = _get_global_gloo_group() - - # if unspecified, we will save to the current python file dir - if filesys_save_dir is not None: - save_dir = filesys_save_dir - elif "EXP_DIR" in os.environ: - save_dir = os.environ["EXP_DIR"] - else: - # try the same directory where the code is stored - save_dir = filesys_save_dir or os.path.dirname(__file__) - save_dir = os.path.join(save_dir, "all_gather_via_filesys") - if is_main_process(): - os.makedirs(save_dir, exist_ok=True) - - # use a timestamp and salt to distinguish different all_gather - timestamp = int(time.time()) if is_main_process() else 0 - salt = random.randint(0, 2**31 - 1) if is_main_process() else 0 - # broadcast the timestamp and salt across ranks - # (all-reduce will do the broadcasting since only rank 0 is non-zero) - timestamp_and_salt = torch.tensor([timestamp, salt], dtype=torch.long) - dist.all_reduce(timestamp_and_salt, group=cpu_group) - timestamp, salt = timestamp_and_salt.tolist() - - # save the data to a file on the disk - rank_save = get_rank() - save_data_filename = f"data_to_gather_{timestamp}_{salt}_{rank_save}.pkl" - save_data_path = os.path.join(save_dir, save_data_filename) - assert not os.path.exists(save_data_path), f"{save_data_path} already exists" - torch.save(data, save_data_path) - dist.barrier(group=cpu_group) - - # read the data from the files - data_list = [] - if rank_save == 0 or not gather_to_rank_0_only: - for rank_load in range(world_size): - load_data_filename = f"data_to_gather_{timestamp}_{salt}_{rank_load}.pkl" - load_data_path = os.path.join(save_dir, load_data_filename) - assert os.path.exists(load_data_path), f"cannot read {save_data_path}" - data_list.append(torch.load(load_data_path, weights_only=False)) - dist.barrier(group=cpu_group) - - # delete the saved file - os.remove(save_data_path) - return data_list - - -def all_gather(data, force_cpu=False, force_filesys=False, filesys_save_dir=None): - """ - Run all_gather on arbitrary picklable data (not necessarily tensors) - Args: - data: any picklable object - Returns: - list[data]: list of data gathered from each rank - """ - - world_size = get_world_size() - if world_size == 1: - return [data] - - if os.getenv("MDETR_FILESYS_REDUCE_RANK_0_ONLY") == "1": - return all_gather_via_filesys( - data, filesys_save_dir, gather_to_rank_0_only=True - ) - - if os.getenv("MDETR_FILESYS_REDUCE") == "1" or force_filesys: - return all_gather_via_filesys(data, filesys_save_dir) - - cpu_group = None - if os.getenv("MDETR_CPU_REDUCE") == "1" or force_cpu: - cpu_group = _get_global_gloo_group() - - buffer = io.BytesIO() - torch.save(data, buffer) - data_view = buffer.getbuffer() - device = "cuda" if cpu_group is None else "cpu" - tensor = torch.ByteTensor(data_view).to(device) - - # obtain Tensor size of each rank - local_size = torch.tensor([tensor.numel()], device=device, dtype=torch.long) - size_list = [ - torch.tensor([0], device=device, dtype=torch.long) for _ in range(world_size) - ] - if cpu_group is None: - dist.all_gather(size_list, local_size) - else: - print("gathering on cpu") - dist.all_gather(size_list, local_size, group=cpu_group) - size_list = [int(size.item()) for size in size_list] - max_size = max(size_list) - assert isinstance(local_size.item(), int) - local_size = int(local_size.item()) - - # receiving Tensor from all ranks - # we pad the tensor because torch all_gather does not support - # gathering tensors of different shapes - tensor_list = [] - for _ in size_list: - tensor_list.append(torch.empty((max_size,), dtype=torch.uint8, device=device)) - if local_size != max_size: - padding = torch.empty( - size=(max_size - local_size,), dtype=torch.uint8, device=device - ) - tensor = torch.cat((tensor, padding), dim=0) - if cpu_group is None: - dist.all_gather(tensor_list, tensor) - else: - dist.all_gather(tensor_list, tensor, group=cpu_group) - - data_list = [] - for size, tensor in zip(size_list, tensor_list): - tensor = torch.split(tensor, [size, max_size - size], dim=0)[0] - buffer = io.BytesIO(tensor.cpu().numpy()) - obj = torch.load(buffer, weights_only=False) - data_list.append(obj) - - return data_list - - -def convert_to_distributed_tensor(tensor: torch.Tensor) -> Tuple[torch.Tensor, str]: - """ - For some backends, such as NCCL, communication only works if the - tensor is on the GPU. This helper function converts to the correct - device and returns the tensor + original device. - """ - orig_device = "cpu" if not tensor.is_cuda else "gpu" - if ( - torch.distributed.is_available() - and torch.distributed.get_backend() == torch.distributed.Backend.NCCL - and not tensor.is_cuda - ): - tensor = tensor.cuda() - return (tensor, orig_device) - - -def convert_to_normal_tensor(tensor: torch.Tensor, orig_device: str) -> torch.Tensor: - """ - For some backends, such as NCCL, communication only works if the - tensor is on the GPU. This converts the tensor back to original device. - """ - if tensor.is_cuda and orig_device == "cpu": - tensor = tensor.cpu() - return tensor - - -def is_distributed_training_run() -> bool: - return ( - torch.distributed.is_available() - and torch.distributed.is_initialized() - and (torch.distributed.get_world_size() > 1) - ) - - -def is_primary() -> bool: - """ - Returns True if this is rank 0 of a distributed training job OR if it is - a single trainer job. Otherwise False. - """ - return get_rank() == _PRIMARY_RANK - - -def all_reduce_mean(tensor: torch.Tensor) -> torch.Tensor: - """ - Wrapper over torch.distributed.all_reduce for performing mean reduction - of tensor over all processes. - """ - return all_reduce_op( - tensor, - torch.distributed.ReduceOp.SUM, - lambda t: t / torch.distributed.get_world_size(), - ) - - -def all_reduce_sum(tensor: torch.Tensor) -> torch.Tensor: - """ - Wrapper over torch.distributed.all_reduce for performing sum - reduction of tensor over all processes in both distributed / - non-distributed scenarios. - """ - return all_reduce_op(tensor, torch.distributed.ReduceOp.SUM) - - -def all_reduce_min(tensor: torch.Tensor) -> torch.Tensor: - """ - Wrapper over torch.distributed.all_reduce for performing min - reduction of tensor over all processes in both distributed / - non-distributed scenarios. - """ - return all_reduce_op(tensor, torch.distributed.ReduceOp.MIN) - - -def all_reduce_max(tensor: torch.Tensor) -> torch.Tensor: - """ - Wrapper over torch.distributed.all_reduce for performing min - reduction of tensor over all processes in both distributed / - non-distributed scenarios. - """ - return all_reduce_op(tensor, torch.distributed.ReduceOp.MAX) - - -def all_reduce_op( - tensor: torch.Tensor, - op: torch.distributed.ReduceOp, - after_op_func: Callable[[torch.Tensor], torch.Tensor] = None, -) -> torch.Tensor: - """ - Wrapper over torch.distributed.all_reduce for performing - reduction of tensor over all processes in both distributed / - non-distributed scenarios. - """ - if is_distributed_training_run(): - tensor, orig_device = convert_to_distributed_tensor(tensor) - torch.distributed.all_reduce(tensor, op) - if after_op_func is not None: - tensor = after_op_func(tensor) - tensor = convert_to_normal_tensor(tensor, orig_device) - return tensor - - -def gather_tensors_from_all(tensor: torch.Tensor) -> List[torch.Tensor]: - """ - Wrapper over torch.distributed.all_gather for performing - 'gather' of 'tensor' over all processes in both distributed / - non-distributed scenarios. - """ - if tensor.ndim == 0: - # 0 dim tensors cannot be gathered. so unsqueeze - tensor = tensor.unsqueeze(0) - - if is_distributed_training_run(): - tensor, orig_device = convert_to_distributed_tensor(tensor) - gathered_tensors = [ - torch.zeros_like(tensor) for _ in range(torch.distributed.get_world_size()) - ] - torch.distributed.all_gather(gathered_tensors, tensor) - gathered_tensors = [ - convert_to_normal_tensor(_tensor, orig_device) - for _tensor in gathered_tensors - ] - else: - gathered_tensors = [tensor] - - return gathered_tensors - - -def gather_from_all(tensor: torch.Tensor) -> torch.Tensor: - gathered_tensors = gather_tensors_from_all(tensor) - gathered_tensor = torch.cat(gathered_tensors, 0) - return gathered_tensor - - -def broadcast(tensor: torch.Tensor, src: int = 0) -> torch.Tensor: - """ - Wrapper over torch.distributed.broadcast for broadcasting a tensor from the source - to all processes in both distributed / non-distributed scenarios. - """ - if is_distributed_training_run(): - tensor, orig_device = convert_to_distributed_tensor(tensor) - torch.distributed.broadcast(tensor, src) - tensor = convert_to_normal_tensor(tensor, orig_device) - return tensor - - -def barrier() -> None: - """ - Wrapper over torch.distributed.barrier, returns without waiting - if the distributed process group is not initialized instead of throwing error. - """ - if not torch.distributed.is_available() or not torch.distributed.is_initialized(): - return - torch.distributed.barrier() - - -def get_world_size() -> int: - """ - Simple wrapper for correctly getting worldsize in both distributed - / non-distributed settings - """ - return ( - torch.distributed.get_world_size() - if torch.distributed.is_available() and torch.distributed.is_initialized() - else 1 - ) - - -def get_rank() -> int: - """ - Simple wrapper for correctly getting rank in both distributed - / non-distributed settings - """ - return ( - torch.distributed.get_rank() - if torch.distributed.is_available() and torch.distributed.is_initialized() - else 0 - ) - - -def get_primary_rank() -> int: - return _PRIMARY_RANK - - -def set_cuda_device_index(idx: int) -> None: - global _cuda_device_index - _cuda_device_index = idx - torch.cuda.set_device(_cuda_device_index) - - -def set_cpu_device() -> None: - global _cuda_device_index - _cuda_device_index = _CPU_DEVICE_INDEX - - -def get_cuda_device_index() -> int: - return _cuda_device_index - - -def init_distributed_data_parallel_model( - model: torch.nn.Module, - broadcast_buffers: bool = False, - find_unused_parameters: bool = True, - bucket_cap_mb: int = 25, -) -> torch.nn.parallel.DistributedDataParallel: - global _cuda_device_index - - if _cuda_device_index == _CPU_DEVICE_INDEX: - # CPU-only model, don't specify device - return torch.nn.parallel.DistributedDataParallel( - model, - broadcast_buffers=broadcast_buffers, - find_unused_parameters=find_unused_parameters, - bucket_cap_mb=bucket_cap_mb, - ) - else: - # GPU model - return torch.nn.parallel.DistributedDataParallel( - model, - device_ids=[_cuda_device_index], - output_device=_cuda_device_index, - broadcast_buffers=broadcast_buffers, - find_unused_parameters=find_unused_parameters, - bucket_cap_mb=bucket_cap_mb, - ) - - -def broadcast_object(obj: Any, src: int = _PRIMARY_RANK, use_disk: bool = True) -> Any: - """Broadcast an object from a source to all workers. - - Args: - obj: Object to broadcast, must be serializable - src: Source rank for broadcast (default is primary) - use_disk: If enabled, removes redundant CPU memory copies by writing to - disk - """ - # Either broadcast from primary to the fleet (default), - # or use the src setting as the original rank - if get_rank() == src: - # Emit data - buffer = io.BytesIO() - torch.save(obj, buffer) - data_view = buffer.getbuffer() - length_tensor = torch.LongTensor([len(data_view)]) - length_tensor = broadcast(length_tensor, src=src) - data_tensor = torch.ByteTensor(data_view) - data_tensor = broadcast(data_tensor, src=src) - else: - # Fetch from the source - length_tensor = torch.LongTensor([0]) - length_tensor = broadcast(length_tensor, src=src) - data_tensor = torch.empty([length_tensor.item()], dtype=torch.uint8) - data_tensor = broadcast(data_tensor, src=src) - if use_disk: - with tempfile.TemporaryFile("r+b") as f: - f.write(data_tensor.numpy()) - # remove reference to the data tensor and hope that Python garbage - # collects it - del data_tensor - f.seek(0) - obj = torch.load(f, weights_only=False) - else: - buffer = io.BytesIO(data_tensor.numpy()) - obj = torch.load(buffer, weights_only=False) - return obj - - -def all_gather_tensor(tensor: torch.Tensor, world_size=None): - if world_size is None: - world_size = get_world_size() - # make contiguous because NCCL won't gather the tensor otherwise - assert tensor.is_contiguous(), f"{tensor.shape} is not contiguous!" - tensor, orig_device = convert_to_distributed_tensor(tensor) - tensor_all = [torch.ones_like(tensor) for _ in range(world_size)] - dist.all_gather(tensor_all, tensor, async_op=False) # performance opt - tensor_all = [ - convert_to_normal_tensor(tensor, orig_device) for tensor in tensor_all - ] - return tensor_all - - -def all_gather_batch(tensors: List[torch.Tensor]): - """ - Performs all_gather operation on the provided tensors. - """ - # Queue the gathered tensors - world_size = get_world_size() - # There is no need for reduction in the single-proc case - if world_size == 1: - return tensors - tensor_list = [] - output_tensor = [] - for tensor in tensors: - tensor_all = all_gather_tensor(tensor, world_size) - tensor_list.append(tensor_all) - - for tensor_all in tensor_list: - output_tensor.append(torch.cat(tensor_all, dim=0)) - return output_tensor - - -class GatherLayer(autograd.Function): - """ - Gather tensors from all workers with support for backward propagation: - This implementation does not cut the gradients as torch.distributed.all_gather does. - """ - - @staticmethod - def forward(ctx, x): - output = [torch.zeros_like(x) for _ in range(dist.get_world_size())] - dist.all_gather(output, x) - return tuple(output) - - @staticmethod - def backward(ctx, *grads): - all_gradients = torch.stack(grads) - dist.all_reduce(all_gradients) - return all_gradients[dist.get_rank()] - - -def all_gather_batch_with_grad(tensors): - """ - Performs all_gather operation on the provided tensors. - Graph remains connected for backward grad computation. - """ - # Queue the gathered tensors - world_size = get_world_size() - # There is no need for reduction in the single-proc case - if world_size == 1: - return tensors - tensor_list = [] - output_tensor = [] - - for tensor in tensors: - tensor_all = GatherLayer.apply(tensor) - tensor_list.append(tensor_all) - - for tensor_all in tensor_list: - output_tensor.append(torch.cat(tensor_all, dim=0)) - return output_tensor - - -def unwrap_ddp_if_wrapped(model): - if isinstance(model, torch.nn.parallel.DistributedDataParallel): - return model.module - return model - - -def create_new_process_group(group_size): - """ - Creates process groups of a gives `group_size` and returns - process group that current GPU participates in. - - `group_size` must divide the total number of GPUs (world_size). - - Modified from - https://github.com/NVIDIA/apex/blob/4e1ae43f7f7ac69113ef426dd15f37123f0a2ed3/apex/parallel/__init__.py#L60 - - Args: - group_size (int): number of GPU's to collaborate for sync bn - """ - - assert group_size > 0 - - world_size = torch.distributed.get_world_size() - if world_size <= 8: - if group_size > world_size: - logging.warning( - f"Requested group size [{group_size}] > world size [{world_size}]. " - "Assuming local debug run and capping it to world size." - ) - group_size = world_size - assert world_size >= group_size - assert world_size % group_size == 0 - - group = None - for group_num in range(world_size // group_size): - group_ids = range(group_num * group_size, (group_num + 1) * group_size) - cur_group = torch.distributed.new_group(ranks=group_ids) - if torch.distributed.get_rank() // group_size == group_num: - group = cur_group - # can not drop out and return here, every process must go through creation of all subgroups - - assert group is not None - return group - - -def is_dist_avail_and_initialized(): - if not dist.is_available(): - return False - if not dist.is_initialized(): - return False - return True - - -def gather_to_rank_0_via_filesys(data, filesys_save_dir=None): - """ - Gather any picklable data to rank 0 via filesystem, using all_gather_via_filesys. - """ - return all_gather_via_filesys(data, filesys_save_dir, gather_to_rank_0_only=True) diff --git a/train/utils/logger.py b/train/utils/logger.py deleted file mode 100644 index 127f6c8..0000000 --- a/train/utils/logger.py +++ /dev/null @@ -1,241 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved - -import atexit -import functools -import logging -import sys -import uuid -from typing import Any, Dict, Optional, Union - -from hydra.utils import instantiate - -from iopath.common.file_io import g_pathmgr -from numpy import ndarray - -from sam3.train.utils.train_utils import get_machine_local_and_dist_rank, makedir -from torch import Tensor -from torch.utils.tensorboard import SummaryWriter - -Scalar = Union[Tensor, ndarray, int, float] - - -def make_tensorboard_logger(log_dir: str, **writer_kwargs: Any): - makedir(log_dir) - summary_writer_method = SummaryWriter - return TensorBoardLogger( - path=log_dir, summary_writer_method=summary_writer_method, **writer_kwargs - ) - - -class TensorBoardWriterWrapper: - """ - A wrapper around a SummaryWriter object. - """ - - def __init__( - self, - path: str, - *args: Any, - filename_suffix: str = None, - summary_writer_method: Any = SummaryWriter, - **kwargs: Any, - ) -> None: - """Create a new TensorBoard logger. - On construction, the logger creates a new events file that logs - will be written to. If the environment variable `RANK` is defined, - logger will only log if RANK = 0. - - NOTE: If using the logger with distributed training: - - This logger can call collective operations - - Logs will be written on rank 0 only - - Logger must be constructed synchronously *after* initializing distributed process group. - - Args: - path (str): path to write logs to - *args, **kwargs: Extra arguments to pass to SummaryWriter - """ - self._writer: Optional[SummaryWriter] = None - _, self._rank = get_machine_local_and_dist_rank() - self._path: str = path - if self._rank == 0: - logging.info( - f"TensorBoard SummaryWriter instantiated. Files will be stored in: {path}" - ) - self._writer = summary_writer_method( - log_dir=path, - *args, - filename_suffix=filename_suffix or str(uuid.uuid4()), - **kwargs, - ) - else: - logging.debug( - f"Not logging meters on this host because env RANK: {self._rank} != 0" - ) - atexit.register(self.close) - - @property - def writer(self) -> Optional[SummaryWriter]: - return self._writer - - @property - def path(self) -> str: - return self._path - - def flush(self) -> None: - """Writes pending logs to disk.""" - - if not self._writer: - return - - self._writer.flush() - - def close(self) -> None: - """Close writer, flushing pending logs to disk. - Logs cannot be written after `close` is called. - """ - - if not self._writer: - return - - self._writer.close() - self._writer = None - - -class TensorBoardLogger(TensorBoardWriterWrapper): - """ - A simple logger for TensorBoard. - """ - - def log_dict(self, payload: Dict[str, Scalar], step: int) -> None: - """Add multiple scalar values to TensorBoard. - - Args: - payload (dict): dictionary of tag name and scalar value - step (int, Optional): step value to record - """ - if not self._writer: - return - for k, v in payload.items(): - self.log(k, v, step) - - def log(self, name: str, data: Scalar, step: int) -> None: - """Add scalar data to TensorBoard. - - Args: - name (string): tag name used to group scalars - data (float/int/Tensor): scalar data to log - step (int, optional): step value to record - """ - if not self._writer: - return - self._writer.add_scalar(name, data, global_step=step, new_style=True) - - def log_hparams( - self, hparams: Dict[str, Scalar], meters: Dict[str, Scalar] - ) -> None: - """Add hyperparameter data to TensorBoard. - - Args: - hparams (dict): dictionary of hyperparameter names and corresponding values - meters (dict): dictionary of name of meter and corersponding values - """ - if not self._writer: - return - self._writer.add_hparams(hparams, meters) - - -class Logger: - """ - A logger class that can interface with multiple loggers. It now supports tensorboard only for simplicity, but you can extend it with your own logger. - """ - - def __init__(self, logging_conf): - # allow turning off TensorBoard with "should_log: false" in config - tb_config = logging_conf.tensorboard_writer - tb_should_log = tb_config and tb_config.pop("should_log", True) - self.tb_logger = instantiate(tb_config) if tb_should_log else None - - def log_dict(self, payload: Dict[str, Scalar], step: int) -> None: - if self.tb_logger: - self.tb_logger.log_dict(payload, step) - - def log(self, name: str, data: Scalar, step: int) -> None: - if self.tb_logger: - self.tb_logger.log(name, data, step) - - def log_hparams( - self, hparams: Dict[str, Scalar], meters: Dict[str, Scalar] - ) -> None: - if self.tb_logger: - self.tb_logger.log_hparams(hparams, meters) - - -# cache the opened file object, so that different calls to `setup_logger` -# with the same file name can safely write to the same file. -@functools.lru_cache(maxsize=None) -def _cached_log_stream(filename): - # we tune the buffering value so that the logs are updated - # frequently. - log_buffer_kb = 10 * 1024 # 10KB - io = g_pathmgr.open(filename, mode="a", buffering=log_buffer_kb) - atexit.register(io.close) - return io - - -def setup_logging( - name, - output_dir=None, - rank=0, - log_level_primary="INFO", - log_level_secondary="ERROR", -): - """ - Setup various logging streams: stdout and file handlers. - For file handlers, we only setup for the master gpu. - """ - # get the filename if we want to log to the file as well - log_filename = None - if output_dir: - makedir(output_dir) - if rank == 0: - log_filename = f"{output_dir}/log.txt" - - logger = logging.getLogger(name) - logger.setLevel(log_level_primary) - - # create formatter - FORMAT = "%(levelname)s %(asctime)s %(filename)s:%(lineno)4d: %(message)s" - formatter = logging.Formatter(FORMAT) - - # Cleanup any existing handlers - for h in logger.handlers: - logger.removeHandler(h) - logger.root.handlers = [] - - # setup the console handler - console_handler = logging.StreamHandler(sys.stdout) - console_handler.setFormatter(formatter) - logger.addHandler(console_handler) - if rank == 0: - console_handler.setLevel(log_level_primary) - else: - console_handler.setLevel(log_level_secondary) - - # we log to file as well if user wants - if log_filename and rank == 0: - file_handler = logging.StreamHandler(_cached_log_stream(log_filename)) - file_handler.setLevel(log_level_primary) - file_handler.setFormatter(formatter) - logger.addHandler(file_handler) - - logging.root = logger - - -def shutdown_logging(): - """ - After training is done, we ensure to shut down all the logger streams. - """ - logging.info("Shutting down loggers...") - handlers = logging.root.handlers - for handler in handlers: - handler.close() diff --git a/train/utils/train_utils.py b/train/utils/train_utils.py deleted file mode 100644 index 43a8078..0000000 --- a/train/utils/train_utils.py +++ /dev/null @@ -1,285 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved - -import logging -import math -import os -import random -import re -from datetime import timedelta -from typing import Optional - -import hydra - -import numpy as np -import omegaconf -import torch -import torch.distributed as dist -from iopath.common.file_io import g_pathmgr -from omegaconf import OmegaConf - - -def multiply_all(*args): - return np.prod(np.array(args)).item() - - -def collect_dict_keys(config): - """This function recursively iterates through a dataset configuration, and collect all the dict_key that are defined""" - val_keys = [] - # If the this config points to the collate function, then it has a key - if "_target_" in config and re.match(r".*collate_fn.*", config["_target_"]): - val_keys.append(config["dict_key"]) - else: - # Recursively proceed - for v in config.values(): - if isinstance(v, type(config)): - val_keys.extend(collect_dict_keys(v)) - elif isinstance(v, omegaconf.listconfig.ListConfig): - for item in v: - if isinstance(item, type(config)): - val_keys.extend(collect_dict_keys(item)) - return val_keys - - -class Phase: - TRAIN = "train" - VAL = "val" - - -def register_omegaconf_resolvers(): - OmegaConf.register_new_resolver("get_method", hydra.utils.get_method) - OmegaConf.register_new_resolver("get_class", hydra.utils.get_class) - OmegaConf.register_new_resolver("add", lambda x, y: x + y) - OmegaConf.register_new_resolver("times", multiply_all) - OmegaConf.register_new_resolver("divide", lambda x, y: x / y) - OmegaConf.register_new_resolver("pow", lambda x, y: x**y) - OmegaConf.register_new_resolver("subtract", lambda x, y: x - y) - OmegaConf.register_new_resolver("range", lambda x: list(range(x))) - OmegaConf.register_new_resolver("int", lambda x: int(x)) - OmegaConf.register_new_resolver("ceil_int", lambda x: int(math.ceil(x))) - OmegaConf.register_new_resolver("merge", lambda *x: OmegaConf.merge(*x)) - OmegaConf.register_new_resolver("string", lambda x: str(x)) - - -def setup_distributed_backend(backend, timeout_mins): - """ - Initialize torch.distributed and set the CUDA device. - Expects environment variables to be set as per - https://pytorch.org/docs/stable/distributed.html#environment-variable-initialization - along with the environ variable "LOCAL_RANK" which is used to set the CUDA device. - """ - # enable TORCH_NCCL_ASYNC_ERROR_HANDLING to ensure dist nccl ops time out after timeout_mins - # of waiting - os.environ["TORCH_NCCL_ASYNC_ERROR_HANDLING"] = "1" - logging.info(f"Setting up torch.distributed with a timeout of {timeout_mins} mins") - dist.init_process_group(backend=backend, timeout=timedelta(minutes=timeout_mins)) - return dist.get_rank() - - -def get_machine_local_and_dist_rank(): - """ - Get the distributed and local rank of the current gpu. - """ - local_rank = int(os.environ.get("LOCAL_RANK", None)) - distributed_rank = int(os.environ.get("RANK", None)) - assert ( - local_rank is not None and distributed_rank is not None - ), "Please the set the RANK and LOCAL_RANK environment variables." - return local_rank, distributed_rank - - -def print_cfg(cfg): - """ - Supports printing both Hydra DictConfig and also the AttrDict config - """ - logging.info("Training with config:") - logging.info(OmegaConf.to_yaml(cfg)) - - -def set_seeds(seed_value, max_epochs, dist_rank): - """ - Set the python random, numpy and torch seed for each gpu. Also set the CUDA - seeds if the CUDA is available. This ensures deterministic nature of the training. - """ - # Since in the pytorch sampler, we increment the seed by 1 for every epoch. - seed_value = (seed_value + dist_rank) * max_epochs - logging.info(f"MACHINE SEED: {seed_value}") - random.seed(seed_value) - np.random.seed(seed_value) - torch.manual_seed(seed_value) - if torch.cuda.is_available(): - torch.cuda.manual_seed_all(seed_value) - - -def makedir(dir_path): - """ - Create the directory if it does not exist. - """ - is_success = False - try: - if not g_pathmgr.exists(dir_path): - g_pathmgr.mkdirs(dir_path) - is_success = True - except BaseException: - logging.info(f"Error creating directory: {dir_path}") - return is_success - - -def is_dist_avail_and_initialized(): - if not dist.is_available(): - return False - if not dist.is_initialized(): - return False - return True - - -def get_amp_type(amp_type: Optional[str] = None): - if amp_type is None: - return None - assert amp_type in ["bfloat16", "float16"], "Invalid Amp type." - if amp_type == "bfloat16": - return torch.bfloat16 - else: - return torch.float16 - - -def log_env_variables(): - env_keys = sorted(list(os.environ.keys())) - st = "" - for k in env_keys: - v = os.environ[k] - st += f"{k}={v}\n" - logging.info("Logging ENV_VARIABLES") - logging.info(st) - - -class AverageMeter: - """Computes and stores the average and current value""" - - def __init__(self, name, device, fmt=":f"): - self.name = name - self.fmt = fmt - self.device = device - self.reset() - - def reset(self): - self.val = 0 - self.avg = 0 - self.sum = 0 - self.count = 0 - self._allow_updates = True - - def update(self, val, n=1): - self.val = val - self.sum += val * n - self.count += n - self.avg = self.sum / self.count - - def __str__(self): - fmtstr = "{name}: {val" + self.fmt + "} ({avg" + self.fmt + "})" - return fmtstr.format(**self.__dict__) - - -class MemMeter: - """Computes and stores the current, avg, and max of peak Mem usage per iteration""" - - def __init__(self, name, device, fmt=":f"): - self.name = name - self.fmt = fmt - self.device = device - self.reset() - - def reset(self): - self.val = 0 # Per iteration max usage - self.avg = 0 # Avg per iteration max usage - self.peak = 0 # Peak usage for lifetime of program - self.sum = 0 - self.count = 0 - self._allow_updates = True - - def update(self, n=1, reset_peak_usage=True): - self.val = torch.cuda.max_memory_allocated() // 1e9 - self.sum += self.val * n - self.count += n - self.avg = self.sum / self.count - self.peak = max(self.peak, self.val) - if reset_peak_usage: - torch.cuda.reset_peak_memory_stats() - - def __str__(self): - fmtstr = ( - "{name}: {val" - + self.fmt - + "} ({avg" - + self.fmt - + "}/{peak" - + self.fmt - + "})" - ) - return fmtstr.format(**self.__dict__) - - -def human_readable_time(time_seconds): - time = int(time_seconds) - minutes, seconds = divmod(time, 60) - hours, minutes = divmod(minutes, 60) - days, hours = divmod(hours, 24) - return f"{days:02}d {hours:02}h {minutes:02}m" - - -class DurationMeter: - def __init__(self, name, device, fmt=":f"): - self.name = name - self.device = device - self.fmt = fmt - self.val = 0 - - def reset(self): - self.val = 0 - - def update(self, val): - self.val = val - - def add(self, val): - self.val += val - - def __str__(self): - return f"{self.name}: {human_readable_time(self.val)}" - - -class ProgressMeter: - def __init__(self, num_batches, meters, real_meters, prefix=""): - self.batch_fmtstr = self._get_batch_fmtstr(num_batches) - self.meters = meters - self.real_meters = real_meters - self.prefix = prefix - - def display(self, batch, enable_print=False): - entries = [self.prefix + self.batch_fmtstr.format(batch)] - entries += [str(meter) for meter in self.meters] - entries += [ - " | ".join( - [ - f"{os.path.join(name, subname)}: {val:.4f}" - for subname, val in meter.compute().items() - ] - ) - for name, meter in self.real_meters.items() - ] - logging.info(" | ".join(entries)) - if enable_print: - print(" | ".join(entries)) - - def _get_batch_fmtstr(self, num_batches): - num_digits = len(str(num_batches // 1)) - fmt = "{:" + str(num_digits) + "d}" - return "[" + fmt + "/" + fmt.format(num_batches) + "]" - - -def get_resume_checkpoint(checkpoint_save_dir): - if not g_pathmgr.isdir(checkpoint_save_dir): - return None - ckpt_file = os.path.join(checkpoint_save_dir, "checkpoint.pt") - if not g_pathmgr.isfile(ckpt_file): - return None - - return ckpt_file