import math from functools import partial import torch import torch.nn as nn class MultiboxLoss(nn.Module): def __init__(self, num_classes, alpha=1.0, neg_pos_ratio=3.0, background_label_id=0, negatives_for_hard=100.0): self.num_classes = num_classes self.alpha = alpha self.neg_pos_ratio = neg_pos_ratio if background_label_id != 0: raise Exception('Only 0 as background label id is supported') self.background_label_id = background_label_id self.negatives_for_hard = torch.FloatTensor([negatives_for_hard])[0] def _l1_smooth_loss(self, y_true, y_pred): abs_loss = torch.abs(y_true - y_pred) sq_loss = 0.5 * (y_true - y_pred)**2 l1_loss = torch.where(abs_loss < 1.0, sq_loss, abs_loss - 0.5) return torch.sum(l1_loss, -1) def _softmax_loss(self, y_true, y_pred): y_pred = torch.clamp(y_pred, min = 1e-7) softmax_loss = -torch.sum(y_true * torch.log(y_pred), axis=-1) return softmax_loss def forward(self, y_true, y_pred): # --------------------------------------------- # # y_true batch_size, 8732, 4 + self.num_classes + 1 # y_pred batch_size, 8732, 4 + self.num_classes # --------------------------------------------- # num_boxes = y_true.size()[1] y_pred = torch.cat([y_pred[0], nn.Softmax(-1)(y_pred[1])], dim = -1) # --------------------------------------------- # # 分类的loss # batch_size,8732,21 -> batch_size,8732 # --------------------------------------------- # conf_loss = self._softmax_loss(y_true[:, :, 4:-1], y_pred[:, :, 4:]) # --------------------------------------------- # # 框的位置的loss # batch_size,8732,4 -> batch_size,8732 # --------------------------------------------- # loc_loss = self._l1_smooth_loss(y_true[:, :, :4], y_pred[:, :, :4]) # --------------------------------------------- # # 获取所有的正标签的loss # --------------------------------------------- # pos_loc_loss = torch.sum(loc_loss * y_true[:, :, -1], axis=1) pos_conf_loss = torch.sum(conf_loss * y_true[:, :, -1], axis=1) # --------------------------------------------- # # 每一张图的正样本的个数 # num_pos [batch_size,] # --------------------------------------------- # num_pos = torch.sum(y_true[:, :, -1], axis=-1) # --------------------------------------------- # # 每一张图的负样本的个数 # num_neg [batch_size,] # --------------------------------------------- # num_neg = torch.min(self.neg_pos_ratio * num_pos, num_boxes - num_pos) # 找到了哪些值是大于0的 pos_num_neg_mask = num_neg > 0 # --------------------------------------------- # # 如果所有的图，正样本的数量均为0 # 那么则默认选取100个先验框作为负样本 # --------------------------------------------- # has_min = torch.sum(pos_num_neg_mask) # --------------------------------------------- # # 从这里往后，与视频中看到的代码有些许不同。 # 由于以前的负样本选取方式存在一些问题， # 我对该部分代码进行重构。 # 求整个batch应该的负样本数量总和 # --------------------------------------------- # num_neg_batch = torch.sum(num_neg) if has_min > 0 else self.negatives_for_hard # --------------------------------------------- # # 对预测结果进行判断，如果该先验框没有包含物体 # 那么它的不属于背景的预测概率过大的话 # 就是难分类样本 # --------------------------------------------- # confs_start = 4 + self.background_label_id + 1 confs_end = confs_start + self.num_classes - 1 # --------------------------------------------- # # batch_size,8732 # 把不是背景的概率求和，求和后的概率越大 # 代表越难分类。 # --------------------------------------------- # max_confs = torch.sum(y_pred[:, :, confs_start:confs_end], dim=2) # --------------------------------------------------- # # 只有没有包含物体的先验框才得到保留 # 我们在整个batch里面选取最难分类的num_neg_batch个 # 先验框作为负样本。 # --------------------------------------------------- # max_confs = (max_confs * (1 - y_true[:, :, -1])).view([-1]) _, indices = torch.topk(max_confs, k = int(num_neg_batch.cpu().numpy().tolist())) neg_conf_loss = torch.gather(conf_loss.view([-1]), 0, indices) # 进行归一化 num_pos = torch.where(num_pos != 0, num_pos, torch.ones_like(num_pos)) total_loss = torch.sum(pos_conf_loss) + torch.sum(neg_conf_loss) + torch.sum(self.alpha * pos_loc_loss) total_loss = total_loss / torch.sum(num_pos) return total_loss 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