489 lines
24 KiB
Python
489 lines
24 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding:utf-8 -*-
|
|
# Copyright (c) Megvii, Inc. and its affiliates.
|
|
import math
|
|
from copy import deepcopy
|
|
from functools import partial
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
|
|
class IOUloss(nn.Module):
|
|
def __init__(self, reduction="none", loss_type="iou"):
|
|
super(IOUloss, self).__init__()
|
|
self.reduction = reduction
|
|
self.loss_type = loss_type
|
|
|
|
def forward(self, pred, target):
|
|
assert pred.shape[0] == target.shape[0]
|
|
|
|
pred = pred.view(-1, 4)
|
|
target = target.view(-1, 4)
|
|
tl = torch.max(
|
|
(pred[:, :2] - pred[:, 2:] / 2), (target[:, :2] - target[:, 2:] / 2)
|
|
)
|
|
br = torch.min(
|
|
(pred[:, :2] + pred[:, 2:] / 2), (target[:, :2] + target[:, 2:] / 2)
|
|
)
|
|
|
|
area_p = torch.prod(pred[:, 2:], 1)
|
|
area_g = torch.prod(target[:, 2:], 1)
|
|
|
|
en = (tl < br).type(tl.type()).prod(dim=1)
|
|
area_i = torch.prod(br - tl, 1) * en
|
|
area_u = area_p + area_g - area_i
|
|
iou = (area_i) / (area_u + 1e-16)
|
|
|
|
if self.loss_type == "iou":
|
|
loss = 1 - iou ** 2
|
|
elif self.loss_type == "giou":
|
|
c_tl = torch.min(
|
|
(pred[:, :2] - pred[:, 2:] / 2), (target[:, :2] - target[:, 2:] / 2)
|
|
)
|
|
c_br = torch.max(
|
|
(pred[:, :2] + pred[:, 2:] / 2), (target[:, :2] + target[:, 2:] / 2)
|
|
)
|
|
area_c = torch.prod(c_br - c_tl, 1)
|
|
giou = iou - (area_c - area_u) / area_c.clamp(1e-16)
|
|
loss = 1 - giou.clamp(min=-1.0, max=1.0)
|
|
|
|
if self.reduction == "mean":
|
|
loss = loss.mean()
|
|
elif self.reduction == "sum":
|
|
loss = loss.sum()
|
|
|
|
return loss
|
|
|
|
class YOLOLoss(nn.Module):
|
|
def __init__(self, num_classes, fp16, strides=[8, 16, 32]):
|
|
super().__init__()
|
|
self.num_classes = num_classes
|
|
self.strides = strides
|
|
|
|
self.bcewithlog_loss = nn.BCEWithLogitsLoss(reduction="none")
|
|
self.iou_loss = IOUloss(reduction="none")
|
|
self.grids = [torch.zeros(1)] * len(strides)
|
|
self.fp16 = fp16
|
|
|
|
def forward(self, inputs, labels=None):
|
|
outputs = []
|
|
x_shifts = []
|
|
y_shifts = []
|
|
expanded_strides = []
|
|
|
|
#-----------------------------------------------#
|
|
# inputs [[batch_size, num_classes + 5, 20, 20]
|
|
# [batch_size, num_classes + 5, 40, 40]
|
|
# [batch_size, num_classes + 5, 80, 80]]
|
|
# outputs [[batch_size, 400, num_classes + 5]
|
|
# [batch_size, 1600, num_classes + 5]
|
|
# [batch_size, 6400, num_classes + 5]]
|
|
# x_shifts [[batch_size, 400]
|
|
# [batch_size, 1600]
|
|
# [batch_size, 6400]]
|
|
#-----------------------------------------------#
|
|
for k, (stride, output) in enumerate(zip(self.strides, inputs)):
|
|
output, grid = self.get_output_and_grid(output, k, stride)
|
|
x_shifts.append(grid[:, :, 0])
|
|
y_shifts.append(grid[:, :, 1])
|
|
expanded_strides.append(torch.ones_like(grid[:, :, 0]) * stride)
|
|
outputs.append(output)
|
|
|
|
return self.get_losses(x_shifts, y_shifts, expanded_strides, labels, torch.cat(outputs, 1))
|
|
|
|
def get_output_and_grid(self, output, k, stride):
|
|
grid = self.grids[k]
|
|
hsize, wsize = output.shape[-2:]
|
|
if grid.shape[2:4] != output.shape[2:4]:
|
|
yv, xv = torch.meshgrid([torch.arange(hsize), torch.arange(wsize)])
|
|
grid = torch.stack((xv, yv), 2).view(1, hsize, wsize, 2).type(output.type())
|
|
self.grids[k] = grid
|
|
grid = grid.view(1, -1, 2)
|
|
|
|
output = output.flatten(start_dim=2).permute(0, 2, 1)
|
|
output[..., :2] = (output[..., :2] + grid.type_as(output)) * stride
|
|
output[..., 2:4] = torch.exp(output[..., 2:4]) * stride
|
|
return output, grid
|
|
|
|
def get_losses(self, x_shifts, y_shifts, expanded_strides, labels, outputs):
|
|
#-----------------------------------------------#
|
|
# [batch, n_anchors_all, 4]
|
|
#-----------------------------------------------#
|
|
bbox_preds = outputs[:, :, :4]
|
|
#-----------------------------------------------#
|
|
# [batch, n_anchors_all, 1]
|
|
#-----------------------------------------------#
|
|
obj_preds = outputs[:, :, 4:5]
|
|
#-----------------------------------------------#
|
|
# [batch, n_anchors_all, n_cls]
|
|
#-----------------------------------------------#
|
|
cls_preds = outputs[:, :, 5:]
|
|
|
|
total_num_anchors = outputs.shape[1]
|
|
#-----------------------------------------------#
|
|
# x_shifts [1, n_anchors_all]
|
|
# y_shifts [1, n_anchors_all]
|
|
# expanded_strides [1, n_anchors_all]
|
|
#-----------------------------------------------#
|
|
x_shifts = torch.cat(x_shifts, 1).type_as(outputs)
|
|
y_shifts = torch.cat(y_shifts, 1).type_as(outputs)
|
|
expanded_strides = torch.cat(expanded_strides, 1).type_as(outputs)
|
|
|
|
cls_targets = []
|
|
reg_targets = []
|
|
obj_targets = []
|
|
fg_masks = []
|
|
|
|
num_fg = 0.0
|
|
for batch_idx in range(outputs.shape[0]):
|
|
num_gt = len(labels[batch_idx])
|
|
if num_gt == 0:
|
|
cls_target = outputs.new_zeros((0, self.num_classes))
|
|
reg_target = outputs.new_zeros((0, 4))
|
|
obj_target = outputs.new_zeros((total_num_anchors, 1))
|
|
fg_mask = outputs.new_zeros(total_num_anchors).bool()
|
|
else:
|
|
#-----------------------------------------------#
|
|
# gt_bboxes_per_image [num_gt, num_classes]
|
|
# gt_classes [num_gt]
|
|
# bboxes_preds_per_image [n_anchors_all, 4]
|
|
# cls_preds_per_image [n_anchors_all, num_classes]
|
|
# obj_preds_per_image [n_anchors_all, 1]
|
|
#-----------------------------------------------#
|
|
gt_bboxes_per_image = labels[batch_idx][..., :4].type_as(outputs)
|
|
gt_classes = labels[batch_idx][..., 4].type_as(outputs)
|
|
bboxes_preds_per_image = bbox_preds[batch_idx]
|
|
cls_preds_per_image = cls_preds[batch_idx]
|
|
obj_preds_per_image = obj_preds[batch_idx]
|
|
|
|
gt_matched_classes, fg_mask, pred_ious_this_matching, matched_gt_inds, num_fg_img = self.get_assignments(
|
|
num_gt, total_num_anchors, gt_bboxes_per_image, gt_classes, bboxes_preds_per_image, cls_preds_per_image, obj_preds_per_image,
|
|
expanded_strides, x_shifts, y_shifts,
|
|
)
|
|
torch.cuda.empty_cache()
|
|
num_fg += num_fg_img
|
|
cls_target = F.one_hot(gt_matched_classes.to(torch.int64), self.num_classes).float() * pred_ious_this_matching.unsqueeze(-1)
|
|
obj_target = fg_mask.unsqueeze(-1)
|
|
reg_target = gt_bboxes_per_image[matched_gt_inds]
|
|
cls_targets.append(cls_target)
|
|
reg_targets.append(reg_target)
|
|
obj_targets.append(obj_target.type(cls_target.type()))
|
|
fg_masks.append(fg_mask)
|
|
|
|
cls_targets = torch.cat(cls_targets, 0)
|
|
reg_targets = torch.cat(reg_targets, 0)
|
|
obj_targets = torch.cat(obj_targets, 0)
|
|
fg_masks = torch.cat(fg_masks, 0)
|
|
|
|
num_fg = max(num_fg, 1)
|
|
loss_iou = (self.iou_loss(bbox_preds.view(-1, 4)[fg_masks], reg_targets)).sum()
|
|
loss_obj = (self.bcewithlog_loss(obj_preds.view(-1, 1), obj_targets)).sum()
|
|
loss_cls = (self.bcewithlog_loss(cls_preds.view(-1, self.num_classes)[fg_masks], cls_targets)).sum()
|
|
reg_weight = 5.0
|
|
loss = reg_weight * loss_iou + loss_obj + loss_cls
|
|
|
|
return loss / num_fg
|
|
|
|
@torch.no_grad()
|
|
def get_assignments(self, num_gt, total_num_anchors, gt_bboxes_per_image, gt_classes, bboxes_preds_per_image, cls_preds_per_image, obj_preds_per_image, expanded_strides, x_shifts, y_shifts):
|
|
#-------------------------------------------------------#
|
|
# fg_mask [n_anchors_all]
|
|
# is_in_boxes_and_center [num_gt, len(fg_mask)]
|
|
#-------------------------------------------------------#
|
|
fg_mask, is_in_boxes_and_center = self.get_in_boxes_info(gt_bboxes_per_image, expanded_strides, x_shifts, y_shifts, total_num_anchors, num_gt)
|
|
|
|
#-------------------------------------------------------#
|
|
# fg_mask [n_anchors_all]
|
|
# bboxes_preds_per_image [fg_mask, 4]
|
|
# cls_preds_ [fg_mask, num_classes]
|
|
# obj_preds_ [fg_mask, 1]
|
|
#-------------------------------------------------------#
|
|
bboxes_preds_per_image = bboxes_preds_per_image[fg_mask]
|
|
cls_preds_ = cls_preds_per_image[fg_mask]
|
|
obj_preds_ = obj_preds_per_image[fg_mask]
|
|
num_in_boxes_anchor = bboxes_preds_per_image.shape[0]
|
|
|
|
#-------------------------------------------------------#
|
|
# pair_wise_ious [num_gt, fg_mask]
|
|
#-------------------------------------------------------#
|
|
pair_wise_ious = self.bboxes_iou(gt_bboxes_per_image, bboxes_preds_per_image, False)
|
|
pair_wise_ious_loss = -torch.log(pair_wise_ious + 1e-8)
|
|
|
|
#-------------------------------------------------------#
|
|
# cls_preds_ [num_gt, fg_mask, num_classes]
|
|
# gt_cls_per_image [num_gt, fg_mask, num_classes]
|
|
#-------------------------------------------------------#
|
|
if self.fp16:
|
|
with torch.cuda.amp.autocast(enabled=False):
|
|
cls_preds_ = cls_preds_.float().unsqueeze(0).repeat(num_gt, 1, 1).sigmoid_() * obj_preds_.unsqueeze(0).repeat(num_gt, 1, 1).sigmoid_()
|
|
gt_cls_per_image = F.one_hot(gt_classes.to(torch.int64), self.num_classes).float().unsqueeze(1).repeat(1, num_in_boxes_anchor, 1)
|
|
pair_wise_cls_loss = F.binary_cross_entropy(cls_preds_.sqrt_(), gt_cls_per_image, reduction="none").sum(-1)
|
|
else:
|
|
cls_preds_ = cls_preds_.float().unsqueeze(0).repeat(num_gt, 1, 1).sigmoid_() * obj_preds_.unsqueeze(0).repeat(num_gt, 1, 1).sigmoid_()
|
|
gt_cls_per_image = F.one_hot(gt_classes.to(torch.int64), self.num_classes).float().unsqueeze(1).repeat(1, num_in_boxes_anchor, 1)
|
|
pair_wise_cls_loss = F.binary_cross_entropy(cls_preds_.sqrt_(), gt_cls_per_image, reduction="none").sum(-1)
|
|
del cls_preds_
|
|
|
|
cost = pair_wise_cls_loss + 3.0 * pair_wise_ious_loss + 100000.0 * (~is_in_boxes_and_center).float()
|
|
|
|
num_fg, gt_matched_classes, pred_ious_this_matching, matched_gt_inds = self.dynamic_k_matching(cost, pair_wise_ious, gt_classes, num_gt, fg_mask)
|
|
del pair_wise_cls_loss, cost, pair_wise_ious, pair_wise_ious_loss
|
|
return gt_matched_classes, fg_mask, pred_ious_this_matching, matched_gt_inds, num_fg
|
|
|
|
def bboxes_iou(self, bboxes_a, bboxes_b, xyxy=True):
|
|
if bboxes_a.shape[1] != 4 or bboxes_b.shape[1] != 4:
|
|
raise IndexError
|
|
|
|
if xyxy:
|
|
tl = torch.max(bboxes_a[:, None, :2], bboxes_b[:, :2])
|
|
br = torch.min(bboxes_a[:, None, 2:], bboxes_b[:, 2:])
|
|
area_a = torch.prod(bboxes_a[:, 2:] - bboxes_a[:, :2], 1)
|
|
area_b = torch.prod(bboxes_b[:, 2:] - bboxes_b[:, :2], 1)
|
|
else:
|
|
tl = torch.max(
|
|
(bboxes_a[:, None, :2] - bboxes_a[:, None, 2:] / 2),
|
|
(bboxes_b[:, :2] - bboxes_b[:, 2:] / 2),
|
|
)
|
|
br = torch.min(
|
|
(bboxes_a[:, None, :2] + bboxes_a[:, None, 2:] / 2),
|
|
(bboxes_b[:, :2] + bboxes_b[:, 2:] / 2),
|
|
)
|
|
|
|
area_a = torch.prod(bboxes_a[:, 2:], 1)
|
|
area_b = torch.prod(bboxes_b[:, 2:], 1)
|
|
en = (tl < br).type(tl.type()).prod(dim=2)
|
|
area_i = torch.prod(br - tl, 2) * en
|
|
return area_i / (area_a[:, None] + area_b - area_i)
|
|
|
|
def get_in_boxes_info(self, gt_bboxes_per_image, expanded_strides, x_shifts, y_shifts, total_num_anchors, num_gt, center_radius = 2.5):
|
|
#-------------------------------------------------------#
|
|
# expanded_strides_per_image [n_anchors_all]
|
|
# x_centers_per_image [num_gt, n_anchors_all]
|
|
# x_centers_per_image [num_gt, n_anchors_all]
|
|
#-------------------------------------------------------#
|
|
expanded_strides_per_image = expanded_strides[0]
|
|
x_centers_per_image = ((x_shifts[0] + 0.5) * expanded_strides_per_image).unsqueeze(0).repeat(num_gt, 1)
|
|
y_centers_per_image = ((y_shifts[0] + 0.5) * expanded_strides_per_image).unsqueeze(0).repeat(num_gt, 1)
|
|
|
|
#-------------------------------------------------------#
|
|
# gt_bboxes_per_image_x [num_gt, n_anchors_all]
|
|
#-------------------------------------------------------#
|
|
gt_bboxes_per_image_l = (gt_bboxes_per_image[:, 0] - 0.5 * gt_bboxes_per_image[:, 2]).unsqueeze(1).repeat(1, total_num_anchors)
|
|
gt_bboxes_per_image_r = (gt_bboxes_per_image[:, 0] + 0.5 * gt_bboxes_per_image[:, 2]).unsqueeze(1).repeat(1, total_num_anchors)
|
|
gt_bboxes_per_image_t = (gt_bboxes_per_image[:, 1] - 0.5 * gt_bboxes_per_image[:, 3]).unsqueeze(1).repeat(1, total_num_anchors)
|
|
gt_bboxes_per_image_b = (gt_bboxes_per_image[:, 1] + 0.5 * gt_bboxes_per_image[:, 3]).unsqueeze(1).repeat(1, total_num_anchors)
|
|
|
|
#-------------------------------------------------------#
|
|
# bbox_deltas [num_gt, n_anchors_all, 4]
|
|
#-------------------------------------------------------#
|
|
b_l = x_centers_per_image - gt_bboxes_per_image_l
|
|
b_r = gt_bboxes_per_image_r - x_centers_per_image
|
|
b_t = y_centers_per_image - gt_bboxes_per_image_t
|
|
b_b = gt_bboxes_per_image_b - y_centers_per_image
|
|
bbox_deltas = torch.stack([b_l, b_t, b_r, b_b], 2)
|
|
|
|
#-------------------------------------------------------#
|
|
# is_in_boxes [num_gt, n_anchors_all]
|
|
# is_in_boxes_all [n_anchors_all]
|
|
#-------------------------------------------------------#
|
|
is_in_boxes = bbox_deltas.min(dim=-1).values > 0.0
|
|
is_in_boxes_all = is_in_boxes.sum(dim=0) > 0
|
|
|
|
gt_bboxes_per_image_l = (gt_bboxes_per_image[:, 0]).unsqueeze(1).repeat(1, total_num_anchors) - center_radius * expanded_strides_per_image.unsqueeze(0)
|
|
gt_bboxes_per_image_r = (gt_bboxes_per_image[:, 0]).unsqueeze(1).repeat(1, total_num_anchors) + center_radius * expanded_strides_per_image.unsqueeze(0)
|
|
gt_bboxes_per_image_t = (gt_bboxes_per_image[:, 1]).unsqueeze(1).repeat(1, total_num_anchors) - center_radius * expanded_strides_per_image.unsqueeze(0)
|
|
gt_bboxes_per_image_b = (gt_bboxes_per_image[:, 1]).unsqueeze(1).repeat(1, total_num_anchors) + center_radius * expanded_strides_per_image.unsqueeze(0)
|
|
|
|
#-------------------------------------------------------#
|
|
# center_deltas [num_gt, n_anchors_all, 4]
|
|
#-------------------------------------------------------#
|
|
c_l = x_centers_per_image - gt_bboxes_per_image_l
|
|
c_r = gt_bboxes_per_image_r - x_centers_per_image
|
|
c_t = y_centers_per_image - gt_bboxes_per_image_t
|
|
c_b = gt_bboxes_per_image_b - y_centers_per_image
|
|
center_deltas = torch.stack([c_l, c_t, c_r, c_b], 2)
|
|
|
|
#-------------------------------------------------------#
|
|
# is_in_centers [num_gt, n_anchors_all]
|
|
# is_in_centers_all [n_anchors_all]
|
|
#-------------------------------------------------------#
|
|
is_in_centers = center_deltas.min(dim=-1).values > 0.0
|
|
is_in_centers_all = is_in_centers.sum(dim=0) > 0
|
|
|
|
#-------------------------------------------------------#
|
|
# is_in_boxes_anchor [n_anchors_all]
|
|
# is_in_boxes_and_center [num_gt, is_in_boxes_anchor]
|
|
#-------------------------------------------------------#
|
|
is_in_boxes_anchor = is_in_boxes_all | is_in_centers_all
|
|
is_in_boxes_and_center = is_in_boxes[:, is_in_boxes_anchor] & is_in_centers[:, is_in_boxes_anchor]
|
|
return is_in_boxes_anchor, is_in_boxes_and_center
|
|
|
|
def dynamic_k_matching(self, cost, pair_wise_ious, gt_classes, num_gt, fg_mask):
|
|
#-------------------------------------------------------#
|
|
# cost [num_gt, fg_mask]
|
|
# pair_wise_ious [num_gt, fg_mask]
|
|
# gt_classes [num_gt]
|
|
# fg_mask [n_anchors_all]
|
|
# matching_matrix [num_gt, fg_mask]
|
|
#-------------------------------------------------------#
|
|
matching_matrix = torch.zeros_like(cost)
|
|
|
|
#------------------------------------------------------------#
|
|
# 选取iou最大的n_candidate_k个点
|
|
# 然后求和，判断应该有多少点用于该框预测
|
|
# topk_ious [num_gt, n_candidate_k]
|
|
# dynamic_ks [num_gt]
|
|
# matching_matrix [num_gt, fg_mask]
|
|
#------------------------------------------------------------#
|
|
n_candidate_k = min(10, pair_wise_ious.size(1))
|
|
topk_ious, _ = torch.topk(pair_wise_ious, n_candidate_k, dim=1)
|
|
dynamic_ks = torch.clamp(topk_ious.sum(1).int(), min=1)
|
|
|
|
for gt_idx in range(num_gt):
|
|
#------------------------------------------------------------#
|
|
# 给每个真实框选取最小的动态k个点
|
|
#------------------------------------------------------------#
|
|
_, pos_idx = torch.topk(cost[gt_idx], k=dynamic_ks[gt_idx].item(), largest=False)
|
|
matching_matrix[gt_idx][pos_idx] = 1.0
|
|
del topk_ious, dynamic_ks, pos_idx
|
|
|
|
#------------------------------------------------------------#
|
|
# anchor_matching_gt [fg_mask]
|
|
#------------------------------------------------------------#
|
|
anchor_matching_gt = matching_matrix.sum(0)
|
|
if (anchor_matching_gt > 1).sum() > 0:
|
|
#------------------------------------------------------------#
|
|
# 当某一个特征点指向多个真实框的时候
|
|
# 选取cost最小的真实框。
|
|
#------------------------------------------------------------#
|
|
_, cost_argmin = torch.min(cost[:, anchor_matching_gt > 1], dim=0)
|
|
matching_matrix[:, anchor_matching_gt > 1] *= 0.0
|
|
matching_matrix[cost_argmin, anchor_matching_gt > 1] = 1.0
|
|
#------------------------------------------------------------#
|
|
# fg_mask_inboxes [fg_mask]
|
|
# num_fg为正样本的特征点个数
|
|
#------------------------------------------------------------#
|
|
fg_mask_inboxes = matching_matrix.sum(0) > 0.0
|
|
num_fg = fg_mask_inboxes.sum().item()
|
|
|
|
#------------------------------------------------------------#
|
|
# 对fg_mask进行更新
|
|
#------------------------------------------------------------#
|
|
fg_mask[fg_mask.clone()] = fg_mask_inboxes
|
|
|
|
#------------------------------------------------------------#
|
|
# 获得特征点对应的物品种类
|
|
#------------------------------------------------------------#
|
|
matched_gt_inds = matching_matrix[:, fg_mask_inboxes].argmax(0)
|
|
gt_matched_classes = gt_classes[matched_gt_inds]
|
|
|
|
pred_ious_this_matching = (matching_matrix * pair_wise_ious).sum(0)[fg_mask_inboxes]
|
|
return num_fg, gt_matched_classes, pred_ious_this_matching, matched_gt_inds
|
|
|
|
def is_parallel(model):
|
|
# Returns True if model is of type DP or DDP
|
|
return type(model) in (nn.parallel.DataParallel, nn.parallel.DistributedDataParallel)
|
|
|
|
def de_parallel(model):
|
|
# De-parallelize a model: returns single-GPU model if model is of type DP or DDP
|
|
return model.module if is_parallel(model) else model
|
|
|
|
def copy_attr(a, b, include=(), exclude=()):
|
|
# Copy attributes from b to a, options to only include [...] and to exclude [...]
|
|
for k, v in b.__dict__.items():
|
|
if (len(include) and k not in include) or k.startswith('_') or k in exclude:
|
|
continue
|
|
else:
|
|
setattr(a, k, v)
|
|
|
|
class ModelEMA:
|
|
""" Updated Exponential Moving Average (EMA) from https://github.com/rwightman/pytorch-image-models
|
|
Keeps a moving average of everything in the model state_dict (parameters and buffers)
|
|
For EMA details see https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverage
|
|
"""
|
|
|
|
def __init__(self, model, decay=0.9999, tau=2000, updates=0):
|
|
# Create EMA
|
|
self.ema = deepcopy(de_parallel(model)).eval() # FP32 EMA
|
|
# if next(model.parameters()).device.type != 'cpu':
|
|
# self.ema.half() # FP16 EMA
|
|
self.updates = updates # number of EMA updates
|
|
self.decay = lambda x: decay * (1 - math.exp(-x / tau)) # decay exponential ramp (to help early epochs)
|
|
for p in self.ema.parameters():
|
|
p.requires_grad_(False)
|
|
|
|
def update(self, model):
|
|
# Update EMA parameters
|
|
with torch.no_grad():
|
|
self.updates += 1
|
|
d = self.decay(self.updates)
|
|
|
|
msd = de_parallel(model).state_dict() # model state_dict
|
|
for k, v in self.ema.state_dict().items():
|
|
if v.dtype.is_floating_point:
|
|
v *= d
|
|
v += (1 - d) * msd[k].detach()
|
|
|
|
def update_attr(self, model, include=(), exclude=('process_group', 'reducer')):
|
|
# Update EMA attributes
|
|
copy_attr(self.ema, model, include, exclude)
|
|
|
|
def weights_init(net, init_type='normal', init_gain = 0.02):
|
|
def init_func(m):
|
|
classname = m.__class__.__name__
|
|
if hasattr(m, 'weight') and classname.find('Conv') != -1:
|
|
if init_type == 'normal':
|
|
torch.nn.init.normal_(m.weight.data, 0.0, init_gain)
|
|
elif init_type == 'xavier':
|
|
torch.nn.init.xavier_normal_(m.weight.data, gain=init_gain)
|
|
elif init_type == 'kaiming':
|
|
torch.nn.init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
|
|
elif init_type == 'orthogonal':
|
|
torch.nn.init.orthogonal_(m.weight.data, gain=init_gain)
|
|
else:
|
|
raise NotImplementedError('initialization method [%s] is not implemented' % init_type)
|
|
elif classname.find('BatchNorm2d') != -1:
|
|
torch.nn.init.normal_(m.weight.data, 1.0, 0.02)
|
|
torch.nn.init.constant_(m.bias.data, 0.0)
|
|
print('initialize network with %s type' % init_type)
|
|
net.apply(init_func)
|
|
|
|
def get_lr_scheduler(lr_decay_type, lr, min_lr, total_iters, warmup_iters_ratio = 0.05, warmup_lr_ratio = 0.1, no_aug_iter_ratio = 0.05, step_num = 10):
|
|
def yolox_warm_cos_lr(lr, min_lr, total_iters, warmup_total_iters, warmup_lr_start, no_aug_iter, iters):
|
|
if iters <= warmup_total_iters:
|
|
# lr = (lr - warmup_lr_start) * iters / float(warmup_total_iters) + warmup_lr_start
|
|
lr = (lr - warmup_lr_start) * pow(iters / float(warmup_total_iters), 2) + warmup_lr_start
|
|
elif iters >= total_iters - no_aug_iter:
|
|
lr = min_lr
|
|
else:
|
|
lr = min_lr + 0.5 * (lr - min_lr) * (
|
|
1.0 + math.cos(math.pi* (iters - warmup_total_iters) / (total_iters - warmup_total_iters - no_aug_iter))
|
|
)
|
|
return lr
|
|
|
|
def step_lr(lr, decay_rate, step_size, iters):
|
|
if step_size < 1:
|
|
raise ValueError("step_size must above 1.")
|
|
n = iters // step_size
|
|
out_lr = lr * decay_rate ** n
|
|
return out_lr
|
|
|
|
if lr_decay_type == "cos":
|
|
warmup_total_iters = min(max(warmup_iters_ratio * total_iters, 1), 3)
|
|
warmup_lr_start = max(warmup_lr_ratio * lr, 1e-6)
|
|
no_aug_iter = min(max(no_aug_iter_ratio * total_iters, 1), 15)
|
|
func = partial(yolox_warm_cos_lr ,lr, min_lr, total_iters, warmup_total_iters, warmup_lr_start, no_aug_iter)
|
|
else:
|
|
decay_rate = (min_lr / lr) ** (1 / (step_num - 1))
|
|
step_size = total_iters / step_num
|
|
func = partial(step_lr, lr, decay_rate, step_size)
|
|
|
|
return func
|
|
|
|
def set_optimizer_lr(optimizer, lr_scheduler_func, epoch):
|
|
lr = lr_scheduler_func(epoch)
|
|
for param_group in optimizer.param_groups:
|
|
param_group['lr'] = lr
|