|
- import argparse
- import torch
- import torch.nn as nn
- from torch.nn import init
- from torchvision import models
- from torch.autograd import Variable
- #from efficientnet_pytorch import EfficientNet
- #import pretrainedmodels
- from torch.nn import functional as F
-
- ######################################################################
- 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') # For old pytorch, you may use kaiming_normal.
- 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)
-
- def fix_relu(m):
- classname = m.__class__.__name__
- if classname.find('ReLU') != -1:
- m.inplace=True
- # Defines the new fc layer and classification layer
- # |--Linear--|--bn--|--relu--|--Linear--|
- class ClassBlock(nn.Module):
- '''
- xjh_notes:
- -----------------------
- Defines the linear layers after the adp avg pooling layer in the backbone
- -----------------------
- Parameters:
- linear - if add a fc layer of size [input_dim, num_bottleneck]
- bnorm - if add a bn layer
- relu - if add a leakyReLU layer
- droprate - if droprate > 0, using dropout in add_block
-
- '''
- def __init__(self, input_dim, class_num, droprate, relu=True, bnorm=True, num_bottleneck=2048, linear=True, return_f = False):
- super(ClassBlock, self).__init__()
- self.return_f = return_f
- add_block = []
- if linear:
- add_block += [nn.Linear(input_dim, num_bottleneck)]
- else:
- num_bottleneck = input_dim
- if bnorm:
- add_block += [nn.BatchNorm1d(num_bottleneck)]
- if relu:
- add_block += [nn.LeakyReLU(0.1)]
- if droprate>0:
- add_block += [nn.Dropout(p=droprate)]
- add_block = nn.Sequential(*add_block)
- add_block.apply(weights_init_kaiming)
-
- classifier = []
- classifier += [nn.Linear(num_bottleneck, class_num)]
- classifier = nn.Sequential(*classifier)
- classifier.apply(weights_init_classifier)
-
- self.add_block = add_block
- self.classifier = classifier
- def forward(self, x):
- f = x
- x = self.add_block(x)
- x = self.classifier(x)
- if self.return_f:
- return x,f
- else:
- return x
-
- # Define the ResNet50-based Model
- class ft_net(nn.Module):
-
- def __init__(self, class_num, droprate=0.5, stride=2, init_model=None, pool='avg', return_f=True):
- self.return_f=return_f
- super(ft_net, self).__init__()
- model_ft = models.resnet50(pretrained=True)
- # avg pooling to global pooling
- if stride == 1:
- model_ft.layer4[0].downsample[0].stride = (1,1)
- model_ft.layer4[0].conv2.stride = (1,1)
-
- self.pool = pool
- if pool =='avg+max':
- model_ft.avgpool2 = nn.AdaptiveAvgPool2d((1,1))
- model_ft.maxpool2 = nn.AdaptiveMaxPool2d((1,1))
- self.model = model_ft
- self.classifier= ClassBlock(2048, class_num, droprate, return_f=self.return_f, linear=False)
- elif pool=='avg':
- model_ft.avgpool = nn.AdaptiveAvgPool2d((1,1))
- self.model = model_ft
- self.classifier= ClassBlock(2048, class_num, droprate, return_f=self.return_f, linear=False)
-
- if init_model!=None:
- self.flag = True
- self.model = init_model.model
- self.pool = init_model.pool
- self.classifier.add_block = init_model.classifier.add_block
- # avg pooling to global pooling
-
- def forward(self, x):
- x = self.model.conv1(x)
- x = self.model.bn1(x)
- x = self.model.relu(x)
- x = self.model.maxpool(x)
- x = self.model.layer1(x)
- x = self.model.layer2(x)
- x = self.model.layer3(x)
- x = self.model.layer4(x)
- if self.pool == 'avg+max':
- x1 = self.model.avgpool2(x)
- x2 = self.model.maxpool2(x)
- x = torch.cat((x1,x2), dim = 1)
- x = x.view(x.size(0), x.size(1))
- elif self.pool == 'avg':
- x = self.model.avgpool(x)
- x = x.view(x.size(0), x.size(1))
-
- output = self.classifier(x) #if return_f: return (fc, feature) else return fc
- return output
-
-
- '''
- # debug model structure
- # Run this code with:
- python model.py
- '''
- if __name__ == '__main__':
- # Here I left a simple forward function.
- # Test the model, before you train it.
-
- '''
- net = CPB(751)
- #net = ft_net_SE(751)
- print(net)
- input = Variable(torch.FloatTensor(4, 3, 320, 320))
- output = net(input)
- print('net output size:')
- #print(output.shape)
- '''
-
- xjh_net = ft_net(333)
- print(xjh_net)
- input = Variable(torch.FloatTensor(2,3,384,384))
- output = xjh_net(input)
- print(output.shape)
|