从共享特征到两个输出头 · 实现代码

%%writefile -a lab10_detector.py
class TinyAnchorDetector(nn.Module):
    def __init__(self):
        super().__init__()
        self.backbone=nn.Sequential(nn.Conv2d(3,16,3,2,1),nn.ReLU(),nn.Conv2d(16,32,3,2,1),nn.ReLU(),
                                    nn.Conv2d(32,48,3,1,1),nn.ReLU(),nn.Conv2d(48,48,3,1,1),nn.ReLU())
        self.class_head=nn.Conv2d(48,2*3,3,padding=1);self.box_head=nn.Conv2d(48,2*4,3,padding=1)
    def forward(self,x):
        z=self.backbone(x);b=x.shape[0]
        return self.class_head(z).permute(0,2,3,1).reshape(b,-1,3),self.box_head(z).permute(0,2,3,1).reshape(b,-1,4)

def detection_loss(logits,offsets,classes,targets):
    ce=F.cross_entropy(logits.transpose(1,2),classes,ignore_index=-1,reduction='none')
    positive=classes>0;chosen=positive.clone()
    for i in range(len(classes)):
        negatives=torch.where(classes[i]==0)[0]; k=min(len(negatives),max(12,3*int(positive[i].sum())))
        chosen[i,negatives[torch.argsort(ce[i,negatives],descending=True)[:k]]]=True
    cls=ce[chosen].mean()
    loc=F.smooth_l1_loss(offsets[positive],targets[positive],reduction='mean') if positive.any() else offsets.sum()*0
    return cls+3*loc,cls,loc

@torch.no_grad()
def predict(model,images,anchors,score=.01,nms_iou=.35):
    model.eval();result=[]
    for batch in images.split(32):
        logits,offsets=model(batch); probs=logits.softmax(-1)
        for p,o in zip(probs,offsets):
            boxes=decode(anchors,o); all_b=[];all_s=[];all_l=[]
            for cls in [1,2]:
                indices=torch.where(p[:,cls]>=score)[0]
                keep=nms(boxes[indices],p[indices,cls],nms_iou);indices=indices[keep]
                all_b.append(boxes[indices]);all_s.append(p[indices,cls]);all_l.append(torch.full((len(indices),),cls))
            result.append({'boxes':torch.cat(all_b),'scores':torch.cat(all_s),'labels':torch.cat(all_l)})
    return result