towhee
/
retinaface-face-detection
copied
4 changed files with 559 additions and 0 deletions
@ -0,0 +1,328 @@ |
|||||
|
import torch |
||||
|
import numpy as np |
||||
|
|
||||
|
|
||||
|
def point_form(boxes): |
||||
|
""" Convert prior_boxes to (xmin, ymin, xmax, ymax) |
||||
|
representation for comparison to point form ground truth data. |
||||
|
Args: |
||||
|
boxes: (tensor) center-size default boxes from priorbox layers. |
||||
|
Return: |
||||
|
boxes: (tensor) Converted xmin, ymin, xmax, ymax form of boxes. |
||||
|
""" |
||||
|
return torch.cat((boxes[:, :2] - boxes[:, 2:]/2, # xmin, ymin |
||||
|
boxes[:, :2] + boxes[:, 2:]/2), 1) # xmax, ymax |
||||
|
|
||||
|
|
||||
|
def center_size(boxes): |
||||
|
""" Convert prior_boxes to (cx, cy, w, h) |
||||
|
representation for comparison to center-size form ground truth data. |
||||
|
Args: |
||||
|
boxes: (tensor) point_form boxes |
||||
|
Return: |
||||
|
boxes: (tensor) Converted xmin, ymin, xmax, ymax form of boxes. |
||||
|
""" |
||||
|
return torch.cat((boxes[:, 2:] + boxes[:, :2])/2, # cx, cy |
||||
|
boxes[:, 2:] - boxes[:, :2], 1) # w, h |
||||
|
|
||||
|
|
||||
|
def intersect(box_a, box_b): |
||||
|
""" We resize both tensors to [A,B,2] without new malloc: |
||||
|
[A,2] -> [A,1,2] -> [A,B,2] |
||||
|
[B,2] -> [1,B,2] -> [A,B,2] |
||||
|
Then we compute the area of intersect between box_a and box_b. |
||||
|
Args: |
||||
|
box_a: (tensor) bounding boxes, Shape: [A,4]. |
||||
|
box_b: (tensor) bounding boxes, Shape: [B,4]. |
||||
|
Return: |
||||
|
(tensor) intersection area, Shape: [A,B]. |
||||
|
""" |
||||
|
A = box_a.size(0) |
||||
|
B = box_b.size(0) |
||||
|
max_xy = torch.min(box_a[:, 2:].unsqueeze(1).expand(A, B, 2), |
||||
|
box_b[:, 2:].unsqueeze(0).expand(A, B, 2)) |
||||
|
min_xy = torch.max(box_a[:, :2].unsqueeze(1).expand(A, B, 2), |
||||
|
box_b[:, :2].unsqueeze(0).expand(A, B, 2)) |
||||
|
inter = torch.clamp((max_xy - min_xy), min=0) |
||||
|
return inter[:, :, 0] * inter[:, :, 1] |
||||
|
|
||||
|
|
||||
|
def jaccard(box_a, box_b): |
||||
|
"""Compute the jaccard overlap of two sets of boxes. The jaccard overlap |
||||
|
is simply the intersection over union of two boxes. Here we operate on |
||||
|
ground truth boxes and default boxes. |
||||
|
E.g.: |
||||
|
A ∩ B / A ∪ B = A ∩ B / (area(A) + area(B) - A ∩ B) |
||||
|
Args: |
||||
|
box_a: (tensor) Ground truth bounding boxes, Shape: [num_objects,4] |
||||
|
box_b: (tensor) Prior boxes from priorbox layers, Shape: [num_priors,4] |
||||
|
Return: |
||||
|
jaccard overlap: (tensor) Shape: [box_a.size(0), box_b.size(0)] |
||||
|
""" |
||||
|
inter = intersect(box_a, box_b) |
||||
|
area_a = ((box_a[:, 2]-box_a[:, 0]) * |
||||
|
(box_a[:, 3]-box_a[:, 1])).unsqueeze(1).expand_as(inter) # [A,B] |
||||
|
area_b = ((box_b[:, 2]-box_b[:, 0]) * |
||||
|
(box_b[:, 3]-box_b[:, 1])).unsqueeze(0).expand_as(inter) # [A,B] |
||||
|
union = area_a + area_b - inter |
||||
|
return inter / union # [A,B] |
||||
|
|
||||
|
|
||||
|
def matrix_iou(a, b): |
||||
|
""" |
||||
|
return iou of a and b, numpy version for data augenmentation |
||||
|
""" |
||||
|
lt = np.maximum(a[:, np.newaxis, :2], b[:, :2]) |
||||
|
rb = np.minimum(a[:, np.newaxis, 2:], b[:, 2:]) |
||||
|
|
||||
|
area_i = np.prod(rb - lt, axis=2) * (lt < rb).all(axis=2) |
||||
|
area_a = np.prod(a[:, 2:] - a[:, :2], axis=1) |
||||
|
area_b = np.prod(b[:, 2:] - b[:, :2], axis=1) |
||||
|
return area_i / (area_a[:, np.newaxis] + area_b - area_i) |
||||
|
|
||||
|
|
||||
|
def matrix_iof(a, b): |
||||
|
""" |
||||
|
return iof of a and b, numpy version for data augenmentation |
||||
|
""" |
||||
|
lt = np.maximum(a[:, np.newaxis, :2], b[:, :2]) |
||||
|
rb = np.minimum(a[:, np.newaxis, 2:], b[:, 2:]) |
||||
|
|
||||
|
area_i = np.prod(rb - lt, axis=2) * (lt < rb).all(axis=2) |
||||
|
area_a = np.prod(a[:, 2:] - a[:, :2], axis=1) |
||||
|
return area_i / np.maximum(area_a[:, np.newaxis], 1) |
||||
|
|
||||
|
|
||||
|
def match(threshold, truths, priors, variances, labels, landms, loc_t, conf_t, landm_t, idx): |
||||
|
"""Match each prior box with the ground truth box of the highest jaccard |
||||
|
overlap, encode the bounding boxes, then return the matched indices |
||||
|
corresponding to both confidence and location preds. |
||||
|
Args: |
||||
|
threshold: (float) The overlap threshold used when mathing boxes. |
||||
|
truths: (tensor) Ground truth boxes, Shape: [num_obj, 4]. |
||||
|
priors: (tensor) Prior boxes from priorbox layers, Shape: [n_priors,4]. |
||||
|
variances: (tensor) Variances corresponding to each prior coord, |
||||
|
Shape: [num_priors, 4]. |
||||
|
labels: (tensor) All the class labels for the image, Shape: [num_obj]. |
||||
|
landms: (tensor) Ground truth landms, Shape [num_obj, 10]. |
||||
|
loc_t: (tensor) Tensor to be filled w/ endcoded location targets. |
||||
|
conf_t: (tensor) Tensor to be filled w/ matched indices for conf preds. |
||||
|
landm_t: (tensor) Tensor to be filled w/ endcoded landm targets. |
||||
|
idx: (int) current batch index |
||||
|
Return: |
||||
|
The matched indices corresponding to 1)location 2)confidence 3)landm preds. |
||||
|
""" |
||||
|
# jaccard index |
||||
|
overlaps = jaccard( |
||||
|
truths, |
||||
|
point_form(priors) |
||||
|
) |
||||
|
# (Bipartite Matching) |
||||
|
# [1,num_objects] best prior for each ground truth |
||||
|
best_prior_overlap, best_prior_idx = overlaps.max(1, keepdim=True) |
||||
|
|
||||
|
# ignore hard gt |
||||
|
valid_gt_idx = best_prior_overlap[:, 0] >= 0.2 |
||||
|
best_prior_idx_filter = best_prior_idx[valid_gt_idx, :] |
||||
|
if best_prior_idx_filter.shape[0] <= 0: |
||||
|
loc_t[idx] = 0 |
||||
|
conf_t[idx] = 0 |
||||
|
return |
||||
|
|
||||
|
# [1,num_priors] best ground truth for each prior |
||||
|
best_truth_overlap, best_truth_idx = overlaps.max(0, keepdim=True) |
||||
|
best_truth_idx.squeeze_(0) |
||||
|
best_truth_overlap.squeeze_(0) |
||||
|
best_prior_idx.squeeze_(1) |
||||
|
best_prior_idx_filter.squeeze_(1) |
||||
|
best_prior_overlap.squeeze_(1) |
||||
|
best_truth_overlap.index_fill_(0, best_prior_idx_filter, 2) # ensure best prior |
||||
|
# TODO refactor: index best_prior_idx with long tensor |
||||
|
# ensure every gt matches with its prior of max overlap |
||||
|
for j in range(best_prior_idx.size(0)): # 判别此anchor是预测哪一个boxes |
||||
|
best_truth_idx[best_prior_idx[j]] = j |
||||
|
matches = truths[best_truth_idx] # Shape: [num_priors,4] 此处为每一个anchor对应的bbox取出来 |
||||
|
conf = labels[best_truth_idx] # Shape: [num_priors] 此处为每一个anchor对应的label取出来 |
||||
|
conf[best_truth_overlap < threshold] = 0 # label as background overlap<0.35的全部作为负样本 |
||||
|
loc = encode(matches, priors, variances) |
||||
|
|
||||
|
matches_landm = landms[best_truth_idx] |
||||
|
landm = encode_landm(matches_landm, priors, variances) |
||||
|
loc_t[idx] = loc # [num_priors,4] encoded offsets to learn |
||||
|
conf_t[idx] = conf # [num_priors] top class label for each prior |
||||
|
landm_t[idx] = landm |
||||
|
|
||||
|
|
||||
|
def encode(matched, priors, variances): |
||||
|
"""Encode the variances from the priorbox layers into the ground truth boxes |
||||
|
we have matched (based on jaccard overlap) with the prior boxes. |
||||
|
Args: |
||||
|
matched: (tensor) Coords of ground truth for each prior in point-form |
||||
|
Shape: [num_priors, 4]. |
||||
|
priors: (tensor) Prior boxes in center-offset form |
||||
|
Shape: [num_priors,4]. |
||||
|
variances: (list[float]) Variances of priorboxes |
||||
|
Return: |
||||
|
encoded boxes (tensor), Shape: [num_priors, 4] |
||||
|
""" |
||||
|
|
||||
|
# dist b/t match center and prior's center |
||||
|
g_cxcy = (matched[:, :2] + matched[:, 2:])/2 - priors[:, :2] |
||||
|
# encode variance |
||||
|
g_cxcy /= (variances[0] * priors[:, 2:]) |
||||
|
# match wh / prior wh |
||||
|
g_wh = (matched[:, 2:] - matched[:, :2]) / priors[:, 2:] |
||||
|
g_wh = torch.log(g_wh) / variances[1] |
||||
|
# return target for smooth_l1_loss |
||||
|
return torch.cat([g_cxcy, g_wh], 1) # [num_priors,4] |
||||
|
|
||||
|
def encode_landm(matched, priors, variances): |
||||
|
"""Encode the variances from the priorbox layers into the ground truth boxes |
||||
|
we have matched (based on jaccard overlap) with the prior boxes. |
||||
|
Args: |
||||
|
matched: (tensor) Coords of ground truth for each prior in point-form |
||||
|
Shape: [num_priors, 10]. |
||||
|
priors: (tensor) Prior boxes in center-offset form |
||||
|
Shape: [num_priors,4]. |
||||
|
variances: (list[float]) Variances of priorboxes |
||||
|
Return: |
||||
|
encoded landm (tensor), Shape: [num_priors, 10] |
||||
|
""" |
||||
|
|
||||
|
# dist b/t match center and prior's center |
||||
|
matched = torch.reshape(matched, (matched.size(0), 5, 2)) |
||||
|
priors_cx = priors[:, 0].unsqueeze(1).expand(matched.size(0), 5).unsqueeze(2) |
||||
|
priors_cy = priors[:, 1].unsqueeze(1).expand(matched.size(0), 5).unsqueeze(2) |
||||
|
priors_w = priors[:, 2].unsqueeze(1).expand(matched.size(0), 5).unsqueeze(2) |
||||
|
priors_h = priors[:, 3].unsqueeze(1).expand(matched.size(0), 5).unsqueeze(2) |
||||
|
priors = torch.cat([priors_cx, priors_cy, priors_w, priors_h], dim=2) |
||||
|
g_cxcy = matched[:, :, :2] - priors[:, :, :2] |
||||
|
# encode variance |
||||
|
g_cxcy /= (variances[0] * priors[:, :, 2:]) |
||||
|
# g_cxcy /= priors[:, :, 2:] |
||||
|
g_cxcy = g_cxcy.reshape(g_cxcy.size(0), -1) |
||||
|
# return target for smooth_l1_loss |
||||
|
return g_cxcy |
||||
|
|
||||
|
|
||||
|
# Adapted from https://github.com/Hakuyume/chainer-ssd |
||||
|
def decode(loc, priors, variances): |
||||
|
"""Decode locations from predictions using priors to undo |
||||
|
the encoding we did for offset regression at train time. |
||||
|
Args: |
||||
|
loc (tensor): location predictions for loc layers, |
||||
|
Shape: [num_priors,4] |
||||
|
priors (tensor): Prior boxes in center-offset form. |
||||
|
Shape: [num_priors,4]. |
||||
|
variances: (list[float]) Variances of priorboxes |
||||
|
Return: |
||||
|
decoded bounding box predictions |
||||
|
""" |
||||
|
|
||||
|
boxes = torch.cat(( |
||||
|
priors[:, :2] + loc[:, :2] * variances[0] * priors[:, 2:], |
||||
|
priors[:, 2:] * torch.exp(loc[:, 2:] * variances[1])), 1) |
||||
|
boxes[:, :2] -= boxes[:, 2:] / 2 |
||||
|
boxes[:, 2:] += boxes[:, :2] |
||||
|
return boxes |
||||
|
|
||||
|
def decode_landm(pre, priors, variances): |
||||
|
"""Decode landm from predictions using priors to undo |
||||
|
the encoding we did for offset regression at train time. |
||||
|
Args: |
||||
|
pre (tensor): landm predictions for loc layers, |
||||
|
Shape: [num_priors,10] |
||||
|
priors (tensor): Prior boxes in center-offset form. |
||||
|
Shape: [num_priors,4]. |
||||
|
variances: (list[float]) Variances of priorboxes |
||||
|
Return: |
||||
|
decoded landm predictions |
||||
|
""" |
||||
|
landms = torch.cat((priors[:, :2] + pre[:, :2] * variances[0] * priors[:, 2:], |
||||
|
priors[:, :2] + pre[:, 2:4] * variances[0] * priors[:, 2:], |
||||
|
priors[:, :2] + pre[:, 4:6] * variances[0] * priors[:, 2:], |
||||
|
priors[:, :2] + pre[:, 6:8] * variances[0] * priors[:, 2:], |
||||
|
priors[:, :2] + pre[:, 8:10] * variances[0] * priors[:, 2:], |
||||
|
), dim=1) |
||||
|
return landms |
||||
|
|
||||
|
|
||||
|
def log_sum_exp(x): |
||||
|
"""Utility function for computing log_sum_exp while determining |
||||
|
This will be used to determine unaveraged confidence loss across |
||||
|
all examples in a batch. |
||||
|
Args: |
||||
|
x (Variable(tensor)): conf_preds from conf layers |
||||
|
""" |
||||
|
x_max = x.data.max() |
||||
|
return torch.log(torch.sum(torch.exp(x-x_max), 1, keepdim=True)) + x_max |
||||
|
|
||||
|
|
||||
|
# Original author: Francisco Massa: |
||||
|
# https://github.com/fmassa/object-detection.torch |
||||
|
# Ported to PyTorch by Max deGroot (02/01/2017) |
||||
|
def nms(boxes, scores, overlap=0.5, top_k=200): |
||||
|
"""Apply non-maximum suppression at test time to avoid detecting too many |
||||
|
overlapping bounding boxes for a given object. |
||||
|
Args: |
||||
|
boxes: (tensor) The location preds for the img, Shape: [num_priors,4]. |
||||
|
scores: (tensor) The class predscores for the img, Shape:[num_priors]. |
||||
|
overlap: (float) The overlap thresh for suppressing unnecessary boxes. |
||||
|
top_k: (int) The Maximum number of box preds to consider. |
||||
|
Return: |
||||
|
The indices of the kept boxes with respect to num_priors. |
||||
|
""" |
||||
|
|
||||
|
keep = torch.Tensor(scores.size(0)).fill_(0).long() |
||||
|
if boxes.numel() == 0: |
||||
|
return keep |
||||
|
x1 = boxes[:, 0] |
||||
|
y1 = boxes[:, 1] |
||||
|
x2 = boxes[:, 2] |
||||
|
y2 = boxes[:, 3] |
||||
|
area = torch.mul(x2 - x1, y2 - y1) |
||||
|
v, idx = scores.sort(0) # sort in ascending order |
||||
|
# I = I[v >= 0.01] |
||||
|
idx = idx[-top_k:] # indices of the top-k largest vals |
||||
|
xx1 = boxes.new() |
||||
|
yy1 = boxes.new() |
||||
|
xx2 = boxes.new() |
||||
|
yy2 = boxes.new() |
||||
|
w = boxes.new() |
||||
|
h = boxes.new() |
||||
|
|
||||
|
# keep = torch.Tensor() |
||||
|
count = 0 |
||||
|
while idx.numel() > 0: |
||||
|
i = idx[-1] # index of current largest val |
||||
|
# keep.append(i) |
||||
|
keep[count] = i |
||||
|
count += 1 |
||||
|
if idx.size(0) == 1: |
||||
|
break |
||||
|
idx = idx[:-1] # remove kept element from view |
||||
|
# load bboxes of next highest vals |
||||
|
torch.index_select(x1, 0, idx, out=xx1) |
||||
|
torch.index_select(y1, 0, idx, out=yy1) |
||||
|
torch.index_select(x2, 0, idx, out=xx2) |
||||
|
torch.index_select(y2, 0, idx, out=yy2) |
||||
|
# store element-wise max with next highest score |
||||
|
xx1 = torch.clamp(xx1, min=x1[i]) |
||||
|
yy1 = torch.clamp(yy1, min=y1[i]) |
||||
|
xx2 = torch.clamp(xx2, max=x2[i]) |
||||
|
yy2 = torch.clamp(yy2, max=y2[i]) |
||||
|
w.resize_as_(xx2) |
||||
|
h.resize_as_(yy2) |
||||
|
w = xx2 - xx1 |
||||
|
h = yy2 - yy1 |
||||
|
# check sizes of xx1 and xx2.. after each iteration |
||||
|
w = torch.clamp(w, min=0.0) |
||||
|
h = torch.clamp(h, min=0.0) |
||||
|
inter = w*h |
||||
|
# IoU = i / (area(a) + area(b) - i) |
||||
|
rem_areas = torch.index_select(area, 0, idx) # load remaining areas) |
||||
|
union = (rem_areas - inter) + area[i] |
||||
|
IoU = inter/union # store result in iou |
||||
|
# keep only elements with an IoU <= overlap |
||||
|
idx = idx[IoU.le(overlap)] |
||||
|
return keep, count |
@ -0,0 +1,122 @@ |
|||||
|
import torch |
||||
|
import torch.nn as nn |
||||
|
import torch.nn.functional as F |
||||
|
from torch.autograd import Variable |
||||
|
from .box_utils import match, log_sum_exp |
||||
|
|
||||
|
class MultiBoxLoss(nn.Module): |
||||
|
"""SSD Weighted Loss Function |
||||
|
Compute Targets: |
||||
|
1) Produce Confidence Target Indices by matching ground truth boxes |
||||
|
with (default) 'priorboxes' that have jaccard index > threshold parameter |
||||
|
(default threshold: 0.5). |
||||
|
2) Produce localization target by 'encoding' variance into offsets of ground |
||||
|
truth boxes and their matched 'priorboxes'. |
||||
|
3) Hard negative mining to filter the excessive number of negative examples |
||||
|
that comes with using a large number of default bounding boxes. |
||||
|
(default negative:positive ratio 3:1) |
||||
|
Objective Loss: |
||||
|
L(x,c,l,g) = (Lconf(x, c) + αLloc(x,l,g)) / N |
||||
|
Where, Lconf is the CrossEntropy Loss and Lloc is the SmoothL1 Loss |
||||
|
weighted by α which is set to 1 by cross val. |
||||
|
Args: |
||||
|
c: class confidences, |
||||
|
l: predicted boxes, |
||||
|
g: ground truth boxes |
||||
|
N: number of matched default boxes |
||||
|
See: https://arxiv.org/pdf/1512.02325.pdf for more details. |
||||
|
""" |
||||
|
|
||||
|
def __init__(self, num_classes, overlap_thresh, prior_for_matching, bkg_label, neg_mining, neg_pos, neg_overlap, encode_target): |
||||
|
super(MultiBoxLoss, self).__init__() |
||||
|
self.num_classes = num_classes |
||||
|
self.threshold = overlap_thresh |
||||
|
self.background_label = bkg_label |
||||
|
self.encode_target = encode_target |
||||
|
self.use_prior_for_matching = prior_for_matching |
||||
|
self.do_neg_mining = neg_mining |
||||
|
self.negpos_ratio = neg_pos |
||||
|
self.neg_overlap = neg_overlap |
||||
|
self.variance = [0.1, 0.2] |
||||
|
self.GPU = False |
||||
|
|
||||
|
def forward(self, predictions, priors, targets): |
||||
|
"""Multibox Loss |
||||
|
Args: |
||||
|
predictions (tuple): A tuple containing loc preds, conf preds, |
||||
|
and prior boxes from SSD net. |
||||
|
conf shape: torch.size(batch_size,num_priors,num_classes) |
||||
|
loc shape: torch.size(batch_size,num_priors,4) |
||||
|
priors shape: torch.size(num_priors,4) |
||||
|
ground_truth (tensor): Ground truth boxes and labels for a batch, |
||||
|
shape: [batch_size,num_objs,5] (last idx is the label). |
||||
|
""" |
||||
|
|
||||
|
loc_data, conf_data, landm_data = predictions |
||||
|
priors = priors |
||||
|
num = loc_data.size(0) |
||||
|
num_priors = (priors.size(0)) |
||||
|
|
||||
|
# match priors (default boxes) and ground truth boxes |
||||
|
loc_t = torch.Tensor(num, num_priors, 4) |
||||
|
landm_t = torch.Tensor(num, num_priors, 10) |
||||
|
conf_t = torch.LongTensor(num, num_priors) |
||||
|
for idx in range(num): |
||||
|
truths = targets[idx][:, :4].data |
||||
|
labels = targets[idx][:, -1].data |
||||
|
landms = targets[idx][:, 4:14].data |
||||
|
defaults = priors.data |
||||
|
match(self.threshold, truths, defaults, self.variance, labels, landms, loc_t, conf_t, landm_t, idx) |
||||
|
if self.GPU: |
||||
|
loc_t = loc_t.cuda() |
||||
|
conf_t = conf_t.cuda() |
||||
|
landm_t = landm_t.cuda() |
||||
|
|
||||
|
zeros = torch.tensor(0).cuda() |
||||
|
# landm Loss (Smooth L1) |
||||
|
# Shape: [batch,num_priors,10] |
||||
|
pos1 = conf_t > zeros |
||||
|
num_pos_landm = pos1.long().sum(1, keepdim=True) |
||||
|
N1 = max(num_pos_landm.data.sum().float(), 1) |
||||
|
pos_idx1 = pos1.unsqueeze(pos1.dim()).expand_as(landm_data) |
||||
|
landm_p = landm_data[pos_idx1].view(-1, 10) |
||||
|
landm_t = landm_t[pos_idx1].view(-1, 10) |
||||
|
loss_landm = F.smooth_l1_loss(landm_p, landm_t, reduction='sum') |
||||
|
|
||||
|
|
||||
|
pos = conf_t != zeros |
||||
|
conf_t[pos] = 1 |
||||
|
|
||||
|
# Localization Loss (Smooth L1) |
||||
|
# Shape: [batch,num_priors,4] |
||||
|
pos_idx = pos.unsqueeze(pos.dim()).expand_as(loc_data) |
||||
|
loc_p = loc_data[pos_idx].view(-1, 4) |
||||
|
loc_t = loc_t[pos_idx].view(-1, 4) |
||||
|
loss_l = F.smooth_l1_loss(loc_p, loc_t, reduction='sum') |
||||
|
|
||||
|
# Compute max conf across batch for hard negative mining |
||||
|
batch_conf = conf_data.view(-1, self.num_classes) |
||||
|
loss_c = log_sum_exp(batch_conf) - batch_conf.gather(1, conf_t.view(-1, 1)) |
||||
|
|
||||
|
# Hard Negative Mining |
||||
|
loss_c[pos.view(-1, 1)] = 0 # filter out pos boxes for now |
||||
|
loss_c = loss_c.view(num, -1) |
||||
|
_, loss_idx = loss_c.sort(1, descending=True) |
||||
|
_, idx_rank = loss_idx.sort(1) |
||||
|
num_pos = pos.long().sum(1, keepdim=True) |
||||
|
num_neg = torch.clamp(self.negpos_ratio*num_pos, max=pos.size(1)-1) |
||||
|
neg = idx_rank < num_neg.expand_as(idx_rank) |
||||
|
|
||||
|
# Confidence Loss Including Positive and Negative Examples |
||||
|
pos_idx = pos.unsqueeze(2).expand_as(conf_data) |
||||
|
neg_idx = neg.unsqueeze(2).expand_as(conf_data) |
||||
|
conf_p = conf_data[(pos_idx+neg_idx).gt(0)].view(-1,self.num_classes) |
||||
|
targets_weighted = conf_t[(pos+neg).gt(0)] |
||||
|
loss_c = F.cross_entropy(conf_p, targets_weighted, reduction='sum') |
||||
|
|
||||
|
# Sum of losses: L(x,c,l,g) = (Lconf(x, c) + αLloc(x,l,g)) / N N = max(num_pos.data.sum().float(), 1) |
||||
|
loss_l /= N |
||||
|
loss_c /= N |
||||
|
loss_landm /= N1 |
||||
|
|
||||
|
return loss_l, loss_c, loss_landm |
@ -0,0 +1,25 @@ |
|||||
|
device: |
||||
|
device_str: null |
||||
|
n_gpu: -1 |
||||
|
sync_bn: true |
||||
|
metrics: |
||||
|
metric: Accuracy |
||||
|
train: |
||||
|
batch_size: 32 |
||||
|
overwrite_output_dir: true |
||||
|
epoch_num: 2 |
||||
|
learning: |
||||
|
optimizer: |
||||
|
name_: SGD |
||||
|
lr: 0.04 |
||||
|
momentum: 0.001 |
||||
|
loss: |
||||
|
name_: CrossEntropyLoss |
||||
|
ignore_index: -1 |
||||
|
logging: |
||||
|
print_steps: 2 |
||||
|
#learning: |
||||
|
# optimizer: |
||||
|
# name_: Adam |
||||
|
# lr: 0.02 |
||||
|
# eps: 0.001 |
@ -0,0 +1,84 @@ |
|||||
|
import numpy as np |
||||
|
from torch.optim import AdamW |
||||
|
from torchvision import transforms |
||||
|
from torchvision.transforms import RandomResizedCrop, Lambda |
||||
|
from towhee.trainer.modelcard import ModelCard |
||||
|
|
||||
|
from towhee.trainer.training_config import TrainingConfig |
||||
|
from towhee.trainer.dataset import get_dataset |
||||
|
from resnet_image_embedding import ResnetImageEmbedding |
||||
|
from towhee.types import Image |
||||
|
from towhee.trainer.training_config import dump_default_yaml |
||||
|
from PIL import Image as PILImage |
||||
|
from timm.models.resnet import ResNet |
||||
|
from torch import nn |
||||
|
|
||||
|
if __name__ == '__main__': |
||||
|
dump_default_yaml(yaml_path='default_config.yaml') |
||||
|
# img = torch.rand([1, 3, 224, 224]) |
||||
|
img_path = './ILSVRC2012_val_00049771.JPEG' |
||||
|
# # logo_path = os.path.join(Path(__file__).parent.parent.parent.parent.resolve(), 'towhee_logo.png') |
||||
|
img = PILImage.open(img_path) |
||||
|
img_bytes = img.tobytes() |
||||
|
img_width = img.width |
||||
|
img_height = img.height |
||||
|
img_channel = len(img.split()) |
||||
|
img_mode = img.mode |
||||
|
img_array = np.array(img) |
||||
|
array_size = np.array(img).shape |
||||
|
towhee_img = Image(img_bytes, img_width, img_height, img_channel, img_mode, img_array) |
||||
|
|
||||
|
op = ResnetImageEmbedding('resnet34') |
||||
|
# op.model_card = ModelCard(model_details="resnet test modelcard", training_data="use resnet test data") |
||||
|
# old_out = op(towhee_img) |
||||
|
# print(old_out.feature_vector[0]) |
||||
|
|
||||
|
training_config = TrainingConfig() |
||||
|
yaml_path = 'resnet_training_yaml.yaml' |
||||
|
# dump_default_yaml(yaml_path=yaml_path) |
||||
|
training_config.load_from_yaml(yaml_path) |
||||
|
# output_dir='./temp_output', |
||||
|
# overwrite_output_dir=True, |
||||
|
# epoch_num=2, |
||||
|
# per_gpu_train_batch_size=16, |
||||
|
# prediction_loss_only=True, |
||||
|
# metric='Accuracy' |
||||
|
# # device_str='cuda', |
||||
|
# # n_gpu=4 |
||||
|
# ) |
||||
|
|
||||
|
mnist_transform = transforms.Compose([transforms.ToTensor(), |
||||
|
RandomResizedCrop(224), |
||||
|
Lambda(lambda x: x.repeat(3, 1, 1)), |
||||
|
transforms.Normalize(mean=[0.5], std=[0.5])]) |
||||
|
train_data = get_dataset('mnist', transform=mnist_transform, download=True, root='data', train=True) |
||||
|
eval_data = get_dataset('mnist', transform=mnist_transform, download=True, root='data', train=False) |
||||
|
# fake_transform = transforms.Compose([transforms.ToTensor(), |
||||
|
# RandomResizedCrop(224),]) |
||||
|
# train_data = get_dataset('fake', size=20, transform=fake_transform) |
||||
|
|
||||
|
op.change_before_train(10) |
||||
|
trainer = op.setup_trainer() |
||||
|
# my_optimimzer = AdamW(op.get_model().parameters(), lr=0.002, betas=(0.91, 0.98), eps=1e-08, weight_decay=0.01, amsgrad=False) |
||||
|
# op.setup_trainer() |
||||
|
|
||||
|
# trainer.add_callback() |
||||
|
# trainer.set_optimizer() |
||||
|
|
||||
|
# op.trainer.set_optimizer(my_optimimzer) |
||||
|
# trainer.configs.save_to_yaml('changed_optimizer_yaml.yaml') |
||||
|
|
||||
|
# my_loss = nn.BCELoss() |
||||
|
# trainer.set_loss(my_loss, 'my_loss111') |
||||
|
# trainer.configs.save_to_yaml('chaned_loss_yaml.yaml') |
||||
|
# op.trainer._create_optimizer() |
||||
|
# op.trainer.set_optimizer() |
||||
|
op.train(training_config, train_dataset=train_data, eval_dataset=eval_data) |
||||
|
# training_config.num_epoch = 3 |
||||
|
# op.train(training_config, train_dataset=train_data, resume_checkpoint_path=training_config.output_dir + '/epoch_2') |
||||
|
|
||||
|
# op.save('./test_save') |
||||
|
# op.load('./test_save') |
||||
|
# new_out = op(towhee_img) |
||||
|
|
||||
|
# assert (new_out[0]!=old_out[0]).all() |
Loading…
Reference in new issue