File size: 2,414 Bytes
3ed0796 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 |
# Ultralytics YOLO 🚀, AGPL-3.0 license
"""
Convolution modules
"""
import math
import numpy as np
import torch
import torch.nn as nn
__all__ = ('Conv', 'LightConv', 'DWConv', 'DWConvTranspose2d', 'ConvTranspose', 'Focus', 'GhostConv',
'ChannelAttention', 'SpatialAttention', 'CBAM', 'Concat', 'RepConv')
def autopad(k, p=None, d=1): # kernel, padding, dilation
"""Pad to 'same' shape outputs."""
if d > 1:
k = d * (k - 1) + 1 if isinstance(k, int) else [d * (x - 1) + 1 for x in k] # actual kernel-size
if p is None:
p = k // 2 if isinstance(k, int) else [x // 2 for x in k] # auto-pad
return p
# Pavlo's implementation with switch to deploy
class Conv(nn.Module):
default_act = nn.SiLU() # default activation
def __init__(self, a, b, kernel_size=1, stride=1, padding=None, g=1, dilation=1, bn_weight_init=1, bias=False, act=True):
super().__init__()
self.conv = torch.nn.Conv2d(a, b, kernel_size, stride, autopad(kernel_size, padding, dilation), dilation, g, bias=False)
if 1:
self.bn = torch.nn.BatchNorm2d(b)
torch.nn.init.constant_(self.bn.weight, bn_weight_init)
torch.nn.init.constant_(self.bn.bias, 0)
self.act = self.default_act if act is True else act if isinstance(act, nn.Module) else nn.Identity()
def forward(self,x):
x = self.conv(x)
x = self.bn(x)
x = self.act(x)
return x
@torch.no_grad()
def switch_to_deploy(self):
if not isinstance(self.bn, nn.Identity):
# return 1
c, bn = self.conv, self.bn
w = bn.weight / (bn.running_var + bn.eps) ** 0.5
w = c.weight * w[:, None, None, None]
b = bn.bias - bn.running_mean * bn.weight / \
(bn.running_var + bn.eps)**0.5
# m = torch.nn.Conv2d(w.size(1) * c.groups,
# w.size(0),
# w.shape[2:],
# stride=c.stride,
# padding=c.padding,
# dilation=c.dilation,
# groups=c.groups)
self.conv.weight.data.copy_(w)
self.conv.bias = nn.Parameter(b)
# self.conv.bias.data.copy_(b)
# self.conv = m.to(c.weight.device)
self.bn = nn.Identity()
|