Initial commit
This commit is contained in:
xixu-me committed
commit
80fa1d394d
77 files changed
+5568
No files matched your search
@@ -0,0 +1,360 @@
|
||||
import copy
|
||||
import math
|
||||
from collections import OrderedDict
|
||||
from functools import partial
|
||||
from typing import Callable, Optional
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch import Tensor
|
||||
from torch.nn import functional as F
|
||||
|
||||
|
||||
def _make_divisible(ch, divisor=8, min_ch=None):
|
||||
|
||||
if min_ch is None:
|
||||
min_ch = divisor
|
||||
new_ch = max(min_ch, int(ch + divisor / 2) // divisor * divisor)
|
||||
if new_ch < 0.9 * ch:
|
||||
new_ch += divisor
|
||||
return new_ch
|
||||
|
||||
|
||||
def drop_path(x, drop_prob: float = 0.0, training: bool = False):
|
||||
if drop_prob == 0.0 or not training:
|
||||
return x
|
||||
keep_prob = 1 - drop_prob
|
||||
shape = (x.shape[0],) + (1,) * (
|
||||
x.ndim - 1
|
||||
) # work with diff dim tensors, not just 2D ConvNets
|
||||
random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
|
||||
random_tensor.floor_() # binarize
|
||||
output = x.div(keep_prob) * random_tensor
|
||||
return output
|
||||
|
||||
|
||||
class DropPath(nn.Module):
|
||||
def __init__(self, drop_prob=None):
|
||||
super(DropPath, self).__init__()
|
||||
self.drop_prob = drop_prob
|
||||
|
||||
def forward(self, x):
|
||||
return drop_path(x, self.drop_prob, self.training)
|
||||
|
||||
|
||||
class ConvBNActivation(nn.Sequential):
|
||||
def __init__(
|
||||
self,
|
||||
in_planes: int,
|
||||
out_planes: int,
|
||||
kernel_size: int = 3,
|
||||
stride: int = 1,
|
||||
groups: int = 1,
|
||||
norm_layer: Optional[Callable[..., nn.Module]] = None,
|
||||
activation_layer: Optional[Callable[..., nn.Module]] = None,
|
||||
):
|
||||
padding = (kernel_size - 1) // 2
|
||||
if norm_layer is None:
|
||||
norm_layer = nn.BatchNorm2d
|
||||
if activation_layer is None:
|
||||
activation_layer = nn.SiLU # alias Swish (torch>=1.7)
|
||||
super(ConvBNActivation, self).__init__(
|
||||
nn.Conv2d(
|
||||
in_channels=in_planes,
|
||||
out_channels=out_planes,
|
||||
kernel_size=kernel_size,
|
||||
stride=stride,
|
||||
padding=padding,
|
||||
groups=groups,
|
||||
bias=False,
|
||||
),
|
||||
norm_layer(out_planes),
|
||||
activation_layer(),
|
||||
)
|
||||
|
||||
|
||||
class SqueezeExcitation(nn.Module):
|
||||
def __init__(self, input_c: int, expand_c: int, squeeze_factor: int = 4):
|
||||
super(SqueezeExcitation, self).__init__()
|
||||
squeeze_c = input_c // squeeze_factor
|
||||
self.fc1 = nn.Conv2d(expand_c, squeeze_c, 1)
|
||||
self.ac1 = nn.SiLU()
|
||||
self.fc2 = nn.Conv2d(squeeze_c, expand_c, 1)
|
||||
self.ac2 = nn.Sigmoid()
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
scale = F.adaptive_avg_pool2d(x, output_size=(1, 1))
|
||||
scale = self.fc1(scale)
|
||||
scale = self.ac1(scale)
|
||||
scale = self.fc2(scale)
|
||||
scale = self.ac2(scale)
|
||||
return scale * x
|
||||
|
||||
|
||||
class InvertedResidualConfig:
|
||||
def __init__(
|
||||
self,
|
||||
kernel: int,
|
||||
input_c: int,
|
||||
out_c: int,
|
||||
expanded_ratio: int,
|
||||
stride: int,
|
||||
use_se: bool,
|
||||
drop_rate: float,
|
||||
index: str,
|
||||
width_coefficient: float,
|
||||
):
|
||||
self.input_c = self.adjust_channels(input_c, width_coefficient)
|
||||
self.kernel = kernel
|
||||
self.expanded_c = self.input_c * expanded_ratio
|
||||
self.out_c = self.adjust_channels(out_c, width_coefficient)
|
||||
self.use_se = use_se
|
||||
self.stride = stride
|
||||
self.drop_rate = drop_rate
|
||||
self.index = index
|
||||
|
||||
@staticmethod
|
||||
def adjust_channels(channels: int, width_coefficient: float):
|
||||
return _make_divisible(channels * width_coefficient, 8)
|
||||
|
||||
|
||||
class InvertedResidual(nn.Module):
|
||||
def __init__(
|
||||
self, cnf: InvertedResidualConfig, norm_layer: Callable[..., nn.Module]
|
||||
):
|
||||
super(InvertedResidual, self).__init__()
|
||||
if cnf.stride not in [1, 2]:
|
||||
raise ValueError("illegal stride value.")
|
||||
self.use_res_connect = cnf.stride == 1 and cnf.input_c == cnf.out_c
|
||||
layers = OrderedDict()
|
||||
activation_layer = nn.SiLU
|
||||
if cnf.expanded_c != cnf.input_c:
|
||||
layers.update(
|
||||
{
|
||||
"expand_conv": ConvBNActivation(
|
||||
cnf.input_c,
|
||||
cnf.expanded_c,
|
||||
kernel_size=1,
|
||||
norm_layer=norm_layer,
|
||||
activation_layer=activation_layer,
|
||||
)
|
||||
}
|
||||
)
|
||||
layers.update(
|
||||
{
|
||||
"dwconv": ConvBNActivation(
|
||||
cnf.expanded_c,
|
||||
cnf.expanded_c,
|
||||
kernel_size=cnf.kernel,
|
||||
stride=cnf.stride,
|
||||
groups=cnf.expanded_c,
|
||||
norm_layer=norm_layer,
|
||||
activation_layer=activation_layer,
|
||||
)
|
||||
}
|
||||
)
|
||||
if cnf.use_se:
|
||||
layers.update({"se": SqueezeExcitation(cnf.input_c, cnf.expanded_c)})
|
||||
layers.update(
|
||||
{
|
||||
"project_conv": ConvBNActivation(
|
||||
cnf.expanded_c,
|
||||
cnf.out_c,
|
||||
kernel_size=1,
|
||||
norm_layer=norm_layer,
|
||||
activation_layer=nn.Identity,
|
||||
)
|
||||
}
|
||||
)
|
||||
self.block = nn.Sequential(layers)
|
||||
self.out_channels = cnf.out_c
|
||||
self.is_strided = cnf.stride > 1
|
||||
if self.use_res_connect and cnf.drop_rate > 0:
|
||||
self.dropout = DropPath(cnf.drop_rate)
|
||||
else:
|
||||
self.dropout = nn.Identity()
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
result = self.block(x)
|
||||
result = self.dropout(result)
|
||||
if self.use_res_connect:
|
||||
result += x
|
||||
return result
|
||||
|
||||
|
||||
class EfficientNet(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
width_coefficient: float,
|
||||
depth_coefficient: float,
|
||||
num_classes: int = 1000,
|
||||
dropout_rate: float = 0.2,
|
||||
drop_connect_rate: float = 0.2,
|
||||
block: Optional[Callable[..., nn.Module]] = None,
|
||||
norm_layer: Optional[Callable[..., nn.Module]] = None,
|
||||
):
|
||||
super(EfficientNet, self).__init__()
|
||||
default_cnf = [
|
||||
[3, 32, 16, 1, 1, True, drop_connect_rate, 1],
|
||||
[3, 16, 24, 6, 2, True, drop_connect_rate, 2],
|
||||
[5, 24, 40, 6, 2, True, drop_connect_rate, 2],
|
||||
[3, 40, 80, 6, 2, True, drop_connect_rate, 3],
|
||||
[5, 80, 112, 6, 1, True, drop_connect_rate, 3],
|
||||
[5, 112, 192, 6, 2, True, drop_connect_rate, 4],
|
||||
[3, 192, 320, 6, 1, True, drop_connect_rate, 1],
|
||||
]
|
||||
|
||||
def round_repeats(repeats):
|
||||
return int(math.ceil(depth_coefficient * repeats))
|
||||
|
||||
if block is None:
|
||||
block = InvertedResidual
|
||||
if norm_layer is None:
|
||||
norm_layer = partial(nn.BatchNorm2d, eps=1e-3, momentum=0.1)
|
||||
|
||||
adjust_channels = partial(
|
||||
InvertedResidualConfig.adjust_channels, width_coefficient=width_coefficient
|
||||
)
|
||||
bneck_conf = partial(
|
||||
InvertedResidualConfig, width_coefficient=width_coefficient
|
||||
)
|
||||
b = 0
|
||||
num_blocks = float(sum(round_repeats(i[-1]) for i in default_cnf))
|
||||
inverted_residual_setting = []
|
||||
for stage, args in enumerate(default_cnf):
|
||||
cnf = copy.copy(args)
|
||||
for i in range(round_repeats(cnf.pop(-1))):
|
||||
if i > 0:
|
||||
cnf[-3] = 1
|
||||
cnf[1] = cnf[2]
|
||||
cnf[-1] = args[-2] * b / num_blocks
|
||||
index = str(stage + 1) + chr(i + 97)
|
||||
inverted_residual_setting.append(bneck_conf(*cnf, index))
|
||||
b += 1
|
||||
layers = OrderedDict()
|
||||
layers.update(
|
||||
{
|
||||
"stem_conv": ConvBNActivation(
|
||||
in_planes=3,
|
||||
out_planes=adjust_channels(32),
|
||||
kernel_size=3,
|
||||
stride=2,
|
||||
norm_layer=norm_layer,
|
||||
)
|
||||
}
|
||||
)
|
||||
for cnf in inverted_residual_setting:
|
||||
layers.update({cnf.index: block(cnf, norm_layer)})
|
||||
last_conv_input_c = inverted_residual_setting[-1].out_c
|
||||
last_conv_output_c = adjust_channels(1280)
|
||||
layers.update(
|
||||
{
|
||||
"top": ConvBNActivation(
|
||||
in_planes=last_conv_input_c,
|
||||
out_planes=last_conv_output_c,
|
||||
kernel_size=1,
|
||||
norm_layer=norm_layer,
|
||||
)
|
||||
}
|
||||
)
|
||||
self.features = nn.Sequential(layers)
|
||||
self.avgpool = nn.AdaptiveAvgPool2d(1)
|
||||
classifier = []
|
||||
if dropout_rate > 0:
|
||||
classifier.append(nn.Dropout(p=dropout_rate, inplace=True))
|
||||
classifier.append(nn.Linear(last_conv_output_c, num_classes))
|
||||
self.classifier = nn.Sequential(*classifier)
|
||||
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
nn.init.kaiming_normal_(m.weight, mode="fan_out")
|
||||
if m.bias is not None:
|
||||
nn.init.zeros_(m.bias)
|
||||
elif isinstance(m, nn.BatchNorm2d):
|
||||
nn.init.ones_(m.weight)
|
||||
nn.init.zeros_(m.bias)
|
||||
elif isinstance(m, nn.Linear):
|
||||
nn.init.normal_(m.weight, 0, 0.01)
|
||||
nn.init.zeros_(m.bias)
|
||||
|
||||
def _forward_impl(self, x: Tensor) -> Tensor:
|
||||
x = self.features(x)
|
||||
x = self.avgpool(x)
|
||||
x = torch.flatten(x, 1)
|
||||
x = self.classifier(x)
|
||||
return x
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
return self._forward_impl(x)
|
||||
|
||||
|
||||
def efficientnet_b0(num_classes=1000):
|
||||
return EfficientNet(
|
||||
width_coefficient=1.0,
|
||||
depth_coefficient=1.0,
|
||||
dropout_rate=0.2,
|
||||
num_classes=num_classes,
|
||||
)
|
||||
|
||||
|
||||
def efficientnet_b1(num_classes=1000):
|
||||
return EfficientNet(
|
||||
width_coefficient=1.0,
|
||||
depth_coefficient=1.1,
|
||||
dropout_rate=0.2,
|
||||
num_classes=num_classes,
|
||||
)
|
||||
|
||||
|
||||
def efficientnet_b2(num_classes=1000):
|
||||
return EfficientNet(
|
||||
width_coefficient=1.1,
|
||||
depth_coefficient=1.2,
|
||||
dropout_rate=0.3,
|
||||
num_classes=num_classes,
|
||||
)
|
||||
|
||||
|
||||
def efficientnet_b3(num_classes=1000):
|
||||
return EfficientNet(
|
||||
width_coefficient=1.2,
|
||||
depth_coefficient=1.4,
|
||||
dropout_rate=0.3,
|
||||
num_classes=num_classes,
|
||||
)
|
||||
|
||||
|
||||
def efficientnet_b4(num_classes=1000):
|
||||
return EfficientNet(
|
||||
width_coefficient=1.4,
|
||||
depth_coefficient=1.8,
|
||||
dropout_rate=0.4,
|
||||
num_classes=num_classes,
|
||||
)
|
||||
|
||||
|
||||
def efficientnet_b5(num_classes=1000):
|
||||
return EfficientNet(
|
||||
width_coefficient=1.6,
|
||||
depth_coefficient=2.2,
|
||||
dropout_rate=0.4,
|
||||
num_classes=num_classes,
|
||||
)
|
||||
|
||||
|
||||
def efficientnet_b6(num_classes=1000):
|
||||
return EfficientNet(
|
||||
width_coefficient=1.8,
|
||||
depth_coefficient=2.6,
|
||||
dropout_rate=0.5,
|
||||
num_classes=num_classes,
|
||||
)
|
||||
|
||||
|
||||
def efficientnet_b7(num_classes=1000):
|
||||
return EfficientNet(
|
||||
width_coefficient=2.0,
|
||||
depth_coefficient=3.1,
|
||||
dropout_rate=0.5,
|
||||
num_classes=num_classes,
|
||||
)
|
||||
Reference in new issue
Block a user