File size: 3,878 Bytes
404d2af |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 |
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
|