OpenI
/
dface
Not watched
1
0
Fork 0
Code
Releases 2 Wiki Activity
Issues 0 Pull Requests 0 Datasets Model Cloudbrain
You can not select more than 25 topics Topics must start with a chinese character,a letter or number, can include dashes ('-') and can be up to 35 characters long.
Tree: 6a8b154a0d
Branches Tags
${ item.name }
Create branch ${ searchTerm }
from '6a8b154a0d'
${ noResults }
dface/dface/train_net/train.py

282 lines
11 kB
Raw Blame History 

  1. from dface.core.image_reader import TrainImageReader

  2. import datetime

  3. import os

  4. from dface.core.models import PNet,RNet,ONet,LossFn

  5. import torch

  6. from torch.autograd import Variable

  7. import dface.core.image_tools as image_tools

  8. 

  9. 

  10. 

  11. 

  12. 

  13. def compute_accuracy(prob_cls, gt_cls):

  14.     prob_cls = torch.squeeze(prob_cls)

  15.     gt_cls = torch.squeeze(gt_cls)

  16. 

  17.     #we only need the detection which >= 0

  18.     mask = torch.ge(gt_cls,0)

  19.     #get valid element

  20.     valid_gt_cls = torch.masked_select(gt_cls,mask)

  21.     valid_prob_cls = torch.masked_select(prob_cls,mask)

  22.     size = min(valid_gt_cls.size()[0], valid_prob_cls.size()[0])

  23.     prob_ones = torch.ge(valid_prob_cls,0.6).float()

  24.     right_ones = torch.eq(prob_ones,valid_gt_cls).float()

  25. 

  26.     return torch.div(torch.mul(torch.sum(right_ones),float(1.0)),float(size))

  27. 

  28. 

  29. def train_pnet(model_store_path, end_epoch,imdb,

  30.               batch_size,frequent=50,base_lr=0.01,use_cuda=True):

  31. 

  32.     if not os.path.exists(model_store_path):

  33.         os.makedirs(model_store_path)

  34. 

  35.     lossfn = LossFn()

  36.     net = PNet(is_train=True, use_cuda=use_cuda)

  37.     net.train()

  38.     if use_cuda:

  39.         net.cuda()

  40. 

  41.     optimizer = torch.optim.Adam(net.parameters(), lr=base_lr)

  42. 

  43.     train_data=TrainImageReader(imdb,12,batch_size,shuffle=True)

  44. 

  45. 

  46.     for cur_epoch in range(1,end_epoch+1):

  47.         train_data.reset()

  48.         accuracy_list=[]

  49.         cls_loss_list=[]

  50.         bbox_loss_list=[]

  51.         # landmark_loss_list=[]

  52. 

  53.         for batch_idx,(image,(gt_label,gt_bbox,gt_landmark))in enumerate(train_data):

  54. 

  55.             im_tensor = [ image_tools.convert_image_to_tensor(image[i,:,:,:]) for i in range(image.shape[0]) ]

  56.             im_tensor = torch.stack(im_tensor)

  57. 

  58.             im_tensor = Variable(im_tensor)

  59.             gt_label = Variable(torch.from_numpy(gt_label).float())

  60. 

  61.             gt_bbox = Variable(torch.from_numpy(gt_bbox).float())

  62.             # gt_landmark = Variable(torch.from_numpy(gt_landmark).float())

  63. 

  64.             if use_cuda:

  65.                 im_tensor = im_tensor.cuda()

  66.                 gt_label = gt_label.cuda()

  67.                 gt_bbox = gt_bbox.cuda()

  68.                 # gt_landmark = gt_landmark.cuda()

  69. 

  70.             cls_pred, box_offset_pred = net(im_tensor)

  71.             # 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)

  72. 

  73.             cls_loss = lossfn.cls_loss(gt_label,cls_pred)

  74.             box_offset_loss = lossfn.box_loss(gt_label,gt_bbox,box_offset_pred)

  75.             # landmark_loss = lossfn.landmark_loss(gt_label,gt_landmark,landmark_offset_pred)

  76. 

  77.             all_loss = cls_loss*1.0+box_offset_loss*0.5

  78. 

  79.             if batch_idx%frequent==0:

  80.                 accuracy=compute_accuracy(cls_pred,gt_label)

  81. 

  82.                 show1 = accuracy.data.tolist()[0]

  83.                 show2 = cls_loss.data.tolist()[0]

  84.                 show3 = box_offset_loss.data.tolist()[0]

  85.                 show5 = all_loss.data.tolist()[0]

  86. 

  87.                 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))

  88.                 accuracy_list.append(accuracy)

  89.                 cls_loss_list.append(cls_loss)

  90.                 bbox_loss_list.append(box_offset_loss)

  91. 

  92.             optimizer.zero_grad()

  93.             all_loss.backward()

  94.             optimizer.step()

  95. 

  96. 

  97.         accuracy_avg = torch.mean(torch.cat(accuracy_list))

  98.         cls_loss_avg = torch.mean(torch.cat(cls_loss_list))

  99.         bbox_loss_avg = torch.mean(torch.cat(bbox_loss_list))

  100.         # landmark_loss_avg = torch.mean(torch.cat(landmark_loss_list))

  101. 

  102.         show6 = accuracy_avg.data.tolist()[0]

  103.         show7 = cls_loss_avg.data.tolist()[0]

  104.         show8 = bbox_loss_avg.data.tolist()[0]

  105. 

  106.         print("Epoch: %d, accuracy: %s, cls loss: %s, bbox loss: %s" % (cur_epoch, show6, show7, show8))

  107.         torch.save(net.state_dict(), os.path.join(model_store_path,"pnet_epoch_%d.pt" % cur_epoch))

  108.         torch.save(net, os.path.join(model_store_path,"pnet_epoch_model_%d.pkl" % cur_epoch))

  109. 

  110. 

  111. 

  112. 

  113. def train_rnet(model_store_path, end_epoch,imdb,

  114.               batch_size,frequent=50,base_lr=0.01,use_cuda=True):

  115. 

  116.     if not os.path.exists(model_store_path):

  117.         os.makedirs(model_store_path)

  118. 

  119.     lossfn = LossFn()

  120.     net = RNet(is_train=True, use_cuda=use_cuda)

  121.     net.train()

  122.     if use_cuda:

  123.         net.cuda()

  124. 

  125.     optimizer = torch.optim.Adam(net.parameters(), lr=base_lr)

  126. 

  127.     train_data=TrainImageReader(imdb,24,batch_size,shuffle=True)

  128. 

  129. 

  130.     for cur_epoch in range(1,end_epoch+1):

  131.         train_data.reset()

  132.         accuracy_list=[]

  133.         cls_loss_list=[]

  134.         bbox_loss_list=[]

  135.         landmark_loss_list=[]

  136. 

  137.         for batch_idx,(image,(gt_label,gt_bbox,gt_landmark))in enumerate(train_data):

  138. 

  139.             im_tensor = [ image_tools.convert_image_to_tensor(image[i,:,:,:]) for i in range(image.shape[0]) ]

  140.             im_tensor = torch.stack(im_tensor)

  141. 

  142.             im_tensor = Variable(im_tensor)

  143.             gt_label = Variable(torch.from_numpy(gt_label).float())

  144. 

  145.             gt_bbox = Variable(torch.from_numpy(gt_bbox).float())

  146.             gt_landmark = Variable(torch.from_numpy(gt_landmark).float())

  147. 

  148.             if use_cuda:

  149.                 im_tensor = im_tensor.cuda()

  150.                 gt_label = gt_label.cuda()

  151.                 gt_bbox = gt_bbox.cuda()

  152.                 gt_landmark = gt_landmark.cuda()

  153. 

  154.             cls_pred, box_offset_pred = net(im_tensor)

  155.             # 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)

  156. 

  157.             cls_loss = lossfn.cls_loss(gt_label,cls_pred)

  158.             box_offset_loss = lossfn.box_loss(gt_label,gt_bbox,box_offset_pred)

  159.             # landmark_loss = lossfn.landmark_loss(gt_label,gt_landmark,landmark_offset_pred)

  160. 

  161.             all_loss = cls_loss*1.0+box_offset_loss*0.5

  162. 

  163.             if batch_idx%frequent==0:

  164.                 accuracy=compute_accuracy(cls_pred,gt_label)

  165. 

  166.                 show1 = accuracy.data.tolist()[0]

  167.                 show2 = cls_loss.data.tolist()[0]

  168.                 show3 = box_offset_loss.data.tolist()[0]

  169.                 # show4 = landmark_loss.data.tolist()[0]

  170.                 show5 = all_loss.data.tolist()[0]

  171. 

  172.                 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))

  173.                 accuracy_list.append(accuracy)

  174.                 cls_loss_list.append(cls_loss)

  175.                 bbox_loss_list.append(box_offset_loss)

  176.                 # landmark_loss_list.append(landmark_loss)

  177. 

  178.             optimizer.zero_grad()

  179.             all_loss.backward()

  180.             optimizer.step()

  181. 

  182. 

  183.         accuracy_avg = torch.mean(torch.cat(accuracy_list))

  184.         cls_loss_avg = torch.mean(torch.cat(cls_loss_list))

  185.         bbox_loss_avg = torch.mean(torch.cat(bbox_loss_list))

  186.         # landmark_loss_avg = torch.mean(torch.cat(landmark_loss_list))

  187. 

  188.         show6 = accuracy_avg.data.tolist()[0]

  189.         show7 = cls_loss_avg.data.tolist()[0]

  190.         show8 = bbox_loss_avg.data.tolist()[0]

  191.         # show9 = landmark_loss_avg.data.tolist()[0]

  192. 

  193.         print("Epoch: %d, accuracy: %s, cls loss: %s, bbox loss: %s" % (cur_epoch, show6, show7, show8))

  194.         torch.save(net.state_dict(), os.path.join(model_store_path,"rnet_epoch_%d.pt" % cur_epoch))

  195.         torch.save(net, os.path.join(model_store_path,"rnet_epoch_model_%d.pkl" % cur_epoch))

  196. 

  197. 

  198. def train_onet(model_store_path, end_epoch,imdb,

  199.               batch_size,frequent=50,base_lr=0.01,use_cuda=True):

  200. 

  201.     if not os.path.exists(model_store_path):

  202.         os.makedirs(model_store_path)

  203. 

  204.     lossfn = LossFn()

  205.     net = ONet(is_train=True)

  206.     net.train()

  207.     if use_cuda:

  208.         net.cuda()

  209. 

  210.     optimizer = torch.optim.Adam(net.parameters(), lr=base_lr)

  211. 

  212.     train_data=TrainImageReader(imdb,48,batch_size,shuffle=True)

  213. 

  214. 

  215.     for cur_epoch in range(1,end_epoch+1):

  216.         train_data.reset()

  217.         accuracy_list=[]

  218.         cls_loss_list=[]

  219.         bbox_loss_list=[]

  220.         landmark_loss_list=[]

  221. 

  222.         for batch_idx,(image,(gt_label,gt_bbox,gt_landmark))in enumerate(train_data):

  223. 

  224.             im_tensor = [ image_tools.convert_image_to_tensor(image[i,:,:,:]) for i in range(image.shape[0]) ]

  225.             im_tensor = torch.stack(im_tensor)

  226. 

  227.             im_tensor = Variable(im_tensor)

  228.             gt_label = Variable(torch.from_numpy(gt_label).float())

  229. 

  230.             gt_bbox = Variable(torch.from_numpy(gt_bbox).float())

  231.             gt_landmark = Variable(torch.from_numpy(gt_landmark).float())

  232. 

  233.             if use_cuda:

  234.                 im_tensor = im_tensor.cuda()

  235.                 gt_label = gt_label.cuda()

  236.                 gt_bbox = gt_bbox.cuda()

  237.                 gt_landmark = gt_landmark.cuda()

  238. 

  239.             cls_pred, box_offset_pred, landmark_offset_pred = net(im_tensor)

  240.             # 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)

  241. 

  242.             cls_loss = lossfn.cls_loss(gt_label,cls_pred)

  243.             box_offset_loss = lossfn.box_loss(gt_label,gt_bbox,box_offset_pred)

  244.             landmark_loss = lossfn.landmark_loss(gt_label,gt_landmark,landmark_offset_pred)

  245. 

  246.             all_loss = cls_loss*0.8+box_offset_loss*0.6+landmark_loss*1.5

  247. 

  248.             if batch_idx%frequent==0:

  249.                 accuracy=compute_accuracy(cls_pred,gt_label)

  250. 

  251.                 show1 = accuracy.data.tolist()[0]

  252.                 show2 = cls_loss.data.tolist()[0]

  253.                 show3 = box_offset_loss.data.tolist()[0]

  254.                 show4 = landmark_loss.data.tolist()[0]

  255.                 show5 = all_loss.data.tolist()[0]

  256. 

  257.                 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))

  258.                 accuracy_list.append(accuracy)

  259.                 cls_loss_list.append(cls_loss)

  260.                 bbox_loss_list.append(box_offset_loss)

  261.                 landmark_loss_list.append(landmark_loss)

  262. 

  263.             optimizer.zero_grad()

  264.             all_loss.backward()

  265.             optimizer.step()

  266. 

  267. 

  268.         accuracy_avg = torch.mean(torch.cat(accuracy_list))

  269.         cls_loss_avg = torch.mean(torch.cat(cls_loss_list))

  270.         bbox_loss_avg = torch.mean(torch.cat(bbox_loss_list))

  271.         landmark_loss_avg = torch.mean(torch.cat(landmark_loss_list))

  272. 

  273.         show6 = accuracy_avg.data.tolist()[0]

  274.         show7 = cls_loss_avg.data.tolist()[0]

  275.         show8 = bbox_loss_avg.data.tolist()[0]

  276.         show9 = landmark_loss_avg.data.tolist()[0]

  277. 

  278.         print("Epoch: %d, accuracy: %s, cls loss: %s, bbox loss: %s, landmark loss: %s " % (cur_epoch, show6, show7, show8, show9))

  279.         torch.save(net.state_dict(), os.path.join(model_store_path,"onet_epoch_%d.pt" % cur_epoch))

  280.         torch.save(net, os.path.join(model_store_path,"onet_epoch_model_%d.pkl" % cur_epoch))