else: pretrained_dict = remove_prefix(pretrained_dict, 'module.') check_keys(model, pretrained_dict) model.load_state_dict(pretrained_dict, strict=False) return model if __name__ == '__main__': # net and model net = FaceBoxes(phase='test', size=None, num_classes=2) # initialize detector net = load_model(net, args.trained_model) net.eval() print('Finished loading model!') print(net) if args.cuda: net = net.cuda() cudnn.benchmark = True else: net = net.cpu() # save file if not os.path.exists(args.save_folder): os.makedirs(args.save_folder) fw = open(os.path.join(args.save_folder, args.dataset + '_dets.txt'), 'w') # testing dataset testset_folder = os.path.join('data', args.dataset, 'images/') testset_list = os.path.join('data', args.dataset, 'img_list.txt') with open(testset_list, 'r') as fr: test_dataset = fr.read().split() num_images = len(test_dataset)
from collections import OrderedDict new_state_dict = OrderedDict() for k, v in state_dict.items(): head = k[:7] if head == 'module.': name = k[7:] # remove `module.` else: name = k new_state_dict[name] = v net.load_state_dict(new_state_dict) if args.ngpu > 1: net = torch.nn.DataParallel(net, device_ids=list(range(args.ngpu))) if args.cuda: net.cuda() cudnn.benchmark = True optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) criterion = MultiBoxLoss(num_classes, 0.35, True, 0, True, 7, 0.35, False) priorbox = PriorBox(cfg) with torch.no_grad(): priors = priorbox.forward() if args.cuda: priors = priors.cuda()
def main(): global args global minmum_loss args.gpu = 0 args.world_size = 1 if args.distributed: args.gpu = args.local_rank % torch.cuda.device_count() torch.cuda.set_device(args.gpu) torch.distributed.init_process_group(backend='nccl', init_method='env://') args.world_size = torch.distributed.get_world_size() args.total_batch_size = args.world_size * args.batch_size model = FaceBoxes('train', args.num_classes) print("Printing net...") if args.distributed: model = torch.nn.parallel.DistributedDataParallel(model) model = model.cuda() # optimizer and loss function optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) criterion = MultiBoxLoss(num_classes=args.num_classes, overlap_thresh=0.35, prior_for_matching=True, bkg_label=0, neg_mining=True, neg_pos=7, neg_overlap=0.35, encode_target=False) # optionally resume from a checkpoint if args.resume: if os.path.isfile(args.resume): print("=> loading checkpoint '{}'".format(args.resume)) checkpoint = torch.load( args.resume, map_location=lambda storage, loc: storage.cuda(args.gpu)) args.start_epoch = checkpoint['epoch'] minmum_loss = checkpoint['minmum_loss'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) print("=> loaded checkpoint '{}' (epoch {})".format( args.resume, checkpoint['epoch'])) else: print("=> no checkpoint found at '{}'".format(args.resume)) # Data loading code print('Loading Dataset...') dataset = VOCDetection(args.training_dataset, preproc(args.img_dim, args.rgb_mean), AnnotationTransform()) train_loader = data.DataLoader(dataset, args.batch_size, num_workers=args.num_workers, shuffle=True, collate_fn=detection_collate, pin_memory=True) priorbox = PriorBox(cfg, image_size=(args.img_dim, args.img_dim)) with torch.no_grad(): priors = priorbox.forward() priors = priors.cuda() for epoch in range(args.start_epoch, args.epochs): # train for one epoch end = time.time() loss = train(train_loader, model, priors, criterion, optimizer, epoch) if args.local_rank == 0: is_best = loss < minmum_loss minmum_loss = min(loss, minmum_loss) save_checkpoint( { 'epoch': epoch + 1, 'state_dict': model.state_dict(), 'best_prec1': minmum_loss, 'optimizer': optimizer.state_dict(), }, is_best, epoch) epoch_time = time.time() - end print('Epoch %s time cost %f' % (epoch, epoch_time))
class FaceBoxDetector: def __init__(self, min_score=0.9, use_gpu=True): # Minimum score to consider as a detection. self.score_min = min_score self.net = FaceBoxes(phase='test', size=None, num_classes=2) self.use_gpu = use_gpu self.logger = Logger() def load_model(self, path_to_model): self.logger.field('Loading pretrained model from', path_to_model) device = torch.cuda.current_device() pretrained_dict = torch.load(path_to_model, map_location=lambda storage, loc: storage.cuda(device)) if "state_dict" in pretrained_dict.keys(): pretrained_dict = self.remove_prefix(pretrained_dict['state_dict'], 'module.') else: pretrained_dict = self.remove_prefix(pretrained_dict, 'module.') self.check_keys(self.net, pretrained_dict) self.net.load_state_dict(pretrained_dict, strict=False) self.net.eval() if self.use_gpu: self.net.cuda() def detect(self, images) -> List[List[TrackingRegion]]: frames = [] for img in images: img = np.float32(img) im_height, im_width, _ = img.shape scale = torch.Tensor([img.shape[1], img.shape[0], img.shape[1], img.shape[0]]) img -= (104, 117, 123) img = img.transpose(2, 0, 1) img = torch.from_numpy(img).unsqueeze(0) if self.use_gpu: img = img.cuda() scale = scale.cuda() out = self.net(img) face_regions = self.nms_process(out, scale, im_height, im_width) frames.append(face_regions) return frames def nms_process(self, network_output, scale, im_height, im_width) -> List[TrackingRegion]: priorbox = PriorBox(cfg, network_output[2], (im_height, im_width), phase='test') priors = priorbox.forward() if self.use_gpu: priors = priors.cuda() loc, conf, _ = network_output prior_data = priors.data boxes = decode(loc.data.squeeze(0), prior_data, cfg['variance']) boxes = boxes * scale boxes = boxes.cpu().numpy() scores = conf.data.cpu().numpy()[:, 1] # ignore low scores inds = np.where(scores > self.score_min)[0] boxes = boxes[inds] scores = scores[inds] # keep top-K before NMS, top_k = 5 order = scores.argsort()[::-1][:5000] boxes = boxes[order] scores = scores[order] # do NMS dets = np.hstack((boxes, scores[:, np.newaxis])).astype(np.float32, copy=False) keep = nms(dets, 0.3, force_cpu=False) dets = dets[keep, :] # keep top-K faster NMS dets = dets[:750, :] regions = [] for i in range(dets.shape[0]): face_region = TrackingRegion() face_region.set_rect(left=dets[i, 0], top=dets[i, 1], right=dets[i, 2], bottom=dets[i, 3]) face_region.confidence = dets[i, 4] face_region.data["class_id"] = "face" regions.append(face_region) return regions def check_keys(self, model, pretrained_state_dict): ckpt_keys = set(pretrained_state_dict.keys()) model_keys = set(model.state_dict().keys()) used_pretrained_keys = model_keys & ckpt_keys unused_pretrained_keys = ckpt_keys - model_keys missing_keys = model_keys - ckpt_keys self.logger.field('Missing keys', len(missing_keys)) self.logger.field('Unused checkpoint keys', len(unused_pretrained_keys)) self.logger.field('Used keys', len(used_pretrained_keys)) assert len(used_pretrained_keys) > 0, 'load NONE from pretrained checkpoint' return True @staticmethod def remove_prefix(state_dict, prefix): """ Old style model is stored with all names of parameters sharing common prefix 'module.' """ print('remove prefix \'{}\''.format(prefix)) f = lambda x: x.split(prefix, 1)[-1] if x.startswith(prefix) else x return {f(key): value for key, value in state_dict.items()}