def create_mtcnn_net(p_model_path=None,
r_model_path=None,
o_model_path=None,
use_cuda=True):
pnet, rnet, onet = None, None, None
if p_model_path is not None:
pnet = PNet(use_cuda=use_cuda)
pnet.load_state_dict(torch.load(p_model_path))
if (use_cuda):
pnet.cuda()
pnet.eval()
if r_model_path is not None:
rnet = RNet(use_cuda=use_cuda)
rnet.load_state_dict(torch.load(r_model_path))
if (use_cuda):
rnet.cuda()
rnet.eval()
if o_model_path is not None:
onet = ONet(use_cuda=use_cuda)
onet.load_state_dict(torch.load(o_model_path))
if (use_cuda):
onet.cuda()
onet.eval()
return pnet, rnet, onet
def create_mtcnn_net(p_model_path=None,
r_model_path=None,
o_model_path=None,
use_cuda=True):
pnet, rnet, onet = None, None, None
if p_model_path is not None:
pnet = PNet(use_cuda=use_cuda)
if (use_cuda):
print('p_model_path:{0}'.format(p_model_path))
pnet.load_state_dict(torch.load(p_model_path))
pnet.cuda()
else:
# forcing all GPU tensors to be in CPU while loading
pnet.load_state_dict(
torch.load(p_model_path,
map_location=lambda storage, loc: storage))
pnet.eval()
if r_model_path is not None:
rnet = RNet(use_cuda=use_cuda)
if (use_cuda):
print('r_model_path:{0}'.format(r_model_path))
rnet.load_state_dict(torch.load(r_model_path))
rnet.cuda()
else:
rnet.load_state_dict(
torch.load(r_model_path,
map_location=lambda storage, loc: storage))
rnet.eval()
if o_model_path is not None:
onet = ONet(use_cuda=use_cuda)
if (use_cuda):
print('o_model_path:{0}'.format(o_model_path))
onet.load_state_dict(torch.load(o_model_path))
onet.cuda()
else:
onet.load_state_dict(
torch.load(o_model_path,
map_location=lambda storage, loc: storage))
onet.eval()
return pnet, rnet, onet
def create_mtcnn_net(self, use_cuda=True):
self.device = torch.device(
"cuda" if use_cuda and torch.cuda.is_available() else "cpu")
pnet = PNet()
summary(pnet.cuda(), (3, 12, 12))
# pnet.load_state_dict(model_zoo.load_url(model_urls['pnet']))
pnet.to(self.device).eval()
rnet = RNet(num_landmarks=config.NUM_LANDMARKS)
summary(rnet.cuda(), (3, 24, 24))
# rnet.load_state_dict(model_zoo.load_url(model_urls['rnet']))
rnet.to(self.device).eval()
onet = ONet(num_landmarks=config.NUM_LANDMARKS)
summary(onet.cuda(), (3, 48, 48))
# onet.load_state_dict(model_zoo.load_url(model_urls['onet']))
onet.to(self.device).eval()
return pnet, rnet, onet
def train_onet(model_store_path,
end_epoch,
imdb,
batch_size,
frequent=50,
base_lr=0.01,
use_cuda=True):
if not os.path.exists(model_store_path):
os.makedirs(model_store_path)
lossfn = LossFn()
net = ONet(is_train=True)
net.train()
print(use_cuda)
if use_cuda:
net.cuda()
optimizer = torch.optim.Adam(net.parameters(), lr=base_lr)
train_data = TrainImageReader(imdb, 48, batch_size, shuffle=True)
for cur_epoch in range(1, end_epoch + 1):
train_data.reset()
for batch_idx, (image, (gt_label, gt_bbox,
gt_landmark)) in enumerate(train_data):
# print("batch id {0}".format(batch_idx))
im_tensor = [
image_tools.convert_image_to_tensor(image[i, :, :, :])
for i in range(image.shape[0])
]
im_tensor = torch.stack(im_tensor)
im_tensor = Variable(im_tensor)
gt_label = Variable(torch.from_numpy(gt_label).float())
gt_bbox = Variable(torch.from_numpy(gt_bbox).float())
gt_landmark = Variable(torch.from_numpy(gt_landmark).float())
if use_cuda:
im_tensor = im_tensor.cuda()
gt_label = gt_label.cuda()
gt_bbox = gt_bbox.cuda()
gt_landmark = gt_landmark.cuda()
cls_pred, box_offset_pred, landmark_offset_pred = net(im_tensor)
# all_loss, cls_loss, offset_loss = lossfn.loss(gt_label=label_y,gt_offset=bbox_y, pred_label=cls_pred, pred_offset=box_offset_pred)
cls_loss = lossfn.cls_loss(gt_label, cls_pred)
box_offset_loss = lossfn.box_loss(gt_label, gt_bbox,
box_offset_pred)
landmark_loss = lossfn.landmark_loss(gt_label, gt_landmark,
landmark_offset_pred)
all_loss = cls_loss * 0.8 + box_offset_loss * 0.6 + landmark_loss * 1.5
if batch_idx % frequent == 0:
accuracy = compute_accuracy(cls_pred, gt_label)
show1 = accuracy.data.cpu().numpy()
show2 = cls_loss.data.cpu().numpy()
show3 = box_offset_loss.data.cpu().numpy()
show4 = landmark_loss.data.cpu().numpy()
show5 = all_loss.data.cpu().numpy()
print(
"%s : Epoch: %d, Step: %d, accuracy: %s, det loss: %s, bbox loss: %s, landmark loss: %s, all_loss: %s, lr:%s "
% (datetime.datetime.now(), cur_epoch, batch_idx, show1,
show2, show3, show4, show5, base_lr))
optimizer.zero_grad()
all_loss.backward()
optimizer.step()
torch.save(
net.state_dict(),
os.path.join(model_store_path, "onet_epoch_%d.pt" % cur_epoch))
torch.save(
net,
os.path.join(model_store_path,
"onet_epoch_model_%d.pkl" % cur_epoch))