File size: 1,679 Bytes
5633819 | 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 | # scripts/criterion_detect.py
from __future__ import annotations
from dataclasses import dataclass
from typing import Dict, List, Optional
import torch
from torch import nn
from scripts.matcher import HungarianMatcher
from scripts.set_criterion_countgd import SetCriterion
@dataclass
class CriterionDetectConfig:
num_classes: int = 1
class_cost: float = 1.0
bbox_cost: float = 5.0
giou_cost: float = 2.0
lambda_cls: float = 1.0
lambda_bbox: float = 5.0
lambda_giou: float = 2.0
eos_coef: float = 0.1
focal_alpha: float = 0.25
def build_criterion_detect(
tokenizer,
num_classes: int = 1,
class_cost: float = 1.0,
bbox_cost: float = 5.0,
giou_cost: float = 2.0,
lambda_cls: float = 1.0,
lambda_bbox: float = 5.0,
lambda_giou: float = 2.0,
eos_coef: float = 0.1,
focal_alpha: float = 0.25,
) -> SetCriterion:
"""
Criterion detect theo đúng CountGD:
- HungarianMatcher trên (cls/token cost + L1 bbox + GIoU)
- SetCriterion dùng create_positive_map_exemplar để map nhãn ↔ tokens
"""
matcher = HungarianMatcher(
cost_class=class_cost,
cost_bbox=bbox_cost,
cost_giou=giou_cost,
)
weight_dict = {
"loss_ce": float(lambda_cls),
"loss_bbox": float(lambda_bbox),
"loss_giou": float(lambda_giou),
}
losses = ["labels", "boxes", "cardinality"]
criterion = SetCriterion(
num_classes=num_classes,
matcher=matcher,
weight_dict=weight_dict,
eos_coef=eos_coef,
losses=losses,
focal_alpha=focal_alpha,
tokenizer=tokenizer,
)
return criterion
|