| 123456789101112131415161718192021222324252627282930313233343536 |
- import torch
- import torch.nn as nn
- try:
- from .yolov6_basic import BasicConv
- except:
- from yolov6_basic import BasicConv
-
- # Spatial Pyramid Pooling - Fast (SPPF) layer for YOLOv5 by Glenn Jocher
- class SPPF(nn.Module):
- """
- This code referenced to https://github.com/ultralytics/yolov5
- """
- def __init__(self, cfg, in_dim, out_dim):
- super().__init__()
- ## ----------- Basic Parameters -----------
- inter_dim = round(in_dim * cfg.neck_expand_ratio)
- self.out_dim = out_dim
- ## ----------- Network Parameters -----------
- self.cv1 = BasicConv(in_dim, inter_dim,
- kernel_size=1, padding=0, stride=1,
- act_type=cfg.neck_act, norm_type=cfg.neck_norm)
- self.cv2 = BasicConv(inter_dim * 4, out_dim,
- kernel_size=1, padding=0, stride=1,
- act_type=cfg.neck_act, norm_type=cfg.neck_norm)
- self.m = nn.MaxPool2d(kernel_size=cfg.spp_pooling_size,
- stride=1,
- padding=cfg.spp_pooling_size // 2)
- def forward(self, x):
- x = self.cv1(x)
- y1 = self.m(x)
- y2 = self.m(y1)
- return self.cv2(torch.cat((x, y1, y2, self.m(y2)), 1))
|
