【光電智造】YOLOv12入門教程

class?AAttn(nn.Module):? ??"""? ? Area-attention module with the requirement of flash attention.? ? Attributes:? ? ? ? dim (int): Number of hidden channels;? ? ? ? num_heads (int): Number of heads into which the attention mechanism is divided;? ? ? ? area (int, optional): Number of areas the feature map is divided. Defaults to 1.? ? Methods:? ? ? ? forward: Performs a forward process of input tensor and outputs a tensor after the execution of the area attention mechanism.? ? Examples:? ? ? ? >>> import torch? ? ? ? >>> from ultralytics.nn.modules import AAttn? ? ? ? >>> model = AAttn(dim=64, num_heads=2, area=4)? ? ? ? >>> x = torch.randn(2, 64, 128, 128)? ? ? ? >>> output = model(x)? ? ? ? >>> print(output.shape)
? ? Notes:?? ? ? ? recommend that dim//num_heads be a multiple of 32 or 64.? ? """
? ??def?__init__(self, dim, num_heads, area=1):? ? ? ??"""Initializes the area-attention module, a simple yet efficient attention module for YOLO."""? ? ? ??super().__init__()? ? ? ? self.area = area
? ? ? ? self.num_heads = num_heads? ? ? ? self.head_dim = head_dim = dim // num_heads? ? ? ? all_head_dim = head_dim * self.num_heads
? ? ? ? self.qkv = Conv(dim, all_head_dim *?3,?1, act=False)? ? ? ? self.proj = Conv(all_head_dim, dim,?1, act=False)? ? ? ? self.pe = Conv(all_head_dim, dim,?7,?1,?3, g=dim, act=False)

? ??def?forward(self, x):? ? ? ??"""Processes the input tensor 'x' through the area-attention"""? ? ? ? B, C, H, W = x.shape? ? ? ? N = H * W
? ? ? ? qkv = self.qkv(x).flatten(2).transpose(1,?2)? ? ? ??if?self.area >?1:? ? ? ? ? ? qkv = qkv.reshape(B * self.area, N // self.area, C *?3)? ? ? ? ? ? B, N, _ = qkv.shape? ? ? ? q, k, v = qkv.view(B, N, self.num_heads, self.head_dim *?3).split(? ? ? ? ? ? [self.head_dim, self.head_dim, self.head_dim], dim=3? ? ? ? )
? ? ? ??# if x.is_cuda:? ? ? ??# ? ? x = flash_attn_func(? ? ? ??# ? ? ? ? q.contiguous().half(),? ? ? ??# ? ? ? ? k.contiguous().half(),? ? ? ??# ? ? ? ? v.contiguous().half()? ? ? ??# ? ? ).to(q.dtype)? ? ? ??# else:? ? ? ? q = q.permute(0,?2,?3,?1)? ? ? ? k = k.permute(0,?2,?3,?1)? ? ? ? v = v.permute(0,?2,?3,?1)? ? ? ? attn = (q.transpose(-2, -1) @ k) * (self.head_dim ** -0.5)? ? ? ? max_attn = attn.max(dim=-1, keepdim=True).values? ? ? ? exp_attn = torch.exp(attn - max_attn)? ? ? ? attn = exp_attn / exp_attn.sum(dim=-1, keepdim=True)? ? ? ? x = (v @ attn.transpose(-2, -1))? ? ? ? x = x.permute(0,?3,?1,?2)? ? ? ? v = v.permute(0,?3,?1,?2)
? ? ? ??if?self.area >?1:? ? ? ? ? ? x = x.reshape(B // self.area, N * self.area, C)? ? ? ? ? ? v = v.reshape(B // self.area, N * self.area, C)? ? ? ? ? ? B, N, _ = x.shape
? ? ? ? x = x.reshape(B, H, W, C).permute(0,?3,?1,?2)? ? ? ? v = v.reshape(B, H, W, C).permute(0,?3,?1,?2)
? ? ? ? x = x + self.pe(v)? ? ? ? x = self.proj(x)? ? ? ??return?x

class?AAttn(nn.Module):? ??"""? ? Area-attention module with the requirement of flash attention.? ? Attributes:? ? ? ? dim (int): Number of hidden channels;? ? ? ? num_heads (int): Number of heads into which the attention mechanism is divided;? ? ? ? area (int, optional): Number of areas the feature map is divided. Defaults to 1.? ? Methods:? ? ? ? forward: Performs a forward process of input tensor and outputs a tensor after the execution of the area attention mechanism.? ? Examples:? ? ? ? >>> import torch? ? ? ? >>> from ultralytics.nn.modules import AAttn? ? ? ? >>> model = AAttn(dim=64, num_heads=2, area=4)? ? ? ? >>> x = torch.randn(2, 64, 128, 128)? ? ? ? >>> output = model(x)? ? ? ? >>> print(output.shape)
? ? Notes:?? ? ? ? recommend that dim//num_heads be a multiple of 32 or 64.? ? """
? ??def?__init__(self, dim, num_heads, area=1):? ? ? ??"""Initializes the area-attention module, a simple yet efficient attention module for YOLO."""? ? ? ??super().__init__()? ? ? ? self.area = area
? ? ? ? self.num_heads = num_heads? ? ? ? self.head_dim = head_dim = dim // num_heads? ? ? ? all_head_dim = head_dim * self.num_heads
? ? ? ? self.qkv = Conv(dim, all_head_dim *?3,?1, act=False)? ? ? ? self.proj = Conv(all_head_dim, dim,?1, act=False)? ? ? ? self.pe = Conv(all_head_dim, dim,?7,?1,?3, g=dim, act=False)

? ??def?forward(self, x):? ? ? ??"""Processes the input tensor 'x' through the area-attention"""? ? ? ? B, C, H, W = x.shape? ? ? ? N = H * W
? ? ? ? qkv = self.qkv(x).flatten(2).transpose(1,?2)? ? ? ??if?self.area >?1:? ? ? ? ? ? qkv = qkv.reshape(B * self.area, N // self.area, C *?3)? ? ? ? ? ? B, N, _ = qkv.shape? ? ? ? q, k, v = qkv.view(B, N, self.num_heads, self.head_dim *?3).split(? ? ? ? ? ? [self.head_dim, self.head_dim, self.head_dim], dim=3? ? ? ? )
? ? ? ??# if x.is_cuda:? ? ? ??# ? ? x = flash_attn_func(? ? ? ??# ? ? ? ? q.contiguous().half(),? ? ? ??# ? ? ? ? k.contiguous().half(),? ? ? ??# ? ? ? ? v.contiguous().half()? ? ? ??# ? ? ).to(q.dtype)? ? ? ??# else:? ? ? ? q = q.permute(0,?2,?3,?1)? ? ? ? k = k.permute(0,?2,?3,?1)? ? ? ? v = v.permute(0,?2,?3,?1)? ? ? ? attn = (q.transpose(-2, -1) @ k) * (self.head_dim ** -0.5)? ? ? ? max_attn = attn.max(dim=-1, keepdim=True).values? ? ? ? exp_attn = torch.exp(attn - max_attn)? ? ? ? attn = exp_attn / exp_attn.sum(dim=-1, keepdim=True)? ? ? ? x = (v @ attn.transpose(-2, -1))? ? ? ? x = x.permute(0,?3,?1,?2)? ? ? ? v = v.permute(0,?3,?1,?2)
? ? ? ??if?self.area >?1:? ? ? ? ? ? x = x.reshape(B // self.area, N * self.area, C)? ? ? ? ? ? v = v.reshape(B // self.area, N * self.area, C)? ? ? ? ? ? B, N, _ = x.shape
? ? ? ? x = x.reshape(B, H, W, C).permute(0,?3,?1,?2)? ? ? ? v = v.reshape(B, H, W, C).permute(0,?3,?1,?2)
? ? ? ? x = x + self.pe(v)? ? ? ? x = self.proj(x)? ? ? ??return?x

class?ABlock(nn.Module):? ??"""? ? ABlock class implementing a Area-Attention block with effective feature extraction.? ? This class encapsulates the functionality for applying multi-head attention with feature map are dividing into areas? ? and feed-forward neural network layers.? ? Attributes:? ? ? ? dim (int): Number of hidden channels;? ? ? ? num_heads (int): Number of heads into which the attention mechanism is divided;? ? ? ? mlp_ratio (float, optional): MLP expansion ratio (or MLP hidden dimension ratio). Defaults to 1.2;? ? ? ? area (int, optional): Number of areas the feature map is divided. ?Defaults to 1.? ? Methods:? ? ? ? forward: Performs a forward pass through the ABlock, applying area-attention and feed-forward layers.? ? Examples:? ? ? ? Create a ABlock and perform a forward pass? ? ? ? >>> model = ABlock(dim=64, num_heads=2, mlp_ratio=1.2, area=4)? ? ? ? >>> x = torch.randn(2, 64, 128, 128)? ? ? ? >>> output = model(x)? ? ? ? >>> print(output.shape)
? ? Notes:?? ? ? ? recommend that dim//num_heads be a multiple of 32 or 64.? ? """
? ??def?__init__(self, dim, num_heads, mlp_ratio=1.2, area=1):? ? ? ??"""Initializes the ABlock with area-attention and feed-forward layers for faster feature extraction."""? ? ? ??super().__init__()
? ? ? ? self.attn = AAttn(dim, num_heads=num_heads, area=area)? ? ? ? mlp_hidden_dim =?int(dim * mlp_ratio)? ? ? ? self.mlp = nn.Sequential(Conv(dim, mlp_hidden_dim,?1), Conv(mlp_hidden_dim, dim,?1, act=False))
? ? ? ? self.apply(self._init_weights)
? ??def?_init_weights(self, m):? ? ? ??"""Initialize weights using a truncated normal distribution."""? ? ? ??if?isinstance(m, nn.Conv2d):? ? ? ? ? ? trunc_normal_(m.weight, std=.02)? ? ? ? ? ??if?isinstance(m, nn.Conv2d)?and?m.bias?is?not?None:? ? ? ? ? ? ? ? nn.init.constant_(m.bias,?0)
? ??def?forward(self, x):? ? ? ??"""Executes a forward pass through ABlock, applying area-attention and feed-forward layers to the input tensor."""? ? ? ? x = x + self.attn(x)? ? ? ? x = x + self.mlp(x)? ? ? ??return?x

class?A2C2f(nn.Module): ?? ??"""? ? A2C2f module with residual enhanced feature extraction using ABlock blocks with area-attention. Also known as R-ELAN? ? This class extends the C2f module by incorporating ABlock blocks for fast attention mechanisms and feature extraction.? ? Attributes:? ? ? ? c1 (int): Number of input channels;? ? ? ? c2 (int): Number of output channels;? ? ? ? n (int, optional): Number of 2xABlock modules to stack. Defaults to 1;? ? ? ? a2 (bool, optional): Whether use area-attention. Defaults to True;? ? ? ? area (int, optional): Number of areas the feature map is divided. Defaults to 1;? ? ? ? residual (bool, optional): Whether use the residual (with layer scale). Defaults to False;? ? ? ? mlp_ratio (float, optional): MLP expansion ratio (or MLP hidden dimension ratio). Defaults to 1.2;? ? ? ? e (float, optional): Expansion ratio for R-ELAN modules. Defaults to 0.5.? ? ? ? g (int, optional): Number of groups for grouped convolution. Defaults to 1;? ? ? ? shortcut (bool, optional): Whether to use shortcut connection. Defaults to True;? ? Methods:? ? ? ? forward: Performs a forward pass through the A2C2f module.? ? Examples:? ? ? ? >>> import torch? ? ? ? >>> from ultralytics.nn.modules import A2C2f? ? ? ? >>> model = A2C2f(c1=64, c2=64, n=2, a2=True, area=4, residual=True, e=0.5)? ? ? ? >>> x = torch.randn(2, 64, 128, 128)? ? ? ? >>> output = model(x)? ? ? ? >>> print(output.shape)? ? """
? ??def?__init__(self, c1, c2, n=1, a2=True, area=1, residual=False, mlp_ratio=2.0, e=0.5, g=1, shortcut=True):? ? ? ??super().__init__()? ? ? ? c_ =?int(c2 * e) ?# hidden channels? ? ? ??assert?c_ %?32?==?0,?"Dimension of ABlock be a multiple of 32."
? ? ? ??# num_heads = c_ // 64 if c_ // 64 >= 2 else c_ // 32? ? ? ? num_heads = c_ //?32
? ? ? ? self.cv1 = Conv(c1, c_,?1,?1)? ? ? ? self.cv2 = Conv((1?+ n) * c_, c2,?1) ?# optional act=FReLU(c2)
? ? ? ? init_values =?0.01??# or smaller? ? ? ? self.gamma = nn.Parameter(init_values * torch.ones((c2)), requires_grad=True)?if?a2?and?residual?else?None
? ? ? ? self.m = nn.ModuleList(? ? ? ? ? ? nn.Sequential(*(ABlock(c_, num_heads, mlp_ratio, area)?for?_?in?range(2)))?if?a2?else?C3k(c_, c_,?2, shortcut, g)?for?_?in?range(n)? ? ? ? )
? ??def?forward(self, x):? ? ? ??"""Forward pass through R-ELAN layer."""? ? ? ? y = [self.cv1(x)]? ? ? ? y.extend(m(y[-1])?for?m?in?self.m)? ? ? ??if?self.gamma?is?not?None:? ? ? ? ? ??return?x + (self.gamma * self.cv2(torch.cat(y,?1)).permute(0,?2,?3,?1)).permute(0,?3,?1,?2)? ? ? ??return?self.cv2(torch.cat(y,?1))

?# dataset pathtrain: ./images/trainval: ./images/testtest: ./images/test
# number of classesnc: 15
# class namesnames: ['car',?'Truck',?'Van',?'Long Vehicle','Bus',?'Airliner',?'Propeller Aircraft',?'Trainer Aircraft',?'Chartered Aircraft',?'Fighter Aircraft',\? ? ? ??'Others',?'Stair Truck',?'Pushback Truck',?'Helicopter',?'Boat']

from?ultralytics.models?import?YOLOimport?osos.environ['KMP_DUPLICATE_LIB_OK'] =?'True'
if?__name__ ==?'__main__':? ? model = YOLO(model='ultralytics/cfg/models/11/yolo11.yaml')? ??# model.load('yolov8n.pt')? ? model.train(data='./data.yaml', epochs=2, batch=1, device='0', imgsz=640, workers=2, cache=False,? ? ? ? ? ? ? ? amp=True, mosaic=False, project='runs/train', name='exp')