Spaces:
Runtime error
Runtime error
File size: 3,604 Bytes
3e99b05 |
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 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 |
import copy
import torch.nn as nn
from detectron2.modeling.backbone import ResNet, BasicStem
from detectron2.layers import ShapeSpec
from detectron2.config import LazyCall as L
from detrex.modeling.matcher import HungarianMatcher
from detrex.modeling.neck import ChannelMapper
from detrex.layers import PositionEmbeddingSine
from projects.dino.modeling import (
DINO,
DINOTransformerEncoder,
DINOTransformerDecoder,
DINOTransformer,
DINOCriterion,
)
model = L(DINO)(
backbone=L(ResNet)(
stem=L(BasicStem)(in_channels=3, out_channels=64, norm="FrozenBN"),
stages=L(ResNet.make_default_stages)(
depth=50,
stride_in_1x1=False,
norm="FrozenBN",
),
out_features=["res3", "res4", "res5"],
freeze_at=1,
),
position_embedding=L(PositionEmbeddingSine)(
num_pos_feats=128,
temperature=10000,
normalize=True,
offset=-0.5,
),
neck=L(ChannelMapper)(
input_shapes={
"res3": ShapeSpec(channels=512),
"res4": ShapeSpec(channels=1024),
"res5": ShapeSpec(channels=2048),
},
in_features=["res3", "res4", "res5"],
out_channels=256,
num_outs=4,
kernel_size=1,
norm_layer=L(nn.GroupNorm)(num_groups=32, num_channels=256),
),
transformer=L(DINOTransformer)(
encoder=L(DINOTransformerEncoder)(
embed_dim=256,
num_heads=8,
feedforward_dim=2048,
attn_dropout=0.0,
ffn_dropout=0.0,
num_layers=6,
post_norm=False,
num_feature_levels="${..num_feature_levels}",
use_checkpoint=False
),
decoder=L(DINOTransformerDecoder)(
embed_dim=256,
num_heads=8,
feedforward_dim=2048,
attn_dropout=0.0,
ffn_dropout=0.0,
num_layers=6,
return_intermediate=True,
num_feature_levels="${..num_feature_levels}",
use_checkpoint=False,
),
num_feature_levels=4,
two_stage_num_proposals="${..num_queries}",
),
embed_dim=256,
num_classes=80,
num_queries=900,
aux_loss=True,
criterion=L(DINOCriterion)(
num_classes="${..num_classes}",
matcher=L(HungarianMatcher)(
cost_class=2.0,
cost_bbox=5.0,
cost_giou=2.0,
cost_class_type="focal_loss_cost",
alpha=0.25,
gamma=2.0,
),
weight_dict={
"loss_class": 1,
"loss_bbox": 5.0,
"loss_giou": 2.0,
"loss_class_dn": 1,
"loss_bbox_dn": 5.0,
"loss_giou_dn": 2.0,
},
loss_class_type="focal_loss",
alpha=0.25,
gamma=2.0,
two_stage_binary_cls=False,
),
dn_number=100,
label_noise_ratio=0.5,
box_noise_scale=1.0,
pixel_mean=[123.675, 116.280, 103.530],
pixel_std=[58.395, 57.120, 57.375],
vis_period=0,
input_format="RGB",
device="cuda",
)
# set aux loss weight dict
base_weight_dict = copy.deepcopy(model.criterion.weight_dict)
if model.aux_loss:
weight_dict = model.criterion.weight_dict
aux_weight_dict = {}
aux_weight_dict.update({k + "_enc": v for k, v in base_weight_dict.items()})
for i in range(model.transformer.decoder.num_layers - 1):
aux_weight_dict.update({k + f"_{i}": v for k, v in base_weight_dict.items()})
weight_dict.update(aux_weight_dict)
model.criterion.weight_dict = weight_dict
|