pengsida/NeuralBody

1import torch2import os3from torch import nn4import numpy as np5import torch.nn.functional6from collections import OrderedDict7from termcolor import colored8 9 10def sigmoid(x):11    y = torch.clamp(x.sigmoid(), min=1e-4, max=1 - 1e-4)12    return y13 14 15def _neg_loss(pred, gt):16    ''' Modified focal loss. Exactly the same as CornerNet.17        Runs faster and costs a little bit more memory18        Arguments:19            pred (batch x c x h x w)20            gt_regr (batch x c x h x w)21    '''22    pos_inds = gt.eq(1).float()23    neg_inds = gt.lt(1).float()24 25    neg_weights = torch.pow(1 - gt, 4)26 27    loss = 028 29    pos_loss = torch.log(pred) * torch.pow(1 - pred, 2) * pos_inds30    neg_loss = torch.log(1 - pred) * torch.pow(pred,31                                               2) * neg_weights * neg_inds32 33    num_pos = pos_inds.float().sum()34    pos_loss = pos_loss.sum()35    neg_loss = neg_loss.sum()36 37    if num_pos == 0:38        loss = loss - neg_loss39    else:40        loss = loss - (pos_loss + neg_loss) / num_pos41    return loss42 43 44class FocalLoss(nn.Module):45    '''nn.Module warpper for focal loss'''46    def __init__(self):47        super(FocalLoss, self).__init__()48        self.neg_loss = _neg_loss49 50    def forward(self, out, target):51        return self.neg_loss(out, target)52 53 54def smooth_l1_loss(vertex_pred,55                   vertex_targets,56                   vertex_weights,57                   sigma=1.0,58                   normalize=True,59                   reduce=True):60    """61    :param vertex_pred:     [b, vn*2, h, w]62    :param vertex_targets:  [b, vn*2, h, w]63    :param vertex_weights:  [b, 1, h, w]64    :param sigma:65    :param normalize:66    :param reduce:67    :return:68    """69    b, ver_dim, _, _ = vertex_pred.shape70    sigma_2 = sigma**271    vertex_diff = vertex_pred - vertex_targets72    diff = vertex_weights * vertex_diff73    abs_diff = torch.abs(diff)74    smoothL1_sign = (abs_diff < 1. / sigma_2).detach().float()75    in_loss = torch.pow(diff, 2) * (sigma_2 / 2.) * smoothL1_sign \76              + (abs_diff - (0.5 / sigma_2)) * (1. - smoothL1_sign)77 78    if normalize:79        in_loss = torch.sum(in_loss.view(b, -1), 1) / (80            ver_dim * torch.sum(vertex_weights.view(b, -1), 1) + 1e-3)81 82    if reduce:83        in_loss = torch.mean(in_loss)84 85    return in_loss86 87 88class SmoothL1Loss(nn.Module):89    def __init__(self):90        super(SmoothL1Loss, self).__init__()91        self.smooth_l1_loss = smooth_l1_loss92 93    def forward(self,94                preds,95                targets,96                weights,97                sigma=1.0,98                normalize=True,99                reduce=True):100        return self.smooth_l1_loss(preds, targets, weights, sigma, normalize,101                                   reduce)102 103 104class AELoss(nn.Module):105    def __init__(self):106        super(AELoss, self).__init__()107 108    def forward(self, ae, ind, ind_mask):109        """110        ae: [b, 1, h, w]111        ind: [b, max_objs, max_parts]112        ind_mask: [b, max_objs, max_parts]113        obj_mask: [b, max_objs]114        """115        # first index116        b, _, h, w = ae.shape117        b, max_objs, max_parts = ind.shape118        obj_mask = torch.sum(ind_mask, dim=2) != 0119 120        ae = ae.view(b, h * w, 1)121        seed_ind = ind.view(b, max_objs * max_parts, 1)122        tag = ae.gather(1, seed_ind).view(b, max_objs, max_parts)123 124        # compute the mean125        tag_mean = tag * ind_mask126        tag_mean = tag_mean.sum(2) / (ind_mask.sum(2) + 1e-4)127 128        # pull ae of the same object to their mean129        pull_dist = (tag - tag_mean.unsqueeze(2)).pow(2) * ind_mask130        obj_num = obj_mask.sum(dim=1).float()131        pull = (pull_dist.sum(dim=(1, 2)) / (obj_num + 1e-4)).sum()132        pull /= b133 134        # push away the mean of different objects135        push_dist = torch.abs(tag_mean.unsqueeze(1) - tag_mean.unsqueeze(2))136        push_dist = 1 - push_dist137        push_dist = nn.functional.relu(push_dist, inplace=True)138        obj_mask = (obj_mask.unsqueeze(1) + obj_mask.unsqueeze(2)) == 2139        push_dist = push_dist * obj_mask.float()140        push = ((push_dist.sum(dim=(1, 2)) - obj_num) /141                (obj_num * (obj_num - 1) + 1e-4)).sum()142        push /= b143        return pull, push144 145 146class PolyMatchingLoss(nn.Module):147    def __init__(self, pnum):148        super(PolyMatchingLoss, self).__init__()149 150        self.pnum = pnum151        batch_size = 1152        pidxall = np.zeros(shape=(batch_size, pnum, pnum), dtype=np.int32)153        for b in range(batch_size):154            for i in range(pnum):155                pidx = (np.arange(pnum) + i) % pnum156                pidxall[b, i] = pidx157 158        device = torch.device('cuda')159        pidxall = torch.from_numpy(160            np.reshape(pidxall, newshape=(batch_size, -1))).to(device)161 162        self.feature_id = pidxall.unsqueeze_(2).long().expand(163            pidxall.size(0), pidxall.size(1), 2).detach()164 165    def forward(self, pred, gt, loss_type="L2"):166        pnum = self.pnum167        batch_size = pred.size()[0]168        feature_id = self.feature_id.expand(batch_size,169                                            self.feature_id.size(1), 2)170        device = torch.device('cuda')171 172        gt_expand = torch.gather(gt, 1,173                                 feature_id).view(batch_size, pnum, pnum, 2)174 175        pred_expand = pred.unsqueeze(1)176 177        dis = pred_expand - gt_expand178 179        if loss_type == "L2":180            dis = (dis**2).sum(3).sqrt().sum(2)181        elif loss_type == "L1":182            dis = torch.abs(dis).sum(3).sum(2)183 184        min_dis, min_id = torch.min(dis, dim=1, keepdim=True)185        # print(min_id)186 187        # min_id = torch.from_numpy(min_id.data.cpu().numpy()).to(device)188        # min_gt_id_to_gather = min_id.unsqueeze_(2).unsqueeze_(3).long().\189        #                         expand(min_id.size(0), min_id.size(1), gt_expand.size(2), gt_expand.size(3))190        # gt_right_order = torch.gather(gt_expand, 1, min_gt_id_to_gather).view(batch_size, pnum, 2)191 192        return torch.mean(min_dis)193 194 195class AttentionLoss(nn.Module):196    def __init__(self, beta=4, gamma=0.5):197        super(AttentionLoss, self).__init__()198 199        self.beta = beta200        self.gamma = gamma201 202    def forward(self, pred, gt):203        num_pos = torch.sum(gt)204        num_neg = torch.sum(1 - gt)205        alpha = num_neg / (num_pos + num_neg)206        edge_beta = torch.pow(self.beta, torch.pow(1 - pred, self.gamma))207        bg_beta = torch.pow(self.beta, torch.pow(pred, self.gamma))208 209        loss = 0210        loss = loss - alpha * edge_beta * torch.log(pred) * gt211        loss = loss - (1 - alpha) * bg_beta * torch.log(1 - pred) * (1 - gt)212        return torch.mean(loss)213 214 215def _gather_feat(feat, ind, mask=None):216    dim = feat.size(2)217    ind = ind.unsqueeze(2).expand(ind.size(0), ind.size(1), dim)218    feat = feat.gather(1, ind)219    if mask is not None:220        mask = mask.unsqueeze(2).expand_as(feat)221        feat = feat[mask]222        feat = feat.view(-1, dim)223    return feat224 225 226def _tranpose_and_gather_feat(feat, ind):227    feat = feat.permute(0, 2, 3, 1).contiguous()228    feat = feat.view(feat.size(0), -1, feat.size(3))229    feat = _gather_feat(feat, ind)230    return feat231 232 233class Ind2dRegL1Loss(nn.Module):234    def __init__(self, type='l1'):235        super(Ind2dRegL1Loss, self).__init__()236        if type == 'l1':237            self.loss = torch.nn.functional.l1_loss238        elif type == 'smooth_l1':239            self.loss = torch.nn.functional.smooth_l1_loss240 241    def forward(self, output, target, ind, ind_mask):242        """ind: [b, max_objs, max_parts]"""243        b, max_objs, max_parts = ind.shape244        ind = ind.view(b, max_objs * max_parts)245        pred = _tranpose_and_gather_feat(output,246                                         ind).view(b, max_objs, max_parts,247                                                   output.size(1))248        mask = ind_mask.unsqueeze(3).expand_as(pred)249        loss = self.loss(pred * mask, target * mask, reduction='sum')250        loss = loss / (mask.sum() + 1e-4)251        return loss252 253 254class IndL1Loss1d(nn.Module):255    def __init__(self, type='l1'):256        super(IndL1Loss1d, self).__init__()257        if type == 'l1':258            self.loss = torch.nn.functional.l1_loss259        elif type == 'smooth_l1':260            self.loss = torch.nn.functional.smooth_l1_loss261 262    def forward(self, output, target, ind, weight):263        """ind: [b, n]"""264        output = _tranpose_and_gather_feat(output, ind)265        weight = weight.unsqueeze(2)266        loss = self.loss(output * weight, target * weight, reduction='sum')267        loss = loss / (weight.sum() * output.size(2) + 1e-4)268        return loss269 270 271class GeoCrossEntropyLoss(nn.Module):272    def __init__(self):273        super(GeoCrossEntropyLoss, self).__init__()274 275    def forward(self, output, target, poly):276        output = torch.nn.functional.softmax(output, dim=1)277        output = torch.log(torch.clamp(output, min=1e-4))278        poly = poly.view(poly.size(0), 4, poly.size(1) // 4, 2)279        target = target[..., None, None].expand(poly.size(0), poly.size(1), 1,280                                                poly.size(3))281        target_poly = torch.gather(poly, 2, target)282        sigma = (poly[:, :, 0] - poly[:, :, 1]).pow(2).sum(2, keepdim=True)283        kernel = torch.exp(-(poly - target_poly).pow(2).sum(3) / (sigma / 3))284        loss = -(output * kernel.transpose(2, 1)).sum(1).mean()285        return loss286 287 288def load_model(net,289               optim,290               scheduler,291               recorder,292               model_dir,293               resume=True,294               epoch=-1):295    if not resume:296        os.system('rm -rf {}'.format(model_dir))297 298    if not os.path.exists(model_dir):299        return 0300 301    pths = [302        int(pth.split('.')[0]) for pth in os.listdir(model_dir)303        if pth != 'latest.pth'304    ]305    if len(pths) == 0 and 'latest.pth' not in os.listdir(model_dir):306        return 0307    if epoch == -1:308        if 'latest.pth' in os.listdir(model_dir):309            pth = 'latest'310        else:311            pth = max(pths)312    else:313        pth = epoch314    print('load model: {}'.format(os.path.join(model_dir,315                                               '{}.pth'.format(pth))))316    pretrained_model = torch.load(317        os.path.join(model_dir, '{}.pth'.format(pth)), 'cpu')318    net.load_state_dict(pretrained_model['net'])319    optim.load_state_dict(pretrained_model['optim'])320    scheduler.load_state_dict(pretrained_model['scheduler'])321    recorder.load_state_dict(pretrained_model['recorder'])322    return pretrained_model['epoch'] + 1323 324 325def save_model(net, optim, scheduler, recorder, model_dir, epoch, last=False):326    os.system('mkdir -p {}'.format(model_dir))327    model = {328        'net': net.state_dict(),329        'optim': optim.state_dict(),330        'scheduler': scheduler.state_dict(),331        'recorder': recorder.state_dict(),332        'epoch': epoch333    }334    if last:335        torch.save(model, os.path.join(model_dir, 'latest.pth'))336    else:337        torch.save(model, os.path.join(model_dir, '{}.pth'.format(epoch)))338 339    # remove previous pretrained model if the number of models is too big340    pths = [341        int(pth.split('.')[0]) for pth in os.listdir(model_dir)342        if pth != 'latest.pth'343    ]344    if len(pths) <= 20:345        return346    os.system('rm {}'.format(347        os.path.join(model_dir, '{}.pth'.format(min(pths)))))348 349 350def load_network(net, model_dir, resume=True, epoch=-1, strict=True):351    if not resume:352        return 0353 354    if not os.path.exists(model_dir):355        print(colored('pretrained model does not exist', 'red'))356        return 0357 358    if os.path.isdir(model_dir):359        pths = [360            int(pth.split('.')[0]) for pth in os.listdir(model_dir)361            if pth != 'latest.pth'362        ]363        if len(pths) == 0 and 'latest.pth' not in os.listdir(model_dir):364            return 0365        if epoch == -1:366            if 'latest.pth' in os.listdir(model_dir):367                pth = 'latest'368            else:369                pth = max(pths)370        else:371            pth = epoch372        model_path = os.path.join(model_dir, '{}.pth'.format(pth))373    else:374        model_path = model_dir375 376    print('load model: {}'.format(model_path))377    pretrained_model = torch.load(model_path)378    net.load_state_dict(pretrained_model['net'], strict=strict)379    return pretrained_model['epoch'] + 1380 381 382def remove_net_prefix(net, prefix):383    net_ = OrderedDict()384    for k in net.keys():385        if k.startswith(prefix):386            net_[k[len(prefix):]] = net[k]387        else:388            net_[k] = net[k]389    return net_390 391 392def add_net_prefix(net, prefix):393    net_ = OrderedDict()394    for k in net.keys():395        net_[prefix + k] = net[k]396    return net_397 398 399def replace_net_prefix(net, orig_prefix, prefix):400    net_ = OrderedDict()401    for k in net.keys():402        if k.startswith(orig_prefix):403            net_[prefix + k[len(orig_prefix):]] = net[k]404        else:405            net_[k] = net[k]406    return net_407 408 409def remove_net_layer(net, layers):410    keys = list(net.keys())411    for k in keys:412        for layer in layers:413            if k.startswith(layer):414                del net[k]415    return net416