-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathLongMIL.py
323 lines (263 loc) · 12.2 KB
/
LongMIL.py
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
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
from __future__ import annotations
import copy
import pdb
import numpy as np
from models.rotary import apply_rotary_position_embeddings, Rotary2D
from torch import autograd
try:
from xformers.ops import memory_efficient_attention
except:
print('please install xformer')
try:
from fla.layers import GatedLinearAttention
except:
print('please install fla from https://github.com/sustcsonglin/flash-linear-attention')
try:
from mamba_ssm import Mamba
except:
print('please install mamba_ssm')
import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange
def exists(val):
return val is not None
def precompute_freqs_cis(dim: int, end: int, pos_idx, theta: float = 10000.0):
freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
t = torch.arange(end, device=freqs.device) # type: ignore
# t = pos_idx.cpu()
freqs = torch.outer(t, freqs).float() # type: ignore
freqs_cis = torch.polar(torch.ones_like(freqs), freqs) # complex64
# pdb.set_trace()
return freqs_cis[pos_idx.long()]
def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):
ndim = x.ndim
assert 0 <= 1 < ndim
assert freqs_cis.shape == (x.shape[1], x.shape[-1])
shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
return freqs_cis.view(*shape)
def apply_rotary_emb(xq: torch.Tensor, xk: torch.Tensor,freqs_cis: torch.Tensor,):
xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
freqs_cis = reshape_for_broadcast(freqs_cis, xq_)
xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
return xq_out.type_as(xq), xk_out.type_as(xk)
class Attention(nn.Module):
def __init__(self, dim, num_heads=4, qkv_bias=False, qk_scale=None, attn_drop=0.05, proj_drop=0.):
super().__init__()
self.num_heads = num_heads
head_dim = dim // num_heads
self.scale = qk_scale or head_dim ** -0.5
self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
self.attn_drop = attn_drop
# self.attn_drop = nn.Dropout(attn_drop)
self.proj = nn.Linear(dim, dim)
self.proj_drop = nn.Dropout(proj_drop)
def forward(self, x, freqs_cis=None, alibi=None):
B, N, C = x.shape
qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 1, 3, 4) #
q, k, v = qkv[0], qkv[1], qkv[2]
q,k = q.view(1,N,1,-1), k.view(1,N,1,-1)
if exists(freqs_cis):
q,k = apply_rotary_position_embeddings(freqs_cis, q, k)
q, k = q.view(1, N, self.num_heads, -1), k.view(1, N, self.num_heads, -1)
if exists(alibi):
try:
x = memory_efficient_attention(q, k, v, alibi,p=self.attn_drop).reshape(B, N, C)
except:
print('xformer error')
else:
x = memory_efficient_attention(q,k,v,p=self.attn_drop).reshape(B, N, C)
x = self.proj(x)
x = self.proj_drop(x)
return x, None
class LinearAttention(nn.Module):
def __init__(self, dim, num_heads=4, qkv_bias=False, qk_scale=None, attn_drop=0.05, proj_drop=0.):
super().__init__()
self.num_heads = num_heads
self.gla = GatedLinearAttention(d_model=dim, num_heads=num_heads,causal=False)#.to(device).to(dtype)
def forward(self, x, freqs_cis=None, alibi=None):
# pdb.set_trace()
return self.gla(x) #, None
class TransLayer(nn.Module):
def __init__(self, norm_layer=nn.LayerNorm, dim=512,num_heads=4,attn_drop=0.05,mlp_ratio=4.):
super().__init__()
self.norm = norm_layer(dim)
self.attn = Attention(dim=dim, num_heads=num_heads,attn_drop=attn_drop)
# self.attn = LinearAttention(dim=dim,num_heads=num_heads)
def forward(self, x, rope=None, alibi=None):
temp = self.attn(self.norm(x), rope, alibi)#[0], None
return x+temp[0], temp[1]
class Mlp(nn.Module):
def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
super().__init__()
out_features = out_features or in_features
hidden_features = hidden_features or in_features
self.fc1 = nn.Linear(in_features, hidden_features)
self.act = act_layer()
self.fc2 = nn.Linear(hidden_features, out_features)
self.drop = nn.Dropout(drop)
def forward(self, x):
x = self.fc1(x)
x = self.act(x)
x = self.drop(x)
x = self.fc2(x)
x = self.drop(x)
return x
def drop_path(x, drop_prob: float = 0., training: bool = False):
if drop_prob == 0. or not training:
return x
keep_prob = 1 - drop_prob
shape = (x.shape[0],) + (1,) * (x.ndim - 1) # work with diff dim tensors, not just 2D ConvNets
random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
random_tensor.floor_() # binarize
output = x.div(keep_prob) * random_tensor
return output
class DropPath(nn.Module):
"""Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
"""
def __init__(self, drop_prob=None):
super(DropPath, self).__init__()
self.drop_prob = drop_prob
def forward(self, x):
return drop_path(x, self.drop_prob, self.training)
class TransBlock(nn.Module):
def __init__(self, norm_layer=nn.LayerNorm, dim=512,num_heads=4,attn_drop=0.05,mlp_ratio=4.,drop_path=0.):
super().__init__()
self.norm = norm_layer(dim)
self.attn = Attention(dim=dim, num_heads=num_heads,attn_drop=attn_drop)
# self.attn = LinearAttention(dim=dim,num_heads=num_heads)
self.norm2 = norm_layer(dim)
mlp_hidden_dim = int(dim * mlp_ratio)
self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=nn.GELU)
self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
def forward(self, x, rope=None, alibi=None):
temp = self.attn(self.norm(x), rope, alibi) # attn for spatial mixture
# return x+temp[0], temp[1]
x = x + self.drop_path(temp[0])
x = x + self.drop_path(self.mlp(self.norm2(x))) # mlp for token mixture
return x, temp[1]
class LinearTransLayer(nn.Module):
def __init__(self, norm_layer=nn.LayerNorm, dim=512,num_heads=4):
super().__init__()
self.norm = norm_layer(dim)
# self.attn = Attention(dim=dim,num_heads=num_heads)
self.attn = LinearAttention(dim=dim,num_heads=num_heads)
def forward(self, x, rope, alibi):
return x + self.attn(self.norm(x)), None
class MambaLayer(nn.Module):
def __init__(self, dim, d_state=16, d_conv=4, expand=2):
super().__init__()
self.dim = dim
self.norm = nn.LayerNorm(dim)
self.mamba = Mamba(
d_model=dim, # Model dimension d_model
d_state=d_state, # SSM state expansion factor
d_conv=d_conv, # Local convolution width
expand=expand, # Block expansion factor
)
# @autocast(enabled=False)
def forward(self, x, scan_idx=None):
x = self.norm(x)
x1 = self.mamba(x)
return x1
# option 2 : multi-scan orientation, need scan idx.
x2 = self.mamba(x[:, scan_idx])[:,np.argsort(scan_idx)]
x3 = self.mamba(x.flip(dims=[1])).flip(dims=[1])
x4 = self.mamba(x[:, scan_idx].flip(dims=[1])).flip(dims=[1])[:,np.argsort(scan_idx)]
return (x1+x2+x3+x4)/4.
import math
def get_slopes(n_heads: int):
n = 2 ** math.floor(math.log2(n_heads))
m_0 = 2.0 ** (-8.0 / n)
m = torch.pow(m_0, torch.arange(1, 1 + n))
if n < n_heads:
m_hat_0 = 2.0 ** (-4.0 / n)
m_hat = torch.pow(m_hat_0, torch.arange(1, 1 + 2 * (n_heads - n), 2))
m = torch.cat([m, m_hat])
return m
class LongMIL(nn.Module):
def __init__(self, n_classes, input_size=384):
super(LongMIL, self).__init__()
self.n_heads = 4
self.input_size = input_size
feat_size = 384
self.feat_size = feat_size
self._fc1 = nn.Sequential(nn.Linear(input_size, feat_size), nn.ReLU())
self.n_classes = n_classes
self.layer1 = TransBlock(dim=feat_size, num_heads=self.n_heads)
self.layer2 = TransBlock(dim=feat_size, num_heads=self.n_heads)
self.layer3 = TransBlock(dim=feat_size, num_heads=self.n_heads)
# optional 1
# self.layer3 = LinearTransLayer(dim=feat_size,num_heads=self.n_heads,)
# optional 2
# self.layer3 = MambaLayer(dim=feat_size)
self.norm = nn.LayerNorm(feat_size)
self._fc2 = nn.Linear(feat_size, self.n_classes)
self.rotary = Rotary2D(dim=feat_size)
self.alibi = torch.load('./alibi_tensor_core.pt').cuda()
self.slope_m = get_slopes(self.n_heads).to(torch.float16)
def forward(self, x):
'''
:param x: shape = N * (feat_size+2), 2 for x-y 2d-position
:return:
'''
h = x[:, :self.input_size].unsqueeze(0) # shape 1*N*feat_size
h = self._fc1(h)
# option 1 use local_alibi +global_alibi
freqs_cis, alibi_local, alibi_global = self.positional_embedding(x, use_alibi=True)
# option 2 combine alibi and rope
# freqs_cis, alibi_local, alibi_global = self.positional_embedding(x,use_rope=True,use_alibi=True)
# option 3 use rope only
# freqs_cis, alibi_local, alibi_global = self.positional_embedding(x,use_rope=True)
h, _ = self.layer1(h, freqs_cis, alibi_local)
h, _ = self.layer2(h, freqs_cis, alibi_local)
h, attn = self.layer3(h, freqs_cis, alibi_global)
h = self.norm(h.mean(1))
# ---->predict
logits = self._fc2(h) # [1, n_classes]
Y_hat = torch.argmax(logits, dim=-1)
Y_prob = F.softmax(logits, dim=-1)
return logits, Y_hat, Y_prob, attn
def positional_embedding(self, x, use_alibi=False, use_rope=False):
# scale = 1 # for 20x 224 with 112 overlap (or 40x 224)
scale = 2 # for 20x 224 with 0 overlap
shape = 112 # or 128
# shape = 128
freqs_cis = None
alibi_bias = None
alibi_bias2 = None
if use_rope or use_alibi:
abs_pos = x[::, -2:]
# pdb.set_trace()
# print(abs_pos)
x_pos, y_pos = abs_pos[:, 0], abs_pos[:, 1]
x_pos = torch.round((x_pos - x_pos.min()) / (shape * scale) / 4)
y_pos = torch.round((y_pos - y_pos.min()) / (shape * scale) / 4)
H, W = 600 // scale, 600 // scale
selected_idx = (x_pos * W + y_pos).to(torch.int)
if use_rope:
pos_cached = self.rotary.forward(torch.tensor([H, W]))
freqs_cis = pos_cached[selected_idx].cuda()
if use_alibi:
alibi_bias = self.alibi[selected_idx, :][:, selected_idx]
alibi_bias = alibi_bias[:, :, None] * self.slope_m[None, None, :].cuda()
# pdb.set_trace()
alibi_bias = alibi_bias.permute(2, 0, 1).unsqueeze(0)#.float()
shape3 = alibi_bias.shape[3]
pad_num = 8 - shape3 % 8 # to tackle xformer problems
padding_bias = torch.zeros(1, alibi_bias.shape[1], alibi_bias.shape[2], pad_num).cuda()
alibi_bias = torch.cat([alibi_bias, padding_bias], dim=-1)
alibi_bias = autograd.Variable(alibi_bias.contiguous())[:, :, :, :shape3]
alibi_bias2 = copy.deepcopy(alibi_bias)
temp_min = alibi_bias2.min()
alibi_bias2[torch.where(alibi_bias2 == temp_min)] = -torch.inf # masked out longer distances
return freqs_cis, alibi_bias2, alibi_bias
if __name__ == "__main__":
data = torch.randn((12800, 384+2)).cuda() # 384 is feature, 2 is coordinates
model = LongMIL(n_classes=2).cuda()
print(model.eval())
results_dict = model(data)
# pdb.set_trace()
print(results_dict)