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