-
Notifications
You must be signed in to change notification settings - Fork 3
Expand file tree
/
Copy pathfid.py
More file actions
153 lines (119 loc) · 5.39 KB
/
fid.py
File metadata and controls
153 lines (119 loc) · 5.39 KB
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
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
# -*- coding:utf-8 -*-
import torch
import torch.nn as nn
import fsspec
class TransformerWithToken(nn.Module):
def __init__(self, d_model, nhead, dim_feedforward, num_layers):
super().__init__()
self.token = nn.Parameter(torch.randn(1, 1, d_model))
token_mask = torch.zeros(1, 1, dtype=torch.bool)
self.register_buffer("token_mask", token_mask)
self.core = nn.TransformerEncoder(
nn.TransformerEncoderLayer(
d_model=d_model,
nhead=nhead,
dim_feedforward=dim_feedforward,
),
num_layers=num_layers,
)
def forward(self, x, src_key_padding_mask):
# x: [N, B, E]
# padding_mask: [B, N]
# `False` for valid values
# `True` for padded values
B = x.size(1)
token = self.token.expand(-1, B, -1)
x = torch.cat([token, x], dim=0)
token_mask = self.token_mask.expand(B, -1)
padding_mask = torch.cat([token_mask, src_key_padding_mask], dim=1)
x = self.core(x, src_key_padding_mask=padding_mask)
return x
class FIDNetV3(nn.Module):
def __init__(self, num_label, d_model=256, nhead=4, num_layers=4, max_bbox=50):
super().__init__()
self.emb_label = nn.Embedding(num_label, d_model)
self.fc_bbox = nn.Linear(4, d_model)
self.enc_fc_in = nn.Linear(d_model * 2, d_model)
self.enc_transformer = TransformerWithToken(
d_model=d_model,
dim_feedforward=d_model // 2,
nhead=nhead,
num_layers=num_layers,
)
self.fc_out_disc = nn.Linear(d_model, 1)
self.pos_token = nn.Parameter(torch.rand(max_bbox, 1, d_model))
self.dec_fc_in = nn.Linear(d_model * 2, d_model)
te = nn.TransformerEncoderLayer(d_model=d_model, nhead=nhead, dim_feedforward=d_model // 2)
self.dec_transformer = nn.TransformerEncoder(te, num_layers=num_layers)
self.fc_out_cls = nn.Linear(d_model, num_label)
self.fc_out_bbox = nn.Linear(d_model, 4)
def extract_features(self, bbox, label, padding_mask): # 可以解释，bbox和label都可以extract出来
b = self.fc_bbox(bbox)
l = self.emb_label(label)
x = self.enc_fc_in(torch.cat([b, l], dim=-1))
x = torch.relu(x).permute(1, 0, 2)
x = self.enc_transformer(x, padding_mask)
return x[0]
def forward(self, bbox, label, padding_mask):
B, N, _ = bbox.size()
x = self.extract_features(bbox, label, padding_mask)
logit_disc = self.fc_out_disc(x).squeeze(-1)
x = x.unsqueeze(0).expand(N, -1, -1)
t = self.pos_token[:N].expand(-1, B, -1)
x = torch.cat([x, t], dim=-1)
x = torch.relu(self.dec_fc_in(x))
x = self.dec_transformer(x, src_key_padding_mask=padding_mask)
# x = x.permute(1, 0, 2)[~padding_mask]
x = x.permute(1, 0, 2)
# logit_cls: [B, N, L] bbox_pred: [B, N, 4]
logit_cls = self.fc_out_cls(x)
bbox_pred = torch.sigmoid(self.fc_out_bbox(x))
return logit_disc, logit_cls, bbox_pred
class FIDNetV3Slide(nn.Module):
def __init__(self, num_label, d_model=256, nhead=4, num_layers=4, max_bbox=50):
super().__init__()
self.emb_label = nn.Embedding(num_label + 1, d_model)
self.fc_bbox = nn.Linear(4, d_model)
self.enc_fc_in = nn.Linear(d_model * 2, d_model)
self.enc_transformer = TransformerWithToken(
d_model=d_model,
dim_feedforward=d_model // 2,
nhead=nhead,
num_layers=num_layers,
)
self.fc_out_disc = nn.Linear(d_model, 1)
self.pos_token = nn.Parameter(torch.rand(max_bbox, 1, d_model))
self.dec_fc_in = nn.Linear(d_model * 2, d_model)
te = nn.TransformerEncoderLayer(d_model=d_model, nhead=nhead, dim_feedforward=d_model // 2)
self.dec_transformer = nn.TransformerEncoder(te, num_layers=num_layers)
self.fc_out_cls = nn.Linear(d_model, num_label)
self.fc_out_bbox = nn.Linear(d_model, 4)
def extract_features(self, bbox, label, padding_mask): # 可以解释，bbox和label都可以extract出来
b = self.fc_bbox(bbox)
l = self.emb_label(label)
x = self.enc_fc_in(torch.cat([b, l], dim=-1))
x = torch.relu(x).permute(1, 0, 2)
x = self.enc_transformer(x, padding_mask)
return x[0]
def forward(self, bbox, label, padding_mask):
B, N, _ = bbox.size()
x = self.extract_features(bbox, label, padding_mask)
logit_disc = self.fc_out_disc(x).squeeze(-1)
x = x.unsqueeze(0).expand(N, -1, -1)
t = self.pos_token[:N].expand(-1, B, -1)
x = torch.cat([x, t], dim=-1)
x = torch.relu(self.dec_fc_in(x))
x = self.dec_transformer(x, src_key_padding_mask=padding_mask)
# x = x.permute(1, 0, 2)[~padding_mask]
x = x.permute(1, 0, 2)
# logit_cls: [B, N, L] bbox_pred: [B, N, 4]
logit_cls = self.fc_out_cls(x)
bbox_pred = torch.sigmoid(self.fc_out_bbox(x))
return logit_disc, logit_cls, bbox_pred
def load_fidnet_v3(weight_path, device: torch.device) -> nn.Module:
fid_model = FIDNetV3(num_label=5, max_bbox=25).to(device)
with fsspec.open(weight_path, "rb") as file_obj:
x = torch.load(file_obj, map_location=device)
fid_model.load_state_dict(x["state_dict"])
fid_model.eval()
return fid_model