birder 0.2.3__py3-none-any.whl → 0.3.0__py3-none-any.whl

birder/common/training_cli.py

@@ -211,6 +211,12 @@ def add_training_schedule_args(parser: argparse.ArgumentParser, default_epochs:
     group.add_argument(
         "--stop-epoch", type=int, metavar="N", help="epoch to stop the training at (multi stage training)"
     )
+    group.add_argument(
+        "--steps-per-epoch",
+        type=int,
+        metavar="N",
+        help="virtual epoch length in steps, leave unset to use the full dataset",
+    )
     group.add_argument("--warmup-epochs", type=int, metavar="N", help="number of warmup epochs")
     group.add_argument("--warmup-steps", type=int, metavar="N", help="number of warmup optimizer steps")
     group.add_argument("--cooldown-epochs", type=int, metavar="N", help="number of cooldown epochs (linear to zero)")

birder/common/training_utils.py

@@ -17,6 +17,7 @@ from typing import Any
 from typing import Literal
 from typing import Optional
 from typing import Sized
+from typing import overload
 
 import numpy as np
 import torch
@@ -70,13 +71,7 @@ class RASampler(torch.utils.data.Sampler):
     """
 
     def __init__(
-        self,
-        dataset: Sized,
-        num_replicas: int,
-        rank: int,
-        shuffle: bool,
-        seed: int = 0,
-        repetitions: int = 3,
+        self, dataset: Sized, num_replicas: int, rank: int, shuffle: bool, seed: int = 0, repetitions: int = 3
     ) -> None:
         super().__init__()
         self.dataset = dataset
@@ -85,12 +80,11 @@ class RASampler(torch.utils.data.Sampler):
         self.epoch = 0
         self.num_samples = int(math.ceil(len(self.dataset) * float(repetitions) / self.num_replicas))
         self.total_size = self.num_samples * self.num_replicas
-        self.num_selected_samples = int(math.floor(len(self.dataset) // 256 * 256 / self.num_replicas))
         self.shuffle = shuffle
         self.seed = seed
         self.repetitions = repetitions
 
-    def __iter__(self) -> Iterator[list[int]]:
+    def __iter__(self) -> Iterator[int]:
         if self.shuffle is True:
             # Deterministically shuffle based on epoch
             g = torch.Generator()
@@ -100,18 +94,148 @@ class RASampler(torch.utils.data.Sampler):
             indices = list(range(len(self.dataset)))
 
         # Add extra samples to make it evenly divisible
-        indices = [ele for ele in indices for i in range(self.repetitions)]
-        indices += indices[: (self.total_size - len(indices))]
-        assert len(indices) == self.total_size
+        indices = [ele for ele in indices for _ in range(self.repetitions)]
+        if len(indices) < self.total_size:
+            indices += indices[: (self.total_size - len(indices))]
+        else:
+            indices = indices[: self.total_size]
 
-        # Subsample
+        # Shard by rank
         indices = indices[self.rank : self.total_size : self.num_replicas]
         assert len(indices) == self.num_samples
 
-        return iter(indices[: self.num_selected_samples])
+        yield from indices
+
+    def __len__(self) -> int:
+        return self.num_samples
+
+    def set_epoch(self, epoch: int) -> None:
+        self.epoch = epoch
+
+
+class InfiniteSampler(torch.utils.data.Sampler):
+    """
+    Infinite sampler that loops indefinitely over the dataset
+    """
+
+    def __init__(self, dataset: Sized, shuffle: bool, seed: int = 0) -> None:
+        super().__init__()
+        self.dataset = dataset
+        self.shuffle = shuffle
+        self.seed = seed
+        self.epoch = 0
+
+    def __iter__(self) -> Iterator[int]:
+        g = torch.Generator()
+        while True:
+            if self.shuffle is True:
+                g.manual_seed(self.seed + self.epoch)
+                indices = torch.randperm(len(self.dataset), generator=g).tolist()
+            else:
+                indices = list(range(len(self.dataset)))
+
+            yield from indices
+
+            logger.debug(f"InfiniteSampler finished epoch {self.epoch}")
+            self.epoch += 1
+
+    def __len__(self) -> int:
+        return len(self.dataset)
+
+    def set_epoch(self, epoch: int) -> None:
+        self.epoch = epoch
+
+
+class InfiniteDistributedSampler(torch.utils.data.Sampler):
+    """
+    Infinite distributed sampler that keeps a continuous shuffled stream per rank
+    """
+
+    def __init__(self, dataset: Sized, num_replicas: int, rank: int, shuffle: bool, seed: int = 0) -> None:
+        super().__init__()
+        self.dataset = dataset
+        self.num_replicas = num_replicas
+        self.rank = rank
+        self.shuffle = shuffle
+        self.seed = seed
+        self.epoch = 0
+        self.num_samples = int(math.ceil(len(self.dataset) / self.num_replicas))
+        self.total_size = self.num_samples * self.num_replicas
+
+    def __iter__(self) -> Iterator[int]:
+        g = torch.Generator()
+        while True:
+            if self.shuffle is True:
+                g.manual_seed(self.seed + self.epoch)
+                indices = torch.randperm(len(self.dataset), generator=g).tolist()
+            else:
+                indices = list(range(len(self.dataset)))
+
+            if len(indices) < self.total_size:
+                indices += indices[: (self.total_size - len(indices))]
+            else:
+                indices = indices[: self.total_size]
+
+            indices = indices[self.rank : self.total_size : self.num_replicas]
+            assert len(indices) == self.num_samples
+
+            yield from indices
+
+            logger.debug(f"InfiniteDistributedSampler finished epoch {self.epoch}")
+            self.epoch += 1
 
     def __len__(self) -> int:
-        return self.num_selected_samples
+        return self.num_samples
+
+    def set_epoch(self, epoch: int) -> None:
+        self.epoch = epoch
+
+
+class InfiniteRASampler(torch.utils.data.Sampler):
+    """
+    Infinite version of the repeated augmentation sampler
+    """
+
+    def __init__(
+        self, dataset: Sized, num_replicas: int, rank: int, shuffle: bool, seed: int = 0, repetitions: int = 3
+    ) -> None:
+        super().__init__()
+        self.dataset = dataset
+        self.num_replicas = num_replicas
+        self.rank = rank
+        self.epoch = 0
+        self.num_samples = int(math.ceil(len(self.dataset) * float(repetitions) / self.num_replicas))
+        self.total_size = self.num_samples * self.num_replicas
+        self.shuffle = shuffle
+        self.seed = seed
+        self.repetitions = repetitions
+
+    def __iter__(self) -> Iterator[int]:
+        g = torch.Generator()
+        while True:
+            if self.shuffle is True:
+                g.manual_seed(self.seed + self.epoch)
+                indices = torch.randperm(len(self.dataset), generator=g).tolist()
+            else:
+                indices = list(range(len(self.dataset)))
+
+            indices = [ele for ele in indices for _ in range(self.repetitions)]
+            if len(indices) < self.total_size:
+                indices += indices[: (self.total_size - len(indices))]
+            else:
+                indices = indices[: self.total_size]
+
+            # Shard by rank
+            indices = indices[self.rank : self.total_size : self.num_replicas]
+            assert len(indices) == self.num_samples
+
+            yield from indices
+
+            logger.debug(f"InfiniteRASampler finished epoch {self.epoch}")
+            self.epoch += 1
+
+    def __len__(self) -> int:
+        return self.num_samples
 
     def set_epoch(self, epoch: int) -> None:
         self.epoch = epoch
@@ -636,27 +760,87 @@ def get_amp_scaler(amp: bool, amp_dtype_str: str) -> tuple[Optional[torch.amp.Gr
     return (scaler, amp_dtype)
 
 
+@overload
 def get_samplers(
-    args: argparse.Namespace, training_dataset: torch.utils.data.Dataset, validation_dataset: torch.utils.data.Dataset
-) -> torch.utils.data.Sampler:
-    if args.distributed is True:
-        if args.ra_sampler is True:
-            train_sampler = RASampler(
-                training_dataset,
-                num_replicas=args.world_size,
-                rank=args.rank,
-                shuffle=True,
-                repetitions=args.ra_reps,
-            )
+    args: argparse.Namespace,
+    training_dataset: torch.utils.data.Dataset,
+    validation_dataset: torch.utils.data.Dataset,
+    infinite: bool = False,
+) -> tuple[torch.utils.data.Sampler, torch.utils.data.Sampler]: ...
 
-        else:
-            train_sampler = torch.utils.data.distributed.DistributedSampler(training_dataset, shuffle=True)
 
-        validation_sampler = torch.utils.data.distributed.DistributedSampler(validation_dataset, shuffle=False)
+@overload
+def get_samplers(
+    args: argparse.Namespace,
+    training_dataset: torch.utils.data.Dataset,
+    validation_dataset: None = None,
+    infinite: bool = False,
+) -> tuple[torch.utils.data.Sampler, None]: ...
+
+
+def get_samplers(
+    args: argparse.Namespace,
+    training_dataset: torch.utils.data.Dataset,
+    validation_dataset: Optional[torch.utils.data.Dataset] = None,
+    infinite: bool = False,
+) -> tuple[torch.utils.data.Sampler, Optional[torch.utils.data.Sampler]]:
+    if args.seed is None:
+        seed = int(torch.empty((), dtype=torch.int64).random_().item())
+        if is_dist_available_and_initialized() is True:
+            seed_tensor = torch.tensor(seed, dtype=torch.int64).cuda()
+            dist.broadcast(seed_tensor, src=0, async_op=False)
+            seed = int(seed_tensor.item())
+    else:
+        seed = args.seed
+
+    ra_sampler = getattr(args, "ra_sampler", False)
+    if args.distributed is True:
+        if infinite is True:
+            if ra_sampler is True:
+                train_sampler = InfiniteRASampler(
+                    training_dataset,
+                    num_replicas=args.world_size,
+                    rank=args.rank,
+                    shuffle=True,
+                    seed=seed,
+                    repetitions=args.ra_reps,
+                )
+            else:
+                train_sampler = InfiniteDistributedSampler(
+                    training_dataset, num_replicas=args.world_size, rank=args.rank, shuffle=True, seed=seed
+                )
+        else:
+            if ra_sampler is True:
+                train_sampler = RASampler(
+                    training_dataset,
+                    num_replicas=args.world_size,
+                    rank=args.rank,
+                    shuffle=True,
+                    seed=seed,
+                    repetitions=args.ra_reps,
+                )
+            else:
+                train_sampler = torch.utils.data.distributed.DistributedSampler(
+                    training_dataset, shuffle=True, seed=seed
+                )
+
+        if validation_dataset is None:
+            validation_sampler = None
+        else:
+            validation_sampler = torch.utils.data.distributed.DistributedSampler(validation_dataset, shuffle=False)
 
     else:
-        train_sampler = torch.utils.data.RandomSampler(training_dataset)
-        validation_sampler = torch.utils.data.SequentialSampler(validation_dataset)
+        if infinite is True:
+            train_sampler = InfiniteSampler(training_dataset, shuffle=True, seed=seed)
+        else:
+            generator = torch.Generator()
+            generator.manual_seed(seed)
+            train_sampler = torch.utils.data.RandomSampler(training_dataset, generator=generator)
+
+        if validation_dataset is None:
+            validation_sampler = None
+        else:
+            validation_sampler = torch.utils.data.SequentialSampler(validation_dataset)
 
     return (train_sampler, validation_sampler)
 

birder/data/collators/detection.py

@@ -98,6 +98,7 @@ class BatchRandomResizeCollator(DetectionCollator):
                 if isinstance(boxes, tv_tensors.BoundingBoxes) is False:
                     if boxes.numel() == 0:
                         boxes = boxes.reshape(0, 4)
+
                     boxes = tv_tensors.BoundingBoxes(
                         boxes, format=tv_tensors.BoundingBoxFormat.XYXY, canvas_size=F.get_size(image)
                     )

birder/data/dataloader/webdataset.py

@@ -22,9 +22,19 @@ def make_wds_loader(
     shuffle: bool = False,
     *,
     exact: bool = False,
+    infinite: bool = False,
 ) -> DataLoader:
+    assert exact is False or infinite is False
+
+    if infinite is True:
+        dataset_iterable = dataset.repeat()
+    elif exact is False:
+        dataset_iterable = dataset.repeat()
+    else:
+        dataset_iterable = dataset
+
     dataloader = wds.WebLoader(
-        dataset.repeat() if exact is False else dataset,
+        dataset_iterable,
         batch_size=batch_size,
         num_workers=num_workers,
         prefetch_factor=prefetch_factor,
@@ -43,7 +53,7 @@ def make_wds_loader(
         epoch_size = math.ceil(len(dataset) / (batch_size * world_size))
 
     dataloader = dataloader.with_length(epoch_size, silent=True)
-    if exact is False:
+    if exact is False and infinite is False:
         dataloader = dataloader.with_epoch(epoch_size)
 
     return dataloader

birder/kernels/load_kernel.py

@@ -14,11 +14,24 @@ logger = logging.getLogger(__name__)
 
 
 _CACHED_KERNELS: dict[str, ModuleType] = {}
+_CUSTOM_KERNELS_ENABLED = True
+
+
+def set_custom_kernels_enabled(enabled: bool) -> None:
+    global _CUSTOM_KERNELS_ENABLED  # pylint: disable=global-statement
+    _CUSTOM_KERNELS_ENABLED = enabled
+
+
+def is_custom_kernels_enabled() -> bool:
+    if os.environ.get("DISABLE_CUSTOM_KERNELS", "0") == "1":
+        return False
+
+    return _CUSTOM_KERNELS_ENABLED
 
 
 def load_msda() -> Optional[ModuleType]:
     name = "msda"
-    if torch.cuda.is_available() is False or os.environ.get("DISABLE_CUSTOM_KERNELS", "0") == "1":
+    if torch.cuda.is_available() is False or is_custom_kernels_enabled() is False:
         return None
 
     if name in _CACHED_KERNELS:
@@ -60,7 +73,7 @@ def load_msda() -> Optional[ModuleType]:
 
 def load_swattention() -> Optional[ModuleType]:
     name = "swattention"
-    if torch.cuda.is_available() is False or os.environ.get("DISABLE_CUSTOM_KERNELS", "0") == "1":
+    if torch.cuda.is_available() is False or is_custom_kernels_enabled() is False:
         return None
 
     if name in _CACHED_KERNELS:
@@ -103,7 +116,7 @@ def load_swattention() -> Optional[ModuleType]:
 
 def load_soft_nms() -> Optional[ModuleType]:
     name = "soft_nms"
-    if os.environ.get("DISABLE_CUSTOM_KERNELS", "0") == "1":
+    if is_custom_kernels_enabled() is False:
         return None
 
     if name in _CACHED_KERNELS:
@@ -120,14 +133,6 @@ def load_soft_nms() -> Optional[ModuleType]:
         soft_nms: Optional[ModuleType] = load(
             "soft_nms",
             src_files,
-            with_cuda=True,
-            extra_cflags=["-DWITH_CUDA=1"],
-            extra_cuda_cflags=[
-                "-DCUDA_HAS_FP16=1",
-                "-D__CUDA_NO_HALF_OPERATORS__",
-                "-D__CUDA_NO_HALF_CONVERSIONS__",
-                "-D__CUDA_NO_HALF2_OPERATORS__",
-            ],
         )
 
     if soft_nms is not None:

birder/kernels/soft_nms/soft_nms.cpp

@@ -61,24 +61,23 @@ void update_sorting_order(torch::Tensor& boxes, torch::Tensor& scores, torch::Te
     std::tie(max_score, t_max_idx) = torch::max(scores.index({Slice(idx + 1, None)}), 0);
 
     // max_idx is computed from sliced data, therefore need to convert it to "global" max idx
-    auto max_idx = t_max_idx.item<int>() + idx + 1;
-
-    if (scores.index({idx}).item<float>() < max_score.item<float>()) {
-        auto boxes_idx = boxes.index({idx}).clone();
-        auto boxes_max = boxes.index({max_idx}).clone();
-        boxes.index({idx}) = boxes_max;
-        boxes.index({max_idx}) = boxes_idx;
-
-        auto scores_idx = scores.index({idx}).clone();
-        auto scores_max = scores.index({max_idx}).clone();
-        scores.index({idx}) = scores_max;
-        scores.index({max_idx}) = scores_idx;
-
-        auto areas_idx = areas.index({idx}).clone();
-        auto areas_max = areas.index({max_idx}).clone();
-        areas.index({idx}) = areas_max;
-        areas.index({max_idx}) = areas_idx;
-    }
+    auto max_idx = t_max_idx + (idx + 1);
+    auto should_swap = scores.index({idx}) < max_score;
+
+    auto boxes_idx = boxes.index({idx}).clone();
+    auto boxes_max = boxes.index({max_idx}).clone();
+    boxes.index_put_({idx}, torch::where(should_swap, boxes_max, boxes_idx));
+    boxes.index_put_({max_idx}, torch::where(should_swap, boxes_idx, boxes_max));
+
+    auto scores_idx = scores.index({idx}).clone();
+    auto scores_max = scores.index({max_idx}).clone();
+    scores.index_put_({idx}, torch::where(should_swap, scores_max, scores_idx));
+    scores.index_put_({max_idx}, torch::where(should_swap, scores_idx, scores_max));
+
+    auto areas_idx = areas.index({idx}).clone();
+    auto areas_max = areas.index({max_idx}).clone();
+    areas.index_put_({idx}, torch::where(should_swap, areas_max, areas_idx));
+    areas.index_put_({max_idx}, torch::where(should_swap, areas_idx, areas_max));
 }
 
 std::tuple<torch::Tensor, torch::Tensor> soft_nms(

birder/net/cait.py

@@ -268,14 +268,15 @@ class CaiT(BaseNet):
         super().adjust_size(new_size)
 
         # Add back class tokens
-        self.pos_embed = nn.Parameter(
-            adjust_position_embedding(
+        with torch.no_grad():
+            pos_embed = adjust_position_embedding(
                 self.pos_embed,
                 (old_size[0] // self.patch_size[0], old_size[1] // self.patch_size[1]),
                 (new_size[0] // self.patch_size[0], new_size[1] // self.patch_size[1]),
                 0,
             )
-        )
+
+        self.pos_embed = nn.Parameter(pos_embed)
 
 
 registry.register_model_config(

birder/net/convnext_v1.py

@@ -195,6 +195,11 @@ class ConvNeXt_v1(DetectorBackbone, PreTrainEncoder, MaskedTokenRetentionMixin):
         return self.features(x)
 
 
+registry.register_model_config(
+    "convnext_v1_nano",  # Not in the original v1, taken from v2
+    ConvNeXt_v1,
+    config={"in_channels": [80, 160, 320, 640], "num_layers": [2, 2, 8, 2], "drop_path_rate": 0.1},
+)
 registry.register_model_config(
     "convnext_v1_tiny",
     ConvNeXt_v1,

birder/net/crossformer.py

@@ -98,15 +98,17 @@ class Attention(nn.Module):
         self.proj_drop = nn.Dropout(proj_drop)
 
     def define_bias_table(self) -> None:
-        position_bias_h = torch.arange(1 - self.group_size[0], self.group_size[0])
-        position_bias_w = torch.arange(1 - self.group_size[1], self.group_size[1])
+        device = next(self.pos.parameters()).device
+        position_bias_h = torch.arange(1 - self.group_size[0], self.group_size[0], device=device)
+        position_bias_w = torch.arange(1 - self.group_size[1], self.group_size[1], device=device)
         biases = torch.stack(torch.meshgrid([position_bias_h, position_bias_w], indexing="ij"))  # 2, 2Wh-1, 2W2-1
         biases = biases.flatten(1).transpose(0, 1).float()
         self.biases = nn.Buffer(biases)
 
     def define_relative_position_index(self) -> None:
-        coords_h = torch.arange(self.group_size[0])
-        coords_w = torch.arange(self.group_size[1])
+        device = self.biases.device
+        coords_h = torch.arange(self.group_size[0], device=device)
+        coords_w = torch.arange(self.group_size[1], device=device)
         coords = torch.stack(torch.meshgrid([coords_h, coords_w], indexing="ij"))  # 2, Wh, Ww
         coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
         relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww
@@ -430,32 +432,33 @@ class CrossFormer(DetectorBackbone):
 
         new_patch_resolution = (new_size[0] // self.patch_sizes[0], new_size[1] // self.patch_sizes[0])
         input_resolution = new_patch_resolution
-        for mod in self.body.modules():
-            if isinstance(mod, CrossFormerStage):
-                for m in mod.modules():
-                    if isinstance(m, PatchMerging):
-                        m.input_resolution = input_resolution
-                        input_resolution = (input_resolution[0] // 2, input_resolution[1] // 2)
-                    elif isinstance(m, CrossFormerBlock):
-                        m.input_resolution = input_resolution
-
-                mod.resolution = input_resolution
-
-        new_group_size = (int(new_size[0] / (2**5)), int(new_size[1] / (2**5)))
-        for m in self.body.modules():
-            if isinstance(m, CrossFormerBlock):
-                m.group_size = new_group_size
-                if m.input_resolution[0] <= m.group_size[0]:
-                    m.use_lda = False
-                    m.group_size = (m.input_resolution[0], m.group_size[1])
-                if m.input_resolution[1] <= m.group_size[1]:
-                    m.use_lda = False
-                    m.group_size = (m.group_size[0], m.input_resolution[1])
-
-            elif isinstance(m, Attention):
-                m.group_size = new_group_size
-                m.define_bias_table()
-                m.define_relative_position_index()
+        with torch.no_grad():
+            for mod in self.body.modules():
+                if isinstance(mod, CrossFormerStage):
+                    for m in mod.modules():
+                        if isinstance(m, PatchMerging):
+                            m.input_resolution = input_resolution
+                            input_resolution = (input_resolution[0] // 2, input_resolution[1] // 2)
+                        elif isinstance(m, CrossFormerBlock):
+                            m.input_resolution = input_resolution
+
+                    mod.resolution = input_resolution
+
+            new_group_size = (int(new_size[0] / (2**5)), int(new_size[1] / (2**5)))
+            for m in self.body.modules():
+                if isinstance(m, CrossFormerBlock):
+                    m.group_size = new_group_size
+                    if m.input_resolution[0] <= m.group_size[0]:
+                        m.use_lda = False
+                        m.group_size = (m.input_resolution[0], m.group_size[1])
+                    if m.input_resolution[1] <= m.group_size[1]:
+                        m.use_lda = False
+                        m.group_size = (m.group_size[0], m.input_resolution[1])
+
+                elif isinstance(m, Attention):
+                    m.group_size = new_group_size
+                    m.define_bias_table()
+                    m.define_relative_position_index()
 
 
 registry.register_model_config(

birder/net/crossvit.py

@@ -359,9 +359,10 @@ class CrossViT(BaseNet):
             old_w = old_size[1] // self.patch_size[i]
             h = new_size[0] // self.patch_size[i]
             w = new_size[1] // self.patch_size[i]
-            self.pos_embed[i] = nn.Parameter(
-                adjust_position_embedding(self.pos_embed[i], (old_h, old_w), (h, w), num_prefix_tokens=1)
-            )
+            with torch.no_grad():
+                pos_embed = adjust_position_embedding(self.pos_embed[i], (old_h, old_w), (h, w), num_prefix_tokens=1)
+
+            self.pos_embed[i] = nn.Parameter(pos_embed)
 
 
 registry.register_model_config(

birder/net/deit.py

@@ -187,14 +187,14 @@ class DeiT(BaseNet):
         num_prefix_tokens = 2
 
         # Add back class tokens
-        self.pos_embedding = nn.Parameter(
-            adjust_position_embedding(
+        with torch.no_grad():
+            pos_embedding = adjust_position_embedding(
                 self.pos_embedding,
                 (old_size[0] // self.patch_size, old_size[1] // self.patch_size),
                 (new_size[0] // self.patch_size, new_size[1] // self.patch_size),
                 num_prefix_tokens,
             )
-        )
+        self.pos_embedding = nn.Parameter(pos_embedding)
 
 
 registry.register_model_config(

birder/net/deit3.py

@@ -355,14 +355,14 @@ class DeiT3(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, MaskedT
             num_prefix_tokens = 0
 
         # Add back class tokens
-        self.pos_embedding = nn.Parameter(
-            adjust_position_embedding(
+        with torch.no_grad():
+            pos_embedding = adjust_position_embedding(
                 self.pos_embedding,
                 (old_size[0] // self.patch_size, old_size[1] // self.patch_size),
                 (new_size[0] // self.patch_size, new_size[1] // self.patch_size),
                 num_prefix_tokens,
             )
-        )
+        self.pos_embedding = nn.Parameter(pos_embedding)
 
 
 registry.register_model_config(

birder/net/detection/deformable_detr.py

@@ -757,11 +757,8 @@ class Deformable_DETR(DetectionBaseNet):
         for s, l, b in zip(scores, labels, boxes):
             # Non-maximum suppression
             if self.soft_nms is not None:
-                # Actually much faster on CPU
-                device = b.device
-                (soft_scores, keep) = self.soft_nms(b.cpu(), s.cpu(), l.cpu(), score_threshold=0.001)
-                keep = keep.to(device)
-                s[keep] = soft_scores.to(device)
+                (soft_scores, keep) = self.soft_nms(b, s, l, score_threshold=0.001)
+                s[keep] = soft_scores
 
                 b = b[keep]
                 s = s[keep]

birder/net/detection/detr.py

@@ -465,11 +465,8 @@ class DETR(DetectionBaseNet):
         for s, l, b in zip(scores, labels, boxes):
             # Non-maximum suppression
             if self.soft_nms is not None:
-                # Actually much faster on CPU
-                device = b.device
-                (soft_scores, keep) = self.soft_nms(b.cpu(), s.cpu(), l.cpu(), score_threshold=0.001)
-                keep = keep.to(device)
-                s[keep] = soft_scores.to(device)
+                (soft_scores, keep) = self.soft_nms(b, s, l, score_threshold=0.001)
+                s[keep] = soft_scores
 
                 b = b[keep]
                 s = s[keep]

birder/net/detection/efficientdet.py

@@ -685,13 +685,8 @@ class EfficientDet(DetectionBaseNet):
 
             # Non-maximum suppression
             if self.soft_nms is not None:
-                # Actually much faster on CPU
-                device = image_boxes.device
-                (soft_scores, keep) = self.soft_nms(
-                    image_boxes.cpu(), image_scores.cpu(), image_labels.cpu(), score_threshold=0.001
-                )
-                keep = keep.to(device)
-                image_scores[keep] = soft_scores.to(device)
+                (soft_scores, keep) = self.soft_nms(image_boxes, image_scores, image_labels, score_threshold=0.001)
+                image_scores[keep] = soft_scores
             else:
                 keep = box_ops.batched_nms(image_boxes, image_scores, image_labels, self.nms_thresh)