| 123456789101112131415161718192021222324252627282930313233343536373839 |
- import torch
- import torch.nn as nn
- from .rtcdet_basic import BasicConv
- # -------------- Neck network --------------
- 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.input_proj = BasicConv(in_dim, inter_dim, kernel_size=1,
- act_type=cfg.neck_act, norm_type=cfg.neck_norm)
- self.output_proj = BasicConv(inter_dim * 4, out_dim, kernel_size=1,
- act_type=cfg.neck_act, norm_type=cfg.neck_norm)
- self.module = nn.MaxPool2d(cfg.spp_pooling_size, stride=1, padding=cfg.spp_pooling_size//2)
- # Initialize all layers
- self.init_weights()
-
- def init_weights(self):
- """Initialize the parameters."""
- for m in self.modules():
- if isinstance(m, torch.nn.Conv2d):
- m.reset_parameters()
- def forward(self, x):
- x = self.input_proj(x)
- y1 = self.module(x)
- y2 = self.module(y1)
- return self.output_proj(torch.cat((x, y1, y2, self.module(y2)), 1))
-
|
