Spaces:
Running
Running
| import torch | |
| import torch.nn as nn | |
| import numpy as np | |
| class APLoss (nn.Module): | |
| """ differentiable AP loss, through quantization. | |
| Input: (N, M) values in [min, max] | |
| label: (N, M) values in {0, 1} | |
| Returns: list of query AP (for each n in {1..N}) | |
| Note: typically, you want to minimize 1 - mean(AP) | |
| """ | |
| def __init__(self, nq=25, min=0, max=1, euc=False): | |
| nn.Module.__init__(self) | |
| assert isinstance(nq, int) and 2 <= nq <= 100 | |
| self.nq = nq | |
| self.min = min | |
| self.max = max | |
| self.euc = euc | |
| gap = max - min | |
| assert gap > 0 | |
| # init quantizer = non-learnable (fixed) convolution | |
| self.quantizer = q = nn.Conv1d(1, 2*nq, kernel_size=1, bias=True) | |
| a = (nq-1) / gap | |
| #1st half = lines passing to (min+x,1) and (min+x+1/a,0) with x = {nq-1..0}*gap/(nq-1) | |
| q.weight.data[:nq] = -a | |
| q.bias.data[:nq] = torch.from_numpy(a*min + np.arange(nq, 0, -1)) # b = 1 + a*(min+x) | |
| #2nd half = lines passing to (min+x,1) and (min+x-1/a,0) with x = {nq-1..0}*gap/(nq-1) | |
| q.weight.data[nq:] = a | |
| q.bias.data[nq:] = torch.from_numpy(np.arange(2-nq, 2, 1) - a*min) # b = 1 - a*(min+x) | |
| # first and last one are special: just horizontal straight line | |
| q.weight.data[0] = q.weight.data[-1] = 0 | |
| q.bias.data[0] = q.bias.data[-1] = 1 | |
| def compute_AP(self, x, label): | |
| N, M = x.shape | |
| # print(x.shape, label.shape) | |
| if self.euc: # euclidean distance in same range than similarities | |
| x = 1 - torch.sqrt(2.001 - 2*x) | |
| # quantize all predictions | |
| q = self.quantizer(x.unsqueeze(1)) | |
| q = torch.min(q[:,:self.nq], q[:,self.nq:]).clamp(min=0) # N x Q x M [1600, 20, 1681] | |
| nbs = q.sum(dim=-1) # number of samples N x Q = c | |
| rec = (q * label.view(N,1,M).float()).sum(dim=-1) # nb of correct samples = c+ N x Q | |
| prec = rec.cumsum(dim=-1) / (1e-16 + nbs.cumsum(dim=-1)) # precision | |
| rec /= rec.sum(dim=-1).unsqueeze(1) # norm in [0,1] | |
| ap = (prec * rec).sum(dim=-1) # per-image AP | |
| return ap | |
| def forward(self, x, label): | |
| assert x.shape == label.shape # N x M | |
| return self.compute_AP(x, label) | |
| class PixelAPLoss (nn.Module): | |
| """ Computes the pixel-wise AP loss: | |
| Given two images and ground-truth optical flow, computes the AP per pixel. | |
| feat1: (B, C, H, W) pixel-wise features extracted from img1 | |
| feat2: (B, C, H, W) pixel-wise features extracted from img2 | |
| aflow: (B, 2, H, W) absolute flow: aflow[...,y1,x1] = x2,y2 | |
| """ | |
| def __init__(self, sampler, nq=20): | |
| nn.Module.__init__(self) | |
| self.aploss = APLoss(nq, min=0, max=1, euc=False) | |
| self.name = 'pixAP' | |
| self.sampler = sampler | |
| def loss_from_ap(self, ap, rel): | |
| return 1 - ap | |
| def forward(self, feat0, feat1, conf0, conf1, pos0, pos1, B, H, W, N=1200): | |
| # subsample things | |
| scores, gt, msk, qconf = self.sampler(feat0, feat1, conf0, conf1, pos0, pos1, B, H, W, N=1200) | |
| # compute pixel-wise AP | |
| n = qconf.numel() | |
| if n == 0: return 0 | |
| scores, gt = scores.view(n,-1), gt.view(n,-1) | |
| ap = self.aploss(scores, gt).view(msk.shape) | |
| pixel_loss = self.loss_from_ap(ap, qconf) | |
| loss = pixel_loss[msk].mean() | |
| return loss | |
| class ReliabilityLoss (PixelAPLoss): | |
| """ same than PixelAPLoss, but also train a pixel-wise confidence | |
| that this pixel is going to have a good AP. | |
| """ | |
| def __init__(self, sampler, base=0.5, **kw): | |
| PixelAPLoss.__init__(self, sampler, **kw) | |
| assert 0 <= base < 1 | |
| self.base = base | |
| def loss_from_ap(self, ap, rel): | |
| return 1 - ap*rel - (1-rel)*self.base | |