第一次提交Yolo项目

mobileclip/modules/image/image_projection.py

@@ -0,0 +1,97 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+# For licensing see accompanying LICENSE file.
+# Copyright (C) 2024 Apple Inc. All Rights Reserved.
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+from mobileclip import logger
+
+
+class GlobalPool(nn.Module):
+    """This layers applies global pooling over a 4D or 5D input tensor.
+
+    Args:
+        pool_type (Optional[str]): Pooling type. It can be mean, rms, or abs. Default: `mean`
+        keep_dim (Optional[bool]): Do not squeeze the dimensions of a tensor. Default: `False`
+
+    Notes:
+        - Input: :math:`(N, C, H, W)` or :math:`(N, C, D, H, W)`
+        - Output: :math:`(N, C, 1, 1)` or :math:`(N, C, 1, 1, 1)` if keep_dim else :math:`(N, C)`
+    """
+
+    pool_types = ["mean", "rms", "abs"]
+
+    def __init__(self, pool_type: str | None = "mean", keep_dim: bool | None = False, *args, **kwargs) -> None:
+        super().__init__()
+        if pool_type not in self.pool_types:
+            logger.error(f"Supported pool types are: {self.pool_types}. Got {pool_type}")
+        self.pool_type = pool_type
+        self.keep_dim = keep_dim
+
+    def _global_pool(self, x: Tensor, dims: list):
+        if self.pool_type == "rms":  # root mean square
+            x = x**2
+            x = torch.mean(x, dim=dims, keepdim=self.keep_dim)
+            x = x**-0.5
+        elif self.pool_type == "abs":  # absolute
+            x = torch.mean(torch.abs(x), dim=dims, keepdim=self.keep_dim)
+        else:
+            # default is mean
+            # same as AdaptiveAvgPool
+            x = torch.mean(x, dim=dims, keepdim=self.keep_dim)
+        return x
+
+    def forward(self, x: Tensor) -> Tensor:
+        if x.dim() == 4:
+            dims = [-2, -1]
+        elif x.dim() == 5:
+            dims = [-3, -2, -1]
+        else:
+            raise NotImplementedError("Currently 2D and 3D global pooling supported")
+        return self._global_pool(x, dims=dims)
+
+
+class GlobalPool2D(nn.Module):
+    """This class implements global pooling with linear projection."""
+
+    def __init__(self, in_dim: int, out_dim: int, *args, **kwargs) -> None:
+        super().__init__()
+        scale = in_dim**-0.5
+        self.pool = GlobalPool(pool_type="mean", keep_dim=False)
+        self.proj = nn.Parameter(scale * torch.randn(size=(in_dim, out_dim)))
+        self.in_dim = in_dim
+        self.out_dim = out_dim
+
+    def forward(self, x: Tensor, *args, **kwargs) -> Tensor:
+        # x is of shape [batch, in_dim]
+        assert x.dim() == 4, f"Input should be 4-dimensional (Batch x in_dim x in_height x in_width). Got: {x.shape}"
+
+        # [batch, in_dim, in_height, in_width] --> [batch, in_dim]
+        x = self.pool(x)
+        # [batch, in_dim]  x [in_dim, out_dim] --> [batch, out_dim]
+        x = x @ self.proj
+        return x
+
+
+class SimpleImageProjectionHead(nn.Module):
+    """This class implements linear projection head."""
+
+    def __init__(self, in_dim: int, out_dim: int) -> None:
+        super().__init__()
+        scale = in_dim**-0.5
+        self.proj = nn.Parameter(scale * torch.randn(size=(in_dim, out_dim)))
+        self.in_dim = in_dim
+        self.out_dim = out_dim
+
+    def forward(self, x: Tensor, *args, **kwargs) -> Tensor:
+        # x is of shape [batch, in_dim]
+        assert x.dim() == 2, f"Input should be 2-dimensional (Batch x in_dim). Got: {x.shape}"
+
+        # [batch, in_dim] x [in_dim, out_dim] --> [batch, out_dim]
+        x = x @ self.proj
+        return x