Implement all missing docstrings (#5298)
Co-authored-by: snyk-bot <snyk-bot@snyk.io> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Glenn Jocher
GitHub
parent
e7f0658744
commit
7fd5dcbd86
26 changed files with 649 additions and 79 deletions
|
|
@ -37,7 +37,12 @@ class DFL(nn.Module):
|
|||
class Proto(nn.Module):
|
||||
"""YOLOv8 mask Proto module for segmentation models."""
|
||||
|
||||
def __init__(self, c1, c_=256, c2=32): # ch_in, number of protos, number of masks
|
||||
def __init__(self, c1, c_=256, c2=32):
|
||||
"""
|
||||
Initializes the YOLOv8 mask Proto module with specified number of protos and masks.
|
||||
|
||||
Input arguments are ch_in, number of protos, number of masks.
|
||||
"""
|
||||
super().__init__()
|
||||
self.cv1 = Conv(c1, c_, k=3)
|
||||
self.upsample = nn.ConvTranspose2d(c_, c_, 2, 2, 0, bias=True) # nn.Upsample(scale_factor=2, mode='nearest')
|
||||
|
|
@ -124,7 +129,12 @@ class SPP(nn.Module):
|
|||
class SPPF(nn.Module):
|
||||
"""Spatial Pyramid Pooling - Fast (SPPF) layer for YOLOv5 by Glenn Jocher."""
|
||||
|
||||
def __init__(self, c1, c2, k=5): # equivalent to SPP(k=(5, 9, 13))
|
||||
def __init__(self, c1, c2, k=5):
|
||||
"""
|
||||
Initializes the SPPF layer with given input/output channels and kernel size.
|
||||
|
||||
This module is equivalent to SPP(k=(5, 9, 13)).
|
||||
"""
|
||||
super().__init__()
|
||||
c_ = c1 // 2 # hidden channels
|
||||
self.cv1 = Conv(c1, c_, 1, 1)
|
||||
|
|
@ -142,7 +152,8 @@ class SPPF(nn.Module):
|
|||
class C1(nn.Module):
|
||||
"""CSP Bottleneck with 1 convolution."""
|
||||
|
||||
def __init__(self, c1, c2, n=1): # ch_in, ch_out, number
|
||||
def __init__(self, c1, c2, n=1):
|
||||
"""Initializes the CSP Bottleneck with configurations for 1 convolution with arguments ch_in, ch_out, number."""
|
||||
super().__init__()
|
||||
self.cv1 = Conv(c1, c2, 1, 1)
|
||||
self.m = nn.Sequential(*(Conv(c2, c2, 3) for _ in range(n)))
|
||||
|
|
@ -156,7 +167,10 @@ class C1(nn.Module):
|
|||
class C2(nn.Module):
|
||||
"""CSP Bottleneck with 2 convolutions."""
|
||||
|
||||
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
|
||||
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
|
||||
"""Initializes the CSP Bottleneck with 2 convolutions module with arguments ch_in, ch_out, number, shortcut,
|
||||
groups, expansion.
|
||||
"""
|
||||
super().__init__()
|
||||
self.c = int(c2 * e) # hidden channels
|
||||
self.cv1 = Conv(c1, 2 * self.c, 1, 1)
|
||||
|
|
@ -173,7 +187,10 @@ class C2(nn.Module):
|
|||
class C2f(nn.Module):
|
||||
"""Faster Implementation of CSP Bottleneck with 2 convolutions."""
|
||||
|
||||
def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
|
||||
def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5):
|
||||
"""Initialize CSP bottleneck layer with two convolutions with arguments ch_in, ch_out, number, shortcut, groups,
|
||||
expansion.
|
||||
"""
|
||||
super().__init__()
|
||||
self.c = int(c2 * e) # hidden channels
|
||||
self.cv1 = Conv(c1, 2 * self.c, 1, 1)
|
||||
|
|
@ -196,7 +213,8 @@ class C2f(nn.Module):
|
|||
class C3(nn.Module):
|
||||
"""CSP Bottleneck with 3 convolutions."""
|
||||
|
||||
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
|
||||
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
|
||||
"""Initialize the CSP Bottleneck with given channels, number, shortcut, groups, and expansion values."""
|
||||
super().__init__()
|
||||
c_ = int(c2 * e) # hidden channels
|
||||
self.cv1 = Conv(c1, c_, 1, 1)
|
||||
|
|
@ -259,7 +277,8 @@ class C3Ghost(C3):
|
|||
class GhostBottleneck(nn.Module):
|
||||
"""Ghost Bottleneck https://github.com/huawei-noah/ghostnet."""
|
||||
|
||||
def __init__(self, c1, c2, k=3, s=1): # ch_in, ch_out, kernel, stride
|
||||
def __init__(self, c1, c2, k=3, s=1):
|
||||
"""Initializes GhostBottleneck module with arguments ch_in, ch_out, kernel, stride."""
|
||||
super().__init__()
|
||||
c_ = c2 // 2
|
||||
self.conv = nn.Sequential(
|
||||
|
|
@ -277,7 +296,10 @@ class GhostBottleneck(nn.Module):
|
|||
class Bottleneck(nn.Module):
|
||||
"""Standard bottleneck."""
|
||||
|
||||
def __init__(self, c1, c2, shortcut=True, g=1, k=(3, 3), e=0.5): # ch_in, ch_out, shortcut, groups, kernels, expand
|
||||
def __init__(self, c1, c2, shortcut=True, g=1, k=(3, 3), e=0.5):
|
||||
"""Initializes a bottleneck module with given input/output channels, shortcut option, group, kernels, and
|
||||
expansion.
|
||||
"""
|
||||
super().__init__()
|
||||
c_ = int(c2 * e) # hidden channels
|
||||
self.cv1 = Conv(c1, c_, k[0], 1)
|
||||
|
|
@ -292,7 +314,8 @@ class Bottleneck(nn.Module):
|
|||
class BottleneckCSP(nn.Module):
|
||||
"""CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks."""
|
||||
|
||||
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
|
||||
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
|
||||
"""Initializes the CSP Bottleneck given arguments for ch_in, ch_out, number, shortcut, groups, expansion."""
|
||||
super().__init__()
|
||||
c_ = int(c2 * e) # hidden channels
|
||||
self.cv1 = Conv(c1, c_, 1, 1)
|
||||
|
|
|
|||
|
|
@ -88,6 +88,7 @@ class DWConv(Conv):
|
|||
"""Depth-wise convolution."""
|
||||
|
||||
def __init__(self, c1, c2, k=1, s=1, d=1, act=True): # ch_in, ch_out, kernel, stride, dilation, activation
|
||||
"""Initialize Depth-wise convolution with given parameters."""
|
||||
super().__init__(c1, c2, k, s, g=math.gcd(c1, c2), d=d, act=act)
|
||||
|
||||
|
||||
|
|
@ -95,6 +96,7 @@ class DWConvTranspose2d(nn.ConvTranspose2d):
|
|||
"""Depth-wise transpose convolution."""
|
||||
|
||||
def __init__(self, c1, c2, k=1, s=1, p1=0, p2=0): # ch_in, ch_out, kernel, stride, padding, padding_out
|
||||
"""Initialize DWConvTranspose2d class with given parameters."""
|
||||
super().__init__(c1, c2, k, s, p1, p2, groups=math.gcd(c1, c2))
|
||||
|
||||
|
||||
|
|
@ -121,12 +123,18 @@ class ConvTranspose(nn.Module):
|
|||
class Focus(nn.Module):
|
||||
"""Focus wh information into c-space."""
|
||||
|
||||
def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups
|
||||
def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True):
|
||||
"""Initializes Focus object with user defined channel, convolution, padding, group and activation values."""
|
||||
super().__init__()
|
||||
self.conv = Conv(c1 * 4, c2, k, s, p, g, act=act)
|
||||
# self.contract = Contract(gain=2)
|
||||
|
||||
def forward(self, x): # x(b,c,w,h) -> y(b,4c,w/2,h/2)
|
||||
def forward(self, x):
|
||||
"""
|
||||
Applies convolution to concatenated tensor and returns the output.
|
||||
|
||||
Input shape is (b,c,w,h) and output shape is (b,4c,w/2,h/2).
|
||||
"""
|
||||
return self.conv(torch.cat((x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]), 1))
|
||||
# return self.conv(self.contract(x))
|
||||
|
||||
|
|
@ -134,7 +142,10 @@ class Focus(nn.Module):
|
|||
class GhostConv(nn.Module):
|
||||
"""Ghost Convolution https://github.com/huawei-noah/ghostnet."""
|
||||
|
||||
def __init__(self, c1, c2, k=1, s=1, g=1, act=True): # ch_in, ch_out, kernel, stride, groups
|
||||
def __init__(self, c1, c2, k=1, s=1, g=1, act=True):
|
||||
"""Initializes the GhostConv object with input channels, output channels, kernel size, stride, groups and
|
||||
activation.
|
||||
"""
|
||||
super().__init__()
|
||||
c_ = c2 // 2 # hidden channels
|
||||
self.cv1 = Conv(c1, c_, k, s, None, g, act=act)
|
||||
|
|
@ -280,7 +291,8 @@ class SpatialAttention(nn.Module):
|
|||
class CBAM(nn.Module):
|
||||
"""Convolutional Block Attention Module."""
|
||||
|
||||
def __init__(self, c1, kernel_size=7): # ch_in, kernels
|
||||
def __init__(self, c1, kernel_size=7):
|
||||
"""Initialize CBAM with given input channel (c1) and kernel size."""
|
||||
super().__init__()
|
||||
self.channel_attention = ChannelAttention(c1)
|
||||
self.spatial_attention = SpatialAttention(kernel_size)
|
||||
|
|
|
|||
|
|
@ -25,7 +25,8 @@ class Detect(nn.Module):
|
|||
anchors = torch.empty(0) # init
|
||||
strides = torch.empty(0) # init
|
||||
|
||||
def __init__(self, nc=80, ch=()): # detection layer
|
||||
def __init__(self, nc=80, ch=()):
|
||||
"""Initializes the YOLOv8 detection layer with specified number of classes and channels."""
|
||||
super().__init__()
|
||||
self.nc = nc # number of classes
|
||||
self.nl = len(ch) # number of detection layers
|
||||
|
|
@ -149,7 +150,10 @@ class Pose(Detect):
|
|||
class Classify(nn.Module):
|
||||
"""YOLOv8 classification head, i.e. x(b,c1,20,20) to x(b,c2)."""
|
||||
|
||||
def __init__(self, c1, c2, k=1, s=1, p=None, g=1): # ch_in, ch_out, kernel, stride, padding, groups
|
||||
def __init__(self, c1, c2, k=1, s=1, p=None, g=1):
|
||||
"""Initializes YOLOv8 classification head with specified input and output channels, kernel size, stride,
|
||||
padding, and groups.
|
||||
"""
|
||||
super().__init__()
|
||||
c_ = 1280 # efficientnet_b0 size
|
||||
self.conv = Conv(c1, c_, k, s, p, g)
|
||||
|
|
@ -166,6 +170,13 @@ class Classify(nn.Module):
|
|||
|
||||
|
||||
class RTDETRDecoder(nn.Module):
|
||||
"""
|
||||
Real-Time Deformable Transformer Decoder (RTDETRDecoder) module for object detection.
|
||||
|
||||
This decoder module utilizes Transformer architecture along with deformable convolutions to predict bounding boxes
|
||||
and class labels for objects in an image. It integrates features from multiple layers and runs through a series of
|
||||
Transformer decoder layers to output the final predictions.
|
||||
"""
|
||||
export = False # export mode
|
||||
|
||||
def __init__(
|
||||
|
|
@ -186,6 +197,26 @@ class RTDETRDecoder(nn.Module):
|
|||
label_noise_ratio=0.5,
|
||||
box_noise_scale=1.0,
|
||||
learnt_init_query=False):
|
||||
"""
|
||||
Initializes the RTDETRDecoder module with the given parameters.
|
||||
|
||||
Args:
|
||||
nc (int): Number of classes. Default is 80.
|
||||
ch (tuple): Channels in the backbone feature maps. Default is (512, 1024, 2048).
|
||||
hd (int): Dimension of hidden layers. Default is 256.
|
||||
nq (int): Number of query points. Default is 300.
|
||||
ndp (int): Number of decoder points. Default is 4.
|
||||
nh (int): Number of heads in multi-head attention. Default is 8.
|
||||
ndl (int): Number of decoder layers. Default is 6.
|
||||
d_ffn (int): Dimension of the feed-forward networks. Default is 1024.
|
||||
dropout (float): Dropout rate. Default is 0.
|
||||
act (nn.Module): Activation function. Default is nn.ReLU.
|
||||
eval_idx (int): Evaluation index. Default is -1.
|
||||
nd (int): Number of denoising. Default is 100.
|
||||
label_noise_ratio (float): Label noise ratio. Default is 0.5.
|
||||
box_noise_scale (float): Box noise scale. Default is 1.0.
|
||||
learnt_init_query (bool): Whether to learn initial query embeddings. Default is False.
|
||||
"""
|
||||
super().__init__()
|
||||
self.hidden_dim = hd
|
||||
self.nhead = nh
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue