def main(args):
# p_model_path = '/data/guoch_workspace/mtcnn-pytorch-master/model_store/pnet_epoch.pt'
# r_model_path = '/data/guoch_workspace/mtcnn-pytorch-master/model_store/rnet_epoch.pt'
# o_model_path = '/data/guoch_workspace/mtcnn-pytorch-master/model_store/onet_epoch.pt'
# print("the version of torch is {}".format(torch.__version__))
dummy_input = getInput(args.img_size) #获得网络的输入
# 加载模型
model = PNet()
#model = RNet()
#model = ONet()
model.load_state_dict(torch.load(args.model_path))
#model_dict = model.state_dict()
#model_dict = pnet.load_state_dict(torch.load(p_model_path))
# if args.model_path:
# if os.path.isfile(args.model_path):
# print(("=> start loading checkpoint '{}'".format(args.model_path)))
# # state_dict = torch.load(args.model_path)
# # print("the best acc is {} in epoch:{}".format(
# # state_dict['epoch_acc'], state_dict['epoch']))
# # params = state_dict["model_state_dict"]
# # # params={k:v for k,v in state_dict.items() if k in model_dict.keys()}
# # # model_dict.update(params)
# # # model.load_state_dict(model_dict)
# model.load_state_dict(args.model_path)
# print("load cls model successfully")
# else:
# print(("=> no checkpoint found at '{}'".format(args.model_path)))
# return
model.to('cpu')
model.eval()
pre = model(dummy_input)
print("the pre:{}".format(pre))
#保存onnx模型
torch2onnx(args, model, dummy_input)
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):
''' Create the mtcnn model '''
pnet, rnet, onet = None, None, None
if len(self.args.pnet_file) > 0:
pnet = PNet(use_cuda=self.args.use_cuda)
if self.args.use_cuda:
pnet.load_state_dict(torch.load(self.args.pnet_file))
pnet = torch.nn.DataParallel(
pnet, device_ids=self.args.gpu_ids).cuda()
else:
pnet.load_state_dict(torch.load(self.args.pnet_file,\
map_location=lambda storage, loc: storage))
pnet.eval()
if len(self.args.rnet_file) > 0:
rnet = RNet(use_cuda=self.args.use_cuda)
if self.args.use_cuda:
rnet.load_state_dict(torch.load(self.args.rnet_file))
rnet = torch.nn.DataParallel(
rnet, device_ids=self.args.gpu_ids).cuda()
else:
rnet.load_state_dict(torch.load(self.args.rnet_file,\
map_location=lambda storage, loc: storage))
rnet.eval()
if len(self.args.onet_file) > 0:
onet = ONet(use_cuda=self.args.use_cuda)
if self.args.use_cuda:
onet.load_state_dict(torch.load(self.args.onet_file))
onet = torch.nn.DataParallel(
onet, device_ids=self.args.gpu_ids).cuda()
else:
onet.load_state_dict(torch.load(self.args.onet_file, \
map_location=lambda storage, loc: storage))
onet.eval()
self.pnet_detector = pnet
self.rnet_detector = rnet
self.onet_detector = onet
def train_pnet(model_store_path, end_epoch,imdb,
batch_size,frequent=10,base_lr=0.01,use_cuda=True):
if not os.path.exists(model_store_path):
os.makedirs(model_store_path)
lossfn = LossFn()
net = PNet(is_train=True, use_cuda=use_cuda)
net.train()
if use_cuda:
net.cuda()
optimizer = torch.optim.Adam(net.parameters(), lr=base_lr)
train_data=TrainImageReader(imdb,12,batch_size,shuffle=True)
frequent = 10
for cur_epoch in range(1,end_epoch+1):
train_data.reset() # shuffle
for batch_idx,(image,(gt_label,gt_bbox,gt_landmark))in enumerate(train_data):
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 = 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*1.0+box_offset_loss*0.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, all_loss: %s, lr:%s "%(datetime.datetime.now(),cur_epoch,batch_idx, show1,show2,show3,show5,base_lr))
optimizer.zero_grad()
all_loss.backward()
optimizer.step()
torch.save(net.state_dict(), os.path.join(model_store_path,"pnet_epoch_%d.pt" % cur_epoch))
torch.save(net, os.path.join(model_store_path,"pnet_epoch_model_%d.pkl" % cur_epoch))
def train_pnet(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 = PNet(is_train=True, use_cuda=use_cuda)
checkpoint = torch.load('model_store/pnet_epoch_4.pt')
net.load_state_dict(checkpoint)
net.train()
if use_cuda:
net.cuda()
optimizer = torch.optim.Adam(net.parameters(), lr=base_lr)
train_data = TrainImageReader(imdb, 12, batch_size, shuffle=True)
for cur_epoch in range(1, end_epoch + 1):
train_data.reset()
accuracy_list = []
cls_loss_list = []
bbox_loss_list = []
# landmark_loss_list=[]
for batch_idx, (image, (gt_label, gt_bbox,
gt_landmark)) in enumerate(train_data):
im_tensor = [
image_tools.convert_image_to_tensor(image[i, :, :, :])
for i in range(image.shape[0])
]
im_tensor = torch.stack(im_tensor).float()
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 = 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 * 1.0 + box_offset_loss * 0.5
if batch_idx % frequent == 0:
accuracy = compute_accuracy(cls_pred, gt_label)
# show1 = accuracy.data.tolist()[0]
# show2 = cls_loss.data.tolist()[0]
# show3 = box_offset_loss.data.tolist()[0]
# show5 = all_loss.data.tolist()[0]
show1 = accuracy.item()
show2 = cls_loss.item()
show3 = box_offset_loss.item()
show5 = all_loss.item()
print(
"%s : Epoch: %d, Step: %d, accuracy: %.4f, det loss: %.4f, bbox loss: %.4f, all_loss: %.4f, lr:%s "
% (datetime.datetime.now(), cur_epoch, batch_idx, show1,
show2, show3, show5, base_lr))
accuracy_list.append(accuracy)
cls_loss_list.append(cls_loss)
bbox_loss_list.append(box_offset_loss)
optimizer.zero_grad()
all_loss.backward()
optimizer.step()
accuracy_avg = torch.mean(torch.tensor(accuracy_list))
cls_loss_avg = torch.mean(torch.tensor(cls_loss_list))
bbox_loss_avg = torch.mean(torch.tensor(bbox_loss_list))
# landmark_loss_avg = torch.mean(torch.cat(landmark_loss_list))
# show6 = accuracy_avg.data.tolist()[0]
# show7 = cls_loss_avg.data.tolist()[0]
# show8 = bbox_loss_avg.data.tolist()[0]
show6 = accuracy_avg.item()
show7 = cls_loss_avg.item()
show8 = bbox_loss_avg.item()
print("Epoch: %d, accuracy: %s, cls loss: %s, bbox loss: %s" %
(cur_epoch, show6, show7, show8))
# state = {'net': net.state_dict(), 'optimizer': optimizer.state_dict(), 'epoch': cur_epoch}
torch.save(
net.state_dict(),
os.path.join(model_store_path, "pnet_epoch_%d.pt" % cur_epoch))
torch.save(
net,
os.path.join(model_store_path,
"pnet_epoch_model_%d.pkl" % cur_epoch))
