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

%%writefile -a lab10_detector.py
def report_predictions(pred,targets):
    empties=[i for i,t in enumerate(targets) if not len(t['boxes'])]
    multi=[i for i,t in enumerate(targets) if len(t['boxes'])>=2]
    return {'AP50':average_precision(pred,targets,.5),'AP75':average_precision(pred,targets,.75),
            'operating_point':average_precision(pred,targets,.5,.4),
            'empty_images':len(empties),'empty_false_positive_images_score04':sum(int((pred[i]['scores']>=.4).any()) for i in empties),
            'multi_object_images':len(multi),'multi_object_AP50':average_precision([pred[i] for i in multi],[targets[i] for i in multi],.5)}

def draw_predictions(images,targets,predictions,path,title,indices=None):
    if indices is None:indices=[0,1,2,3,4,5]
    fig,axes=plt.subplots(2,3,figsize=(11,7))
    for ax,i in zip(axes.flat,indices):
        ax.imshow(images[i].permute(1,2,0));ax.axis('off')
        for b,c in zip(targets[i]['boxes'],targets[i]['labels']):
            x,y,x2,y2=b.numpy()*48;ax.add_patch(Rectangle((x,y),x2-x,y2-y,fill=False,color='lime',lw=1.8))
        count=0
        for b,s,c in zip(predictions[i]['boxes'],predictions[i]['scores'],predictions[i]['labels']):
            if s<.4:continue
            x,y,x2,y2=b.numpy()*48;ax.add_patch(Rectangle((x,y),x2-x,y2-y,fill=False,color='yellow',ls='--',lw=1.1))
            ax.text(x,y,f'{int(c)}:{s:.2f}',color='yellow',fontsize=7);count+=1
        ax.set_title(f"sample {i}: GT={len(targets[i]['boxes'])}, pred={count}")
    fig.suptitle(title+' | green=GT; yellow=prediction score>=0.4');fig.tight_layout();fig.savefig(path,dpi=150);plt.close(fig)

def run(out='lab10-results',epochs=28):
    started=datetime.now(timezone.utc).isoformat();start=time.time();out=Path(out);out.mkdir(parents=True,exist_ok=True);seed_all(7)
    train,ttrain,_=make_scenes(768,1101); val,tval,_=make_scenes(96,3101);test,ttest,_=make_scenes(192,5901)
    small,tsmall,_=make_scenes(192,5901,scale=.5);anchors=anchors_grid()
    matched=[match_targets(anchors,t) for t in ttrain];classes=torch.stack([m[0] for m in matched]);targets=torch.stack([m[1] for m in matched])
    model=TinyAnchorDetector();random_pred=predict(model,test,anchors,score=.01)
    baseline=report_predictions(random_pred,ttest);opt=torch.optim.Adam(model.parameters(),lr=.002);history=[]
    for epoch in range(epochs):
        model.train(); losses=[]; order=torch.randperm(len(train))
        for ids in order.split(32):
            logits,offset=model(train[ids]);loss,cls,loc=detection_loss(logits,offset,classes[ids],targets[ids])
            opt.zero_grad();loss.backward();opt.step();losses.append([float(loss.detach()),float(cls.detach()),float(loc.detach())])
        item={'epoch':epoch+1,'loss':float(np.mean(losses,axis=0)[0]),'classification_loss':float(np.mean(losses,axis=0)[1]),'box_loss':float(np.mean(losses,axis=0)[2])};history.append(item)
        if (epoch+1)%7==0 or epoch==0:print(item,flush=True)
    pred=predict(model,test,anchors);small_pred=predict(model,small,anchors);val_pred=predict(model,val,anchors)
    loose_pred=predict(model,test,anchors,nms_iou=.65)
    results={'provenance':provenance(__file__,started),'config':{'seed':7,'train_count':768,'validation_count':96,'heldout_count':192,
             'split_seeds':[1101,3101,5901],'image_shape':[3,48,48],'num_anchors':len(anchors),'epochs':epochs,'batch_size':32,
             'optimizer':'Adam lr=.002','matching':'positive IoU>=.45 plus distinct forced GT anchors; ignore [.25,.45); background<.25',
             'loss':'hard-negative CE (3:1,min12) + 3*positive smoothL1','box_parameterization':'dx/dw,dy/dh,log width ratio,log height ratio',
             'AP':'per-class global score sort; maximum-IoU GT, one match per GT; precision-envelope all-points integral',
             'AP_candidate_score_floor':.01,'random_baseline_candidate_score_floor':.01,'NMS_IoU':.35,'display_score':.4,
             'model_parameters':sum(p.numel() for p in model.parameters()),'data_identity':'generated images; color/shape correlated; distinct original sample seed splits; no external data'},
             'target_counts':{'positive':int((classes>0).sum()),'ignore':int((classes==-1).sum()),'background':int((classes==0).sum())},
             'random_initialization':baseline,'validation':report_predictions(val_pred,tval),'heldout':report_predictions(pred,ttest),
             'condition_half_object_scale':report_predictions(small_pred,tsmall),'condition_NMS_065':report_predictions(loose_pred,ttest),
             'training_history':history,'seconds':time.time()-start}
    # Select actual worst examples by mean best GT overlap, plus an empty image.
    failures=[]
    for i,t in enumerate(ttest):
        p=pred[i];b=p['boxes'][p['scores']>=.4]
        quality=float(iou_matrix(t['boxes'],b).max(1).values.mean()) if len(t['boxes']) and len(b) else 0
        if len(t['boxes']):failures.append((quality,i))
    failures.sort();indices=[i for _,i in failures[:5]]+[next(i for i,t in enumerate(ttest) if not len(t['boxes']))]
    results['failure_examples']=[{'sample_index':i,'mean_best_box_IoU_unconditioned_on_class':q} for q,i in failures[:5]]
    torch.save({'state_dict':model.state_dict(),'anchors':anchors,'config':results['config']},out/'detector.pt')
    np.savez_compressed(out/'heldout-inputs-and-targets.npz',images=test.numpy(),half_scale_images=small.numpy(),
                        target_boxes=np.array([t['boxes'].numpy() for t in ttest],dtype=object),target_labels=np.array([t['labels'].numpy() for t in ttest],dtype=object))
    serialize=lambda ps:[{k:v.tolist() for k,v in p.items()} for p in ps]
    save_json(out/'predictions.json',{'heldout':serialize(pred),'half_scale':serialize(small_pred),'targets':serialize([{k:v for k,v in t.items() if k!='sample_id'} for t in ttest]),'half_scale_targets':serialize([{k:v for k,v in t.items() if k!='sample_id'} for t in tsmall])})
    draw_predictions(test,ttest,pred,out/'heldout-predictions.png','Heldout image-to-box predictions')
    draw_predictions(test,ttest,pred,out/'failure-predictions.png','Lowest geometric overlap examples',indices)
    draw_predictions(small,tsmall,small_pred,out/'half-scale-predictions.png','Same original scenes, object size x0.5')
    fig,ax=plt.subplots(figsize=(7,4));ax.plot([h['loss'] for h in history],label='total');ax.plot([h['classification_loss'] for h in history],label='classification');ax.plot([3*h['box_loss'] for h in history],label='3 x box');ax.set(xlabel='epoch (zero based)',ylabel='training loss');ax.legend();fig.tight_layout();fig.savefig(out/'training-curve.png',dpi=150);plt.close(fig)
    save_json(out/'run.json',results)
    print(json.dumps({'heldout_mAP50':results['heldout']['AP50']['mAP'],'half_scale_mAP50':results['condition_half_object_scale']['AP50']['mAP'],'seconds':results['seconds']},indent=2),flush=True)
    return results

if __name__=='__main__':
    p=argparse.ArgumentParser();p.add_argument('--out',default='lab10-results');p.add_argument('--epochs',type=int,default=28);args=p.parse_args();run(args.out,args.epochs)