You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
90 lines
2.2 KiB
90 lines
2.2 KiB
"""squeezenet in pytorch
|
|
"""
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
|
|
class Fire(nn.Module):
|
|
|
|
def __init__(self, in_channel, out_channel, squzee_channel):
|
|
|
|
super().__init__()
|
|
self.squeeze = nn.Sequential(
|
|
nn.Conv2d(in_channel, squzee_channel, 1),
|
|
nn.BatchNorm2d(squzee_channel),
|
|
nn.ReLU(inplace=True)
|
|
)
|
|
|
|
self.expand_1x1 = nn.Sequential(
|
|
nn.Conv2d(squzee_channel, int(out_channel / 2), 1),
|
|
nn.BatchNorm2d(int(out_channel / 2)),
|
|
nn.ReLU(inplace=True)
|
|
)
|
|
|
|
self.expand_3x3 = nn.Sequential(
|
|
nn.Conv2d(squzee_channel, int(out_channel / 2), 3, padding=1),
|
|
nn.BatchNorm2d(int(out_channel / 2)),
|
|
nn.ReLU(inplace=True)
|
|
)
|
|
|
|
def forward(self, x):
|
|
|
|
x = self.squeeze(x)
|
|
x = torch.cat([
|
|
self.expand_1x1(x),
|
|
self.expand_3x3(x)
|
|
], 1)
|
|
|
|
return x
|
|
|
|
class SqueezeNet(nn.Module):
|
|
|
|
"""mobile net with simple bypass"""
|
|
def __init__(self, class_num=100):
|
|
|
|
super().__init__()
|
|
self.stem = nn.Sequential(
|
|
nn.Conv2d(3, 96, 3, padding=1),
|
|
nn.BatchNorm2d(96),
|
|
nn.ReLU(inplace=True),
|
|
nn.MaxPool2d(2, 2)
|
|
)
|
|
|
|
self.fire2 = Fire(96, 128, 16)
|
|
self.fire3 = Fire(128, 128, 16)
|
|
self.fire4 = Fire(128, 256, 32)
|
|
self.fire5 = Fire(256, 256, 32)
|
|
self.fire6 = Fire(256, 384, 48)
|
|
self.fire7 = Fire(384, 384, 48)
|
|
self.fire8 = Fire(384, 512, 64)
|
|
self.fire9 = Fire(512, 512, 64)
|
|
|
|
self.conv10 = nn.Conv2d(512, class_num, 1)
|
|
self.avg = nn.AdaptiveAvgPool2d(1)
|
|
self.maxpool = nn.MaxPool2d(2, 2)
|
|
|
|
def forward(self, x):
|
|
x = self.stem(x)
|
|
|
|
f2 = self.fire2(x)
|
|
f3 = self.fire3(f2) + f2
|
|
f4 = self.fire4(f3)
|
|
f4 = self.maxpool(f4)
|
|
|
|
f5 = self.fire5(f4) + f4
|
|
f6 = self.fire6(f5)
|
|
f7 = self.fire7(f6) + f6
|
|
f8 = self.fire8(f7)
|
|
f8 = self.maxpool(f8)
|
|
|
|
f9 = self.fire9(f8)
|
|
c10 = self.conv10(f9)
|
|
|
|
x = self.avg(c10)
|
|
x = x.view(x.size(0), -1)
|
|
|
|
return x
|
|
|
|
def squeezenet(class_num=1):
|
|
return SqueezeNet(class_num=class_num)
|