def get_model():
pretrain_model = "/Users/yangjiang/temp/gpu1/FaceBoxes.pth"
num_classes = CONFIG["num_classes"]
params = torch.load(pretrain_model, map_location="cpu")
model = FaceBoxes(num_classes, "test")
model.load_state_dict(params)
model.eval()
return model
print(net)
if args.resume_net is not None:
print('Loading resume network...')
state_dict = torch.load(args.resume_net)
# create new OrderedDict that does not contain `module.`
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 and gpu_train:
net = torch.nn.DataParallel(net, device_ids=list(range(args.ngpu)))
device = torch.device('cuda:0' if gpu_train else 'cpu')
cudnn.benchmark = True
net = net.to(device)
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()
priors = priors.to(device)
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()}
