dk - edit

net/__init__.py

@@ -0,0 +1,9 @@
+from .models import Backbone_nFC, Backbone_nFC_Id
+
+
+def get_model(model_name, num_label, use_id=False, num_id=None):
+    if not use_id:
+        return Backbone_nFC(num_label, model_name)
+    else:
+        return Backbone_nFC_Id(num_label, num_id, model_name)
+

net/models.py

@@ -0,0 +1,76 @@
+import torch
+from torch import nn
+from torch.nn import init
+from torchvision import models
+from net.utils import ClassBlock
+from torch.nn import functional as F
+
+
+class Backbone_nFC(nn.Module):
+    def __init__(self, class_num, model_name='resnet50_nfc'):
+        super(Backbone_nFC, self).__init__()
+        self.model_name = model_name
+        self.backbone_name = model_name.split('_')[0]
+        self.class_num = class_num
+
+        model_ft = getattr(models, self.backbone_name)(pretrained=True)
+        if 'resnet' in self.backbone_name:
+            model_ft.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+            model_ft.fc = nn.Sequential()
+            self.features = model_ft
+            self.num_ftrs = 2048
+        elif 'densenet' in self.backbone_name:
+            model_ft.features.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+            model_ft.fc = nn.Sequential()
+            self.features = model_ft.features
+            self.num_ftrs = 1024
+        else:
+            raise NotImplementedError
+        
+        for c in range(self.class_num):
+            self.__setattr__('class_%d' % c, ClassBlock(input_dim=self.num_ftrs, class_num=1, activ='sigmoid') )
+
+    def forward(self, x):
+        x = self.features(x)
+        x = x.view(x.size(0), -1)
+        pred_label = [self.__getattr__('class_%d' % c)(x) for c in range(self.class_num)]
+        pred_label = torch.cat(pred_label, dim=1)
+        return pred_label
+
+
+class Backbone_nFC_Id(nn.Module):
+    def __init__(self, class_num, id_num, model_name='resnet50_nfc_id'):
+        super(Backbone_nFC_Id, self).__init__()
+        self.model_name = model_name
+        self.backbone_name = model_name.split('_')[0]
+        self.class_num = class_num
+        self.id_num = id_num
+        
+        model_ft = getattr(models, self.backbone_name)(pretrained=True)
+        if 'resnet' in self.backbone_name:
+            model_ft.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+            model_ft.fc = nn.Sequential()
+            self.features = model_ft
+            self.num_ftrs = 2048
+        elif 'densenet' in self.backbone_name:
+            model_ft.features.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+            model_ft.fc = nn.Sequential()
+            self.features = model_ft.features
+            self.num_ftrs = 1024
+        else:
+            raise NotImplementedError
+        
+        for c in range(self.class_num+1):
+            if c == self.class_num:
+                self.__setattr__('class_%d' % c, ClassBlock(self.num_ftrs, class_num=self.id_num, activ='none'))
+            else:
+                self.__setattr__('class_%d' % c, ClassBlock(self.num_ftrs, class_num=1, activ='sigmoid'))
+
+    def forward(self, x):
+        x = self.features(x)
+        x = x.view(x.size(0), -1)
+        pred_label = [self.__getattr__('class_%d' % c)(x) for c in range(self.class_num)]
+        pred_label = torch.cat(pred_label, dim=1)
+        pred_id = self.__getattr__('class_%d' % self.class_num)(x)
+        return pred_label, pred_id
+

net/utils.py

@@ -0,0 +1,60 @@
+from torch import nn
+from torch.nn import init
+
+
+def weights_init_kaiming(m):
+    classname = m.__class__.__name__
+    # print(classname)
+    if classname.find('Conv') != -1:
+        init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
+    elif classname.find('Linear') != -1:
+        init.kaiming_normal_(m.weight.data, a=0, mode='fan_out')
+        init.constant_(m.bias.data, 0.0)
+    elif classname.find('BatchNorm1d') != -1:
+        init.normal_(m.weight.data, 1.0, 0.02)
+        init.constant_(m.bias.data, 0.0)
+
+
+def weights_init_classifier(m):
+    classname = m.__class__.__name__
+    if classname.find('Linear') != -1:
+        init.normal_(m.weight.data, std=0.001)
+        init.constant_(m.bias.data, 0.0)
+
+
+# Defines the new fc layer and classification layer
+# |--Linear--|--bn--|--relu--|--Linear--|
+class ClassBlock(nn.Module):
+    def __init__(self, input_dim, class_num=1, activ='sigmoid', num_bottleneck=512):
+        super(ClassBlock, self).__init__()
+
+        add_block = []
+        add_block += [nn.Linear(input_dim, num_bottleneck)]
+        add_block += [nn.BatchNorm1d(num_bottleneck)]
+        add_block += [nn.LeakyReLU(0.1)]
+        add_block += [nn.Dropout(p=0.5)]
+
+        add_block = nn.Sequential(*add_block)
+        add_block.apply(weights_init_kaiming)
+
+        classifier = []
+        classifier += [nn.Linear(num_bottleneck, class_num)]
+        if activ == 'sigmoid':
+            classifier += [nn.Sigmoid()]
+        elif activ == 'softmax':
+            classifier += [nn.Softmax()]
+        elif activ == 'none':
+            classifier += []
+        else:
+            raise AssertionError("Unsupported activation: {}".format(activ))
+        classifier = nn.Sequential(*classifier)
+        classifier.apply(weights_init_classifier)
+
+        self.add_block = add_block
+        self.classifier = classifier
+
+    def forward(self, x):
+        x = self.add_block(x)
+        x = self.classifier(x)
+        return x
+