第一次提交Yolo项目
This commit is contained in:
lhr
6
mobileclip/modules/image/__init__.py
Normal file
6
mobileclip/modules/image/__init__.py
Normal file
@@ -0,0 +1,6 @@
|
||||
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
||||
|
||||
#
|
||||
# For licensing see accompanying LICENSE file.
|
||||
# Copyright (C) 2024 Apple Inc. All rights reserved.
|
||||
#
|
||||
97
mobileclip/modules/image/image_projection.py
Normal file
97
mobileclip/modules/image/image_projection.py
Normal file
@@ -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
|
||||
177
mobileclip/modules/image/replknet.py
Normal file
177
mobileclip/modules/image/replknet.py
Normal file
@@ -0,0 +1,177 @@
|
||||
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
||||
|
||||
#
|
||||
# For acknowledgment see accompanying ACKNOWLEDGMENTS file.
|
||||
# Copyright (C) 2024 Apple Inc. All rights reserved.
|
||||
#
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from timm.models.layers import SqueezeExcite
|
||||
|
||||
__all__ = ["ReparamLargeKernelConv"]
|
||||
|
||||
|
||||
class ReparamLargeKernelConv(nn.Module):
|
||||
"""Building Block of RepLKNet.
|
||||
|
||||
This class defines overparameterized large kernel conv block introduced in `RepLKNet
|
||||
<https://arxiv.org/abs/2203.06717>`_
|
||||
|
||||
Reference: https://github.com/DingXiaoH/RepLKNet-pytorch
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
kernel_size: int,
|
||||
stride: int,
|
||||
groups: int,
|
||||
small_kernel: int,
|
||||
inference_mode: bool = False,
|
||||
use_se: bool = False,
|
||||
activation: nn.Module = nn.GELU(),
|
||||
) -> None:
|
||||
"""Construct a ReparamLargeKernelConv module.
|
||||
|
||||
Args:
|
||||
in_channels: Number of input channels.
|
||||
out_channels: Number of output channels.
|
||||
kernel_size: Kernel size of the large kernel conv branch.
|
||||
stride: Stride size. Default: 1
|
||||
groups: Group number. Default: 1
|
||||
small_kernel: Kernel size of small kernel conv branch.
|
||||
inference_mode: If True, instantiates model in inference mode. Default: ``False``
|
||||
activation: Activation module. Default: ``nn.GELU``
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
self.stride = stride
|
||||
self.groups = groups
|
||||
self.in_channels = in_channels
|
||||
self.out_channels = out_channels
|
||||
self.activation = activation
|
||||
|
||||
self.kernel_size = kernel_size
|
||||
self.small_kernel = small_kernel
|
||||
self.padding = kernel_size // 2
|
||||
|
||||
# Check if SE is requested
|
||||
if use_se:
|
||||
self.se = SqueezeExcite(out_channels, rd_ratio=0.25)
|
||||
else:
|
||||
self.se = nn.Identity()
|
||||
|
||||
if inference_mode:
|
||||
self.lkb_reparam = nn.Conv2d(
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
kernel_size=kernel_size,
|
||||
stride=stride,
|
||||
padding=self.padding,
|
||||
dilation=1,
|
||||
groups=groups,
|
||||
bias=True,
|
||||
)
|
||||
else:
|
||||
self.lkb_origin = self._conv_bn(kernel_size=kernel_size, padding=self.padding)
|
||||
if small_kernel is not None:
|
||||
assert small_kernel <= kernel_size, (
|
||||
"The kernel size for re-param cannot be larger than the large kernel!"
|
||||
)
|
||||
self.small_conv = self._conv_bn(kernel_size=small_kernel, padding=small_kernel // 2)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""Apply forward pass."""
|
||||
if hasattr(self, "lkb_reparam"):
|
||||
out = self.lkb_reparam(x)
|
||||
else:
|
||||
out = self.lkb_origin(x)
|
||||
if hasattr(self, "small_conv"):
|
||||
out += self.small_conv(x)
|
||||
|
||||
return self.activation(self.se(out))
|
||||
|
||||
def get_kernel_bias(self) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Method to obtain re-parameterized kernel and bias. Reference: https://github.com/DingXiaoH/RepLKNet-pytorch.
|
||||
|
||||
Returns:
|
||||
Tuple of (kernel, bias) after fusing branches.
|
||||
"""
|
||||
eq_k, eq_b = self._fuse_bn(self.lkb_origin.conv, self.lkb_origin.bn)
|
||||
if hasattr(self, "small_conv"):
|
||||
small_k, small_b = self._fuse_bn(self.small_conv.conv, self.small_conv.bn)
|
||||
eq_b += small_b
|
||||
eq_k += nn.functional.pad(small_k, [(self.kernel_size - self.small_kernel) // 2] * 4)
|
||||
return eq_k, eq_b
|
||||
|
||||
def reparameterize(self) -> None:
|
||||
"""Following works like `RepVGG: Making VGG-style ConvNets Great Again` - https://arxiv.org/pdf/2101.03697.pdf.
|
||||
We re-parameterize multi-branched architecture used at training time to obtain a plain CNN-like
|
||||
structure for inference.
|
||||
"""
|
||||
eq_k, eq_b = self.get_kernel_bias()
|
||||
self.lkb_reparam = nn.Conv2d(
|
||||
in_channels=self.in_channels,
|
||||
out_channels=self.out_channels,
|
||||
kernel_size=self.kernel_size,
|
||||
stride=self.stride,
|
||||
padding=self.padding,
|
||||
dilation=self.lkb_origin.conv.dilation,
|
||||
groups=self.groups,
|
||||
bias=True,
|
||||
)
|
||||
|
||||
self.lkb_reparam.weight.data = eq_k
|
||||
self.lkb_reparam.bias.data = eq_b
|
||||
self.__delattr__("lkb_origin")
|
||||
if hasattr(self, "small_conv"):
|
||||
self.__delattr__("small_conv")
|
||||
|
||||
@staticmethod
|
||||
def _fuse_bn(conv: torch.Tensor, bn: nn.BatchNorm2d) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Method to fuse batchnorm layer with conv layer.
|
||||
|
||||
Args:
|
||||
conv: Convolutional kernel weights.
|
||||
bn: Batchnorm 2d layer.
|
||||
|
||||
Returns:
|
||||
Tuple of (kernel, bias) after fusing batchnorm.
|
||||
"""
|
||||
kernel = conv.weight
|
||||
running_mean = bn.running_mean
|
||||
running_var = bn.running_var
|
||||
gamma = bn.weight
|
||||
beta = bn.bias
|
||||
eps = bn.eps
|
||||
std = (running_var + eps).sqrt()
|
||||
t = (gamma / std).reshape(-1, 1, 1, 1)
|
||||
return kernel * t, beta - running_mean * gamma / std
|
||||
|
||||
def _conv_bn(self, kernel_size: int, padding: int = 0) -> nn.Sequential:
|
||||
"""Helper method to construct conv-batchnorm layers.
|
||||
|
||||
Args:
|
||||
kernel_size: Size of the convolution kernel.
|
||||
padding: Zero-padding size.
|
||||
|
||||
Returns:
|
||||
A nn.Sequential Conv-BN module.
|
||||
"""
|
||||
mod_list = nn.Sequential()
|
||||
mod_list.add_module(
|
||||
"conv",
|
||||
nn.Conv2d(
|
||||
in_channels=self.in_channels,
|
||||
out_channels=self.out_channels,
|
||||
kernel_size=kernel_size,
|
||||
stride=self.stride,
|
||||
padding=padding,
|
||||
groups=self.groups,
|
||||
bias=False,
|
||||
),
|
||||
)
|
||||
mod_list.add_module("bn", nn.BatchNorm2d(num_features=self.out_channels))
|
||||
return mod_list
|
||||
Reference in New Issue
Block a user