98 lines
3.4 KiB
Python
98 lines
3.4 KiB
Python
# 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
|