ATSS implementation help

Peter_Featherstone: <p>I’m trying to implement ATSS (Adaptive training sample selection) in Pytorch using the smallest amount of code possible. Here is my implementation so far:</p>
<pre><code class="lang-auto">from einops import repeat, unpack

def assign_atss (
    anchors, # Tensor[N,4]
    targets, # Tensor[B,D,5]
    ps,      # List[3]
    nc,      # int
    topk     # int
):
    # Prep
    B, D    = targets.shape[:2]
    mask_gt = targets[...,-1] &gt; -1
    
    # Calculate all IOUs and distances between targets and anchors
    ious        = torchvision.ops.box_iou(targets[...,:4].reshape([-1, 4]), anchors).reshape([B,D,-1]) # [B,D,N]
    sxy         = (anchors[:,0:2] + anchors[:,2:4]) / 2
    tlt, trb    = targets[...,0:2], targets[...,2:4]
    txy         = (tlt + trb) / 2
    dists       = torch.cdist(txy.reshape([-1,2]), sxy).reshape([B,D,-1])

    # Select topk for each level
    indices = []
    offset  = 0
    for distl, num in zip(unpack(dists, ps, 'b a *'), ps):
        indices.append(distl.topk(topk, dim=-1, largest=False)[1] + offset)
        offset += num[0]
    indices = torch.cat(indices, -1)

    # IOU thresh mask
    ious_idx = ious.gather(2, indices)
    thresh   = ious_idx.mean(2, keepdim=True) + ious_idx.std(2, keepdim=True)
    mask_iou = ious &gt; thresh

    # Centre mask
    bbox_deltas = torch.cat((sxy[None,None] - tlt.unsqueeze(2), trb.unsqueeze(2) - sxy[None,None]), -1)
    mask_centre = bbox_deltas.amin(3) &gt; 1e-9

    # Best IOU for each anchor
    indices = (ious * mask_centre).argmax(1)
    indices = repeat(indices, 'b n -&gt; b n f', f=5)

    # Combine everything
    mask            = mask_gt.unsqueeze(-1) * mask_iou * mask_centre
    targets         = targets.gather(1, indices)
    targets_bbox    = targets[...,:4]
    targets_cls     = F.one_hot(targets[...,4].long(), num_classes=nc)
    target_obj      = mask.gather(1, indices[...,0:1])
    targets         = torch.cat((targets_bbox, target_obj.float(), targets_cls), -1)

    return targets
</code></pre>
<p>Where <code>anchors</code> has N anchor points (<code>= h*w*a</code>) where <code>a == 3</code> and is the number of priors per cell, and dim <code>4</code> contains values x0,y0,x1,y1.<br>
<code>targets</code> has a maximum of <code>D</code> ground truth labels per batch and is padded with <code>-1</code> values. The last dim of 5 contains values x0,y0,x1,y1,cls_index.<br>
<code>ps</code> is a list of 3 numbers containing the number of detections per feature pyramid level. For example if your levels stride by 8, 16 and 32 and <code>N == 10647</code> then <code>ps = [8112, 2028, 507]</code>.</p>
<p>Something somewhere is wrong since i get <code>target_obj</code> always being full of <code>False</code> values.<br>
Can someone help?<br>
I don’t really want to copy paste implementations from mmdetection, ultralytics or other repositories. The point here is to come up with a tiny implementation of ATSS in pure Pytorch while keeping it fairly readable.</p>
<p>Cheers</p>

Peter_Featherstone: <p>I think the following is correct. I missing <code>mask_topk</code>.</p>
<pre><code class="lang-auto">@torch.no_grad()
def assign_atss (
    anchors, # Tensor[N,4]
    targets, # Tensor[B,D,5]
    ps,      # List[3]
    nc,      # int
    topk     # int
):
    # Prep
    B, D, N, device = targets.shape[0], targets.shape[1], anchors.shape[0], anchors.device
    
    # Calculate all IOUs and distances between targets and anchors
    ious        = torchvision.ops.box_iou(targets[...,:4].reshape([-1, 4]), anchors).reshape([B,D,-1]) # [B,D,N]
    sxy         = (anchors[:,0:2] + anchors[:,2:4]) / 2
    tlt, trb    = targets[...,0:2], targets[...,2:4]
    txy         = (tlt + trb) / 2
    dists       = torch.cdist(txy.reshape([-1,2]), sxy).reshape([B,D,-1])

    # Select topk for each level
    indices = []
    offset  = 0
    for distl, num in zip(unpack(dists, ps, 'b a *'), ps):
        indices.append(distl.topk(topk, dim=-1, largest=False)[1] + offset)
        offset += num[0]
    indices = torch.cat(indices, -1)
    mask_topk = torch.zeros((B, D, N), dtype=torch.bool, device=device)
    mask_topk.scatter_(2, indices, torch.ones_like(indices, dtype=torch.bool))

    # IOU thresh mask
    ious_idx = ious.gather(2, indices)
    thresh   = ious_idx.mean(2, keepdim=True) + ious_idx.std(2, keepdim=True, correction=0)
    mask_iou = ious &gt; thresh

    # Centre mask
    bbox_deltas = torch.cat((sxy[None,None] - tlt.unsqueeze(2), trb.unsqueeze(2) - sxy[None,None]), -1)
    mask_centre = bbox_deltas.amin(3) &gt; 1e-9

    # Combine masks
    mask_gt = (targets[...,-1] &gt; -1).unsqueeze(-1)
    mask    = mask_gt * mask_centre * mask_iou * mask_topk

    # Best IOU for each anchor
    scores, indices = (ious * mask).max(1)
    mask = scores &gt; 1e-9
    
    # Combine
    targets_box     = targets.gather(1, repeat(indices, 'b n -&gt; b n f', f=4))
    targets_cls     = targets[...,4].long().gather(1, indices)
    targets_cls     = F.one_hot(targets_cls, num_classes=nc)
    targets_score   = mask.unsqueeze(-1).float()
    targets         = torch.cat([targets_box, targets_score, targets_cls], -1)
    targets[~mask] = 0

    return targets
</code></pre>

Peter_Featherstone: <p>This implementation is still quite slow.<br>
Possible reasons:</p>
<ul>
<li>We’re calculating IOUs on all anchors, not just the topk candidates.</li>
<li>The line <code>targets[~mask] = 0</code> takes a long time to complete. I’ve seen it take longer than the rest of the algorithm. It’s strictly not necessary since the mask (<code>targets[...,4] &gt; 0</code>) is correct.</li>
</ul>