mirror of
https://github.com/deepseek-ai/DreamCraft3D.git
synced 2025-02-23 14:28:55 -05:00
254 lines
7.9 KiB
Python
254 lines
7.9 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
__all__ = ["MobileNetV3", "mobilenetv3"]
|
|
|
|
|
|
def conv_bn(
|
|
inp,
|
|
oup,
|
|
stride,
|
|
conv_layer=nn.Conv2d,
|
|
norm_layer=nn.BatchNorm2d,
|
|
nlin_layer=nn.ReLU,
|
|
):
|
|
return nn.Sequential(
|
|
conv_layer(inp, oup, 3, stride, 1, bias=False),
|
|
norm_layer(oup),
|
|
nlin_layer(inplace=True),
|
|
)
|
|
|
|
|
|
def conv_1x1_bn(
|
|
inp, oup, conv_layer=nn.Conv2d, norm_layer=nn.BatchNorm2d, nlin_layer=nn.ReLU
|
|
):
|
|
return nn.Sequential(
|
|
conv_layer(inp, oup, 1, 1, 0, bias=False),
|
|
norm_layer(oup),
|
|
nlin_layer(inplace=True),
|
|
)
|
|
|
|
|
|
class Hswish(nn.Module):
|
|
def __init__(self, inplace=True):
|
|
super(Hswish, self).__init__()
|
|
self.inplace = inplace
|
|
|
|
def forward(self, x):
|
|
return x * F.relu6(x + 3.0, inplace=self.inplace) / 6.0
|
|
|
|
|
|
class Hsigmoid(nn.Module):
|
|
def __init__(self, inplace=True):
|
|
super(Hsigmoid, self).__init__()
|
|
self.inplace = inplace
|
|
|
|
def forward(self, x):
|
|
return F.relu6(x + 3.0, inplace=self.inplace) / 6.0
|
|
|
|
|
|
class SEModule(nn.Module):
|
|
def __init__(self, channel, reduction=4):
|
|
super(SEModule, self).__init__()
|
|
self.avg_pool = nn.AdaptiveAvgPool2d(1)
|
|
self.fc = nn.Sequential(
|
|
nn.Linear(channel, channel // reduction, bias=False),
|
|
nn.ReLU(inplace=True),
|
|
nn.Linear(channel // reduction, channel, bias=False),
|
|
Hsigmoid()
|
|
# nn.Sigmoid()
|
|
)
|
|
|
|
def forward(self, x):
|
|
b, c, _, _ = x.size()
|
|
y = self.avg_pool(x).view(b, c)
|
|
y = self.fc(y).view(b, c, 1, 1)
|
|
return x * y.expand_as(x)
|
|
|
|
|
|
class Identity(nn.Module):
|
|
def __init__(self, channel):
|
|
super(Identity, self).__init__()
|
|
|
|
def forward(self, x):
|
|
return x
|
|
|
|
|
|
def make_divisible(x, divisible_by=8):
|
|
import numpy as np
|
|
|
|
return int(np.ceil(x * 1.0 / divisible_by) * divisible_by)
|
|
|
|
|
|
class MobileBottleneck(nn.Module):
|
|
def __init__(self, inp, oup, kernel, stride, exp, se=False, nl="RE"):
|
|
super(MobileBottleneck, self).__init__()
|
|
assert stride in [1, 2]
|
|
assert kernel in [3, 5]
|
|
padding = (kernel - 1) // 2
|
|
self.use_res_connect = stride == 1 and inp == oup
|
|
|
|
conv_layer = nn.Conv2d
|
|
norm_layer = nn.BatchNorm2d
|
|
if nl == "RE":
|
|
nlin_layer = nn.ReLU # or ReLU6
|
|
elif nl == "HS":
|
|
nlin_layer = Hswish
|
|
else:
|
|
raise NotImplementedError
|
|
if se:
|
|
SELayer = SEModule
|
|
else:
|
|
SELayer = Identity
|
|
|
|
self.conv = nn.Sequential(
|
|
# pw
|
|
conv_layer(inp, exp, 1, 1, 0, bias=False),
|
|
norm_layer(exp),
|
|
nlin_layer(inplace=True),
|
|
# dw
|
|
conv_layer(exp, exp, kernel, stride, padding, groups=exp, bias=False),
|
|
norm_layer(exp),
|
|
SELayer(exp),
|
|
nlin_layer(inplace=True),
|
|
# pw-linear
|
|
conv_layer(exp, oup, 1, 1, 0, bias=False),
|
|
norm_layer(oup),
|
|
)
|
|
|
|
def forward(self, x):
|
|
if self.use_res_connect:
|
|
return x + self.conv(x)
|
|
else:
|
|
return self.conv(x)
|
|
|
|
|
|
class MobileNetV3(nn.Module):
|
|
def __init__(
|
|
self, n_class=1000, input_size=224, dropout=0.0, mode="small", width_mult=1.0
|
|
):
|
|
super(MobileNetV3, self).__init__()
|
|
input_channel = 16
|
|
last_channel = 1280
|
|
if mode == "large":
|
|
# refer to Table 1 in paper
|
|
mobile_setting = [
|
|
# k, exp, c, se, nl, s,
|
|
[3, 16, 16, False, "RE", 1],
|
|
[3, 64, 24, False, "RE", 2],
|
|
[3, 72, 24, False, "RE", 1],
|
|
[5, 72, 40, True, "RE", 2],
|
|
[5, 120, 40, True, "RE", 1],
|
|
[5, 120, 40, True, "RE", 1],
|
|
[3, 240, 80, False, "HS", 2],
|
|
[3, 200, 80, False, "HS", 1],
|
|
[3, 184, 80, False, "HS", 1],
|
|
[3, 184, 80, False, "HS", 1],
|
|
[3, 480, 112, True, "HS", 1],
|
|
[3, 672, 112, True, "HS", 1],
|
|
[5, 672, 160, True, "HS", 2],
|
|
[5, 960, 160, True, "HS", 1],
|
|
[5, 960, 160, True, "HS", 1],
|
|
]
|
|
elif mode == "small":
|
|
# refer to Table 2 in paper
|
|
mobile_setting = [
|
|
# k, exp, c, se, nl, s,
|
|
[3, 16, 16, True, "RE", 2],
|
|
[3, 72, 24, False, "RE", 2],
|
|
[3, 88, 24, False, "RE", 1],
|
|
[5, 96, 40, True, "HS", 2],
|
|
[5, 240, 40, True, "HS", 1],
|
|
[5, 240, 40, True, "HS", 1],
|
|
[5, 120, 48, True, "HS", 1],
|
|
[5, 144, 48, True, "HS", 1],
|
|
[5, 288, 96, True, "HS", 2],
|
|
[5, 576, 96, True, "HS", 1],
|
|
[5, 576, 96, True, "HS", 1],
|
|
]
|
|
else:
|
|
raise NotImplementedError
|
|
|
|
# building first layer
|
|
assert input_size % 32 == 0
|
|
last_channel = (
|
|
make_divisible(last_channel * width_mult)
|
|
if width_mult > 1.0
|
|
else last_channel
|
|
)
|
|
self.features = [conv_bn(3, input_channel, 2, nlin_layer=Hswish)]
|
|
self.classifier = []
|
|
|
|
# building mobile blocks
|
|
for k, exp, c, se, nl, s in mobile_setting:
|
|
output_channel = make_divisible(c * width_mult)
|
|
exp_channel = make_divisible(exp * width_mult)
|
|
self.features.append(
|
|
MobileBottleneck(
|
|
input_channel, output_channel, k, s, exp_channel, se, nl
|
|
)
|
|
)
|
|
input_channel = output_channel
|
|
|
|
# building last several layers
|
|
if mode == "large":
|
|
last_conv = make_divisible(960 * width_mult)
|
|
self.features.append(
|
|
conv_1x1_bn(input_channel, last_conv, nlin_layer=Hswish)
|
|
)
|
|
self.features.append(nn.AdaptiveAvgPool2d(1))
|
|
self.features.append(nn.Conv2d(last_conv, last_channel, 1, 1, 0))
|
|
self.features.append(Hswish(inplace=True))
|
|
elif mode == "small":
|
|
last_conv = make_divisible(576 * width_mult)
|
|
self.features.append(
|
|
conv_1x1_bn(input_channel, last_conv, nlin_layer=Hswish)
|
|
)
|
|
# self.features.append(SEModule(last_conv)) # refer to paper Table2, but I think this is a mistake
|
|
self.features.append(nn.AdaptiveAvgPool2d(1))
|
|
self.features.append(nn.Conv2d(last_conv, last_channel, 1, 1, 0))
|
|
self.features.append(Hswish(inplace=True))
|
|
else:
|
|
raise NotImplementedError
|
|
|
|
# make it nn.Sequential
|
|
self.features = nn.Sequential(*self.features)
|
|
|
|
# building classifier
|
|
self.classifier = nn.Sequential(
|
|
nn.Dropout(p=dropout), # refer to paper section 6
|
|
nn.Linear(last_channel, n_class),
|
|
)
|
|
|
|
self._initialize_weights()
|
|
|
|
def forward(self, x):
|
|
x = self.features(x)
|
|
x = x.mean(3).mean(2)
|
|
x = self.classifier(x)
|
|
return x
|
|
|
|
def _initialize_weights(self):
|
|
# weight initialization
|
|
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)
|
|
if m.bias is not None:
|
|
nn.init.zeros_(m.bias)
|
|
|
|
|
|
def mobilenetv3(pretrained=False, **kwargs):
|
|
model = MobileNetV3(**kwargs)
|
|
if pretrained:
|
|
state_dict = torch.load("mobilenetv3_small_67.4.pth.tar")
|
|
model.load_state_dict(state_dict, strict=True)
|
|
# raise NotImplementedError
|
|
return model |