import torch.nn as nn
def conv_bn(inp, oup, stride = 1):
return nn.Sequential(
nn.Conv2d(inp, oup, 3, stride, 1, bias=False),
nn.BatchNorm2d(oup),
nn.ReLU6()
)
def conv_dw(inp, oup, stride = 1):
return nn.Sequential(
nn.Conv2d(inp, inp, 3, stride, 1, groups=inp, bias=False),
nn.BatchNorm2d(inp),
nn.ReLU6(),
nn.Conv2d(inp, oup, 1, 1, 0, bias=False),
nn.BatchNorm2d(oup),
nn.ReLU6(),
)
class MobileNetV1(nn.Module):
def __init__(self):
super(MobileNetV1, self).__init__()
self.stage1 = nn.Sequential(
# 160,160,3 -> 80,80,32
conv_bn(3, 32, 2),
# 80,80,32 -> 80,80,64
conv_dw(32, 64, 1),
# 80,80,64 -> 40,40,128
conv_dw(64, 128, 2),
conv_dw(128, 128, 1),
# 40,40,128 -> 20,20,256
conv_dw(128, 256, 2),
conv_dw(256, 256, 1),
)
self.stage2 = nn.Sequential(
# 20,20,256 -> 10,10,512
conv_dw(256, 512, 2),
conv_dw(512, 512, 1),
conv_dw(512, 512, 1),
conv_dw(512, 512, 1),
conv_dw(512, 512, 1),
conv_dw(512, 512, 1),
)
self.stage3 = nn.Sequential(
# 10,10,512 -> 5,5,1024
conv_dw(512, 1024, 2),
conv_dw(1024, 1024, 1),
)
self.avg = nn.AdaptiveAvgPool2d((1,1))
self.fc = nn.Linear(1024, 1000)
def forward(self, x):
x = self.stage1(x)
x = self.stage2(x)
x = self.stage3(x)
x = self.avg(x)
# x = self.model(x)
x = x.view(-1, 1024)
x = self.fc(x)
return x
class Facenet(nn.Module):
def __init__(self, backbone="mobilenet", dropout_keep_prob=0.5, embedding_size=128, num_classes=None, mode="train"):
super(Facenet, self).__init__()
if backbone == "mobilenet":
self.backbone = mobilenet()
flat_shape = 1024
elif backbone == "inception_resnetv1":
self.backbone = inception_resnet()
flat_shape = 1792
else:
raise ValueError('Unsupported backbone - `{}`, Use mobilenet, inception_resnetv1.'.format(backbone))
self.avg = nn.AdaptiveAvgPool2d((1,1))
self.Dropout = nn.Dropout(1 - dropout_keep_prob)
self.Bottleneck = nn.Linear(flat_shape, embedding_size,bias=False)
self.last_bn = nn.BatchNorm1d(embedding_size, eps=0.001, momentum=0.1, affine=True)
if mode == "train":
self.classifier = nn.Linear(embedding_size, num_classes)
def forward(self, x):
x = self.backbone(x)
x = self.avg(x)
x = x.view(x.size(0), -1)
x = self.Dropout(x)
x = self.Bottleneck(x)
x = self.last_bn(x)
x = F.normalize(x, p=2, dim=1)
return x
def forward_feature(self, x):
x = self.backbone(x)
x = self.avg(x)
x = x.view(x.size(0), -1)
x = self.Dropout(x)
x = self.Bottleneck(x)
before_normalize = self.last_bn(x)
x = F.normalize(before_normalize, p=2, dim=1)
return before_normalize, x
def forward_classifier(self, x):
x = self.classifier(x)
return x
在pytorch代码中,只需要一行就可以实现l2标准化的层。
class Facenet(nn.Module):
def __init__(self, backbone="mobilenet", dropout_keep_prob=0.5, embedding_size=128, num_classes=None, mode="train"):
super(Facenet, self).__init__()
if backbone == "mobilenet":
self.backbone = mobilenet()
flat_shape = 1024
elif backbone == "inception_resnetv1":
self.backbone = inception_resnet()
flat_shape = 1792
else:
raise ValueError('Unsupported backbone - `{}`, Use mobilenet, inception_resnetv1.'.format(backbone))
self.avg = nn.AdaptiveAvgPool2d((1,1))
self.Dropout = nn.Dropout(1 - dropout_keep_prob)
self.Bottleneck = nn.Linear(flat_shape, embedding_size,bias=False)
self.last_bn = nn.BatchNorm1d(embedding_size, eps=0.001, momentum=0.1, affine=True)
if mode == "train":
self.classifier = nn.Linear(embedding_size, num_classes)
def forward(self, x):
x = self.backbone(x)
x = self.avg(x)
x = x.view(x.size(0), -1)
x = self.Dropout(x)
x = self.Bottleneck(x)
x = self.last_bn(x)
x = F.normalize(x, p=2, dim=1)
return x
def forward_feature(self, x):
x = self.backbone(x)
x = self.avg(x)
x = x.view(x.size(0), -1)
x = self.Dropout(x)
x = self.Bottleneck(x)
before_normalize = self.last_bn(x)
x = F.normalize(before_normalize, p=2, dim=1)
return before_normalize, x
def forward_classifier(self, x):
x = self.classifier(x)
return x
class Facenet(nn.Module):
def __init__(self, backbone="mobilenet", dropout_keep_prob=0.5, embedding_size=128, num_classes=None, mode="train"):
super(Facenet, self).__init__()
if backbone == "mobilenet":
self.backbone = mobilenet()
flat_shape = 1024
elif backbone == "inception_resnetv1":
self.backbone = inception_resnet()
flat_shape = 1792
else:
raise ValueError('Unsupported backbone - `{}`, Use mobilenet, inception_resnetv1.'.format(backbone))
self.avg = nn.AdaptiveAvgPool2d((1,1))
self.Dropout = nn.Dropout(1 - dropout_keep_prob)
self.Bottleneck = nn.Linear(flat_shape, embedding_size,bias=False)
self.last_bn = nn.BatchNorm1d(embedding_size, eps=0.001, momentum=0.1, affine=True)
if mode == "train":
self.classifier = nn.Linear(embedding_size, num_classes)
def forward(self, x):
x = self.backbone(x)
x = self.avg(x)
x = x.view(x.size(0), -1)
x = self.Dropout(x)
x = self.Bottleneck(x)
x = self.last_bn(x)
x = F.normalize(x, p=2, dim=1)
return x
def forward_feature(self, x):
x = self.backbone(x)
x = self.avg(x)
x = x.view(x.size(0), -1)
x = self.Dropout(x)
x = self.Bottleneck(x)
before_normalize = self.last_bn(x)
x = F.normalize(before_normalize, p=2, dim=1)
return before_normalize, x
def forward_classifier(self, x):
x = self.classifier(x)
return x