import torch import torch.nn as nn import torch.nn.functional as F '''模型使用老师提供的示例代码,仅修改了三处版本改动''' def weights_init(m): if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear): nn.init.xavier_uniform_(m.weight.data) nn.init.constant_(m.bias, 0.1) class LossFn: def __init__(self, cls_factor=1, box_factor=1, landmark_factor=1): # loss function self.cls_factor = cls_factor self.box_factor = box_factor self.land_factor = landmark_factor self.loss_cls = nn.BCELoss() # binary cross entropy self.loss_box = nn.MSELoss() # mean square error self.loss_landmark = nn.MSELoss() def cls_loss(self,gt_label,pred_label): pred_label = torch.squeeze(pred_label) gt_label = torch.squeeze(gt_label) # get the mask element which >= 0, only 0 and 1 can effect the detection loss mask = torch.ge(gt_label,0) valid_gt_label = torch.masked_select(gt_label,mask) valid_pred_label = torch.masked_select(pred_label,mask) return self.loss_cls(valid_pred_label,valid_gt_label)*self.cls_factor def box_loss(self,gt_label,gt_offset,pred_offset): pred_offset = torch.squeeze(pred_offset) gt_offset = torch.squeeze(gt_offset) gt_label = torch.squeeze(gt_label) #get the mask element which != 0 unmask = torch.eq(gt_label,0) mask = torch.eq(unmask,0) #convert mask to dim index chose_index = torch.nonzero(mask.data) chose_index = torch.squeeze(chose_index) #only valid element can effect the loss valid_gt_offset = gt_offset[chose_index,:] valid_pred_offset = pred_offset[chose_index,:] return self.loss_box(valid_pred_offset,valid_gt_offset)*self.box_factor def landmark_loss(self,gt_label,gt_landmark,pred_landmark): pred_landmark = torch.squeeze(pred_landmark) gt_landmark = torch.squeeze(gt_landmark) gt_label = torch.squeeze(gt_label) mask = torch.eq(gt_label,-2) chose_index = torch.nonzero(mask.data) chose_index = torch.squeeze(chose_index) valid_gt_landmark = gt_landmark[chose_index, :] valid_pred_landmark = pred_landmark[chose_index, :] return self.loss_landmark(valid_pred_landmark,valid_gt_landmark)*self.land_factor class PNet(nn.Module): ''' PNet ''' def __init__(self, is_train=False, use_cuda=True): super(PNet, self).__init__() self.is_train = is_train self.use_cuda = use_cuda # backend self.pre_layer = nn.Sequential( nn.Conv2d(3, 10, kernel_size=3, stride=1), # conv1 nn.PReLU(), # PReLU1 nn.MaxPool2d(kernel_size=2, stride=2), # pool1 nn.Conv2d(10, 16, kernel_size=3, stride=1), # conv2 nn.PReLU(), # PReLU2 nn.Conv2d(16, 32, kernel_size=3, stride=1), # conv3 nn.PReLU() # PReLU3 ) # detection self.conv4_1 = nn.Conv2d(32, 1, kernel_size=1, stride=1) # bounding box regresion self.conv4_2 = nn.Conv2d(32, 4, kernel_size=1, stride=1) # landmark localization self.conv4_3 = nn.Conv2d(32, 10, kernel_size=1, stride=1) # weight initiation with xavier self.apply(weights_init) def forward(self, x): x = self.pre_layer(x) label = torch.sigmoid(self.conv4_1(x)) offset = self.conv4_2(x) # landmark = self.conv4_3(x) if self.is_train is True: # label_loss = LossUtil.label_loss(self.gt_label,torch.squeeze(label)) # bbox_loss = LossUtil.bbox_loss(self.gt_bbox,torch.squeeze(offset)) return label,offset #landmark = self.conv4_3(x) return label, offset class RNet(nn.Module): ''' RNet ''' def __init__(self,is_train=False, use_cuda=True): super(RNet, self).__init__() self.is_train = is_train self.use_cuda = use_cuda # backend self.pre_layer = nn.Sequential( nn.Conv2d(3, 28, kernel_size=3, stride=1), # conv1 nn.PReLU(), # prelu1 nn.MaxPool2d(kernel_size=3, stride=2), # pool1 nn.Conv2d(28, 48, kernel_size=3, stride=1), # conv2 nn.PReLU(), # prelu2 nn.MaxPool2d(kernel_size=3, stride=2), # pool2 nn.Conv2d(48, 64, kernel_size=2, stride=1), # conv3 nn.PReLU() # prelu3 ) self.conv4 = nn.Linear(64*2*2, 128) # conv4 self.prelu4 = nn.PReLU() # prelu4 # detection self.conv5_1 = nn.Linear(128, 1) # bounding box regression self.conv5_2 = nn.Linear(128, 4) # lanbmark localization self.conv5_3 = nn.Linear(128, 10) # weight initiation weih xavier self.apply(weights_init) def forward(self, x): # backend x = self.pre_layer(x) x = x.view(x.size(0), -1) x = self.conv4(x) x = self.prelu4(x) # detection det = torch.sigmoid(self.conv5_1(x)) box = self.conv5_2(x) # landmark = self.conv5_3(x) if self.is_train is True: return det, box #landmard = self.conv5_3(x) return det, box class ONet(nn.Module): ''' RNet ''' def __init__(self,is_train=False, use_cuda=True): super(ONet, self).__init__() self.is_train = is_train self.use_cuda = use_cuda # backend self.pre_layer = nn.Sequential( nn.Conv2d(3, 32, kernel_size=3, stride=1), # conv1 nn.PReLU(), # prelu1 nn.MaxPool2d(kernel_size=3, stride=2), # pool1 nn.Conv2d(32, 64, kernel_size=3, stride=1), # conv2 nn.PReLU(), # prelu2 nn.MaxPool2d(kernel_size=3, stride=2), # pool2 nn.Conv2d(64, 64, kernel_size=3, stride=1), # conv3 nn.PReLU(), # prelu3 nn.MaxPool2d(kernel_size=2,stride=2), # pool3 nn.Conv2d(64,128,kernel_size=2,stride=1), # conv4 nn.PReLU() # prelu4 ) self.conv5 = nn.Linear(128*2*2, 256) # conv5 self.prelu5 = nn.PReLU() # prelu5 # detection self.conv6_1 = nn.Linear(256, 1) # bounding box regression self.conv6_2 = nn.Linear(256, 4) # lanbmark localization self.conv6_3 = nn.Linear(256, 10) # weight initiation weih xavier self.apply(weights_init) def forward(self, x): # backend x = self.pre_layer(x) x = x.view(x.size(0), -1) x = self.conv5(x) x = self.prelu5(x) # detection det = torch.sigmoid(self.conv6_1(x)) box = self.conv6_2(x) landmark = self.conv6_3(x) if self.is_train is True: return det, box, landmark #landmard = self.conv5_3(x) return det, box, landmark