FazBrowse GitHub Viewer
|
Trending
|
URL:
|
Home
Tools:
[Download Repo ZIP]
[View Raw Code]
[Original HTTPS Page]
RadioDiff/prune.py at main · UNIC-Lab/RadioDiff · GitHub
UNIC-Lab
RadioDiff
Repository navigation
Code
Issues
2
(2)
Pull requests
Actions
Projects
Security and quality
Insights
Expand file tree
Breadcrumbs
RadioDiff
/
prune.py
Copy path
More file actions
More file actions
Latest commit
History
History
History
272 lines (240 loc) · 11.1 KB
Breadcrumbs
RadioDiff
/
prune.py
Copy path
File metadata and controls
272 lines (240 loc) · 11.1 KB
Raw
Copy raw file
Download raw file
Open symbols panel
Edit and raw actions
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
from
fnmatch
import
fnmatch
import
copy
import
math
import
numpy
as
np
import
torch
import
torch
.
nn
as
nn
def
replace_module
(
model
,
name
,
new_m
):
r"""Replace the module <name> in <model> with <new_m>
E.g., 'module.layer1.0.conv1'
==> model.__getattr__('module').__getattr__("layer1").__getitem__(0).__setattr__('conv1', new_m)
"""
obj
=
model
segs
=
name
.
split
(
"."
)
for
ix
,
s
in
enumerate
(
segs
):
if
ix
==
len
(
segs
)
-
1
:
# the last one
if
s
.
isdigit
():
obj
.
__setitem__
(
int
(
s
),
new_m
)
else
:
obj
.
__setattr__
(
s
,
new_m
)
return
if
s
.
isdigit
():
obj
=
obj
.
__getitem__
(
int
(
s
))
else
:
obj
=
obj
.
__getattr__
(
s
)
def
structured_pruning
(
model
,
pr
=
0.4375
,
reinit
=
False
,
layer_prefix
=
'model.'
):
"""Structured pruning of a given model.
Args:
model: nn.Module, model to prune.
pr: Pruning ratio. A float to indiate layerwise sparsity.
reinit: Reinitialize the weights of the pruned model.
Returns:
new_model: The pruned model.
"""
weight_types
=
(
nn
.
Conv2d
,
nn
.
Linear
)
embedding_type
=
(
nn
.
Embedding
,)
norm_types
=
(
nn
.
LayerNorm
,
nn
.
BatchNorm2d
,
nn
.
GroupNorm
)
learnable_types
=
weight_types
+
embedding_type
+
norm_types
# Name of the layers that are prunable.
target_layers
=
(
layer_prefix
+
'init_conv_mask.'
,
layer_prefix
+
'final_layernorm'
,
layer_prefix
+
'embed_tokens*'
,
layer_prefix
+
'mm_projector.2*'
,
'lm_head'
)
# For first and last layers, only one dimension is shrinked.
first_layers
=
(
layer_prefix
+
'embed_tokens'
,
layer_prefix
+
'mm_projector.2*'
,
)
last_layers
=
(
'lm_head'
,
)
# new_model = copy.deepcopy(model) # @huanwangx: Cannot use this. Don't know why.
new_model
=
model
# Adjust num of heads in attention.
for
name
,
m
in
new_model
.
named_modules
():
if
fnmatch
(
name
,
layer_prefix
+
'layers.*.self_attn'
):
num_heads
=
m
.
num_heads
-
int
(
np
.
ceil
(
m
.
num_heads
*
pr
))
num_key_value_heads
=
m
.
num_key_value_heads
-
int
(
np
.
ceil
(
m
.
num_key_value_heads
*
pr
))
hidden_size
=
m
.
hidden_size
-
int
(
np
.
ceil
(
m
.
hidden_size
*
pr
))
print
(
f'
{
name
}
: num_heads adjusted from
{
m
.
num_heads
}
to
{
num_heads
}
'
)
print
(
f'
{
name
}
: num_key_value_heads adjusted from
{
m
.
num_key_value_heads
}
to
{
num_key_value_heads
}
'
)
print
(
f'
{
name
}
: hidden_size adjusted from
{
m
.
hidden_size
}
to
{
hidden_size
}
'
)
m
.
num_heads
=
num_heads
m
.
num_key_value_heads
=
num_key_value_heads
m
.
hidden_size
=
hidden_size
# Pruning layer by layer.
old_params
,
new_params
=
{}, {}
for
name
,
m
in
new_model
.
named_modules
():
if
isinstance
(
m
,
learnable_types
):
old_params
[
name
]
=
m
.
weight
.
numel
()
prune_current_layer
=
any
([
fnmatch
(
name
,
pattern
)
for
pattern
in
target_layers
])
#print(f"Pruning layer: {name}, Prune current layer: {prune_current_layer}")
prune_current_layer
=
True
if
not
prune_current_layer
:
new_params
[
name
]
=
m
.
weight
.
numel
()
else
:
w
=
m
.
weight
bias
=
hasattr
(
m
,
"bias"
)
and
m
.
bias
is
not
None
num_filters
=
w
.
shape
[
0
]
num_channels
=
w
.
shape
[
1
]
if
len
(
w
.
shape
)
>
1
else
0
num_kept_filters
=
int
(
num_filters
-
np
.
ceil
(
pr
*
num_filters
))
num_kept_channels
=
int
(
num_channels
-
np
.
ceil
(
pr
*
num_channels
))
# Handle the first / last layers.
for
pattern
in
first_layers
:
if
fnmatch
(
name
,
pattern
):
if
isinstance
(
m
,
norm_types
+
embedding_type
):
num_kept_filters
=
num_filters
else
:
num_kept_channels
=
num_channels
for
pattern
in
last_layers
:
if
fnmatch
(
name
,
pattern
):
num_kept_filters
=
num_filters
print
(
f'Current layer to prune:
{
name
}
| Shape:
{
list
(
w
.
shape
)
}
| '
f'num_kept_filters:
{
num_kept_filters
}
'
f'num_kept_channels:
{
num_kept_channels
}
'
f'PR:
{
pr
}
'
)
# Evenly-spaced pruning.
kept_filters
=
[]
space
=
num_filters
//
num_kept_filters
for
i
in
range
(
num_kept_filters
):
kept_filters
.
append
(
i
*
space
)
kept_channels
=
[]
if
num_channels
>
0
:
space
=
num_channels
//
num_kept_channels
for
i
in
range
(
num_kept_channels
):
kept_channels
.
append
(
i
*
space
)
if
isinstance
(
m
,
nn
.
Conv2d
):
new_layer
=
nn
.
Conv2d
(
num_channels
,
len
(
kept_filters
),
m
.
kernel_size
,
m
.
stride
,
m
.
padding
,
m
.
dilation
,
m
.
groups
,
bias
,
)
kept_weights
=
m
.
weight
.
data
[
kept_filters
][:,
kept_channels
, :, :]
if
not
reinit
:
new_layer
.
weight
.
data
.
copy_
(
kept_weights
)
# load weights into the new module
if
bias
:
kept_bias
=
m
.
bias
.
data
[
kept_filters
]
new_layer
.
bias
.
data
.
copy_
(
kept_bias
)
else
:
print
(
f"Layer
{
name
}
is reinited when building the new model!"
)
elif
isinstance
(
m
,
nn
.
Linear
):
new_layer
=
nn
.
Linear
(
in_features
=
len
(
kept_channels
),
out_features
=
len
(
kept_filters
),
bias
=
bias
,
)
kept_weights
=
m
.
weight
.
data
[
kept_filters
][:,
kept_channels
]
if
not
reinit
:
new_layer
.
weight
.
data
.
copy_
(
kept_weights
)
# load weights into the new module
if
bias
:
kept_bias
=
m
.
bias
.
data
[
kept_filters
]
new_layer
.
bias
.
data
.
copy_
(
kept_bias
)
else
:
print
(
f"Layer
{
name
}
is reinited when building the new model!"
)
elif
isinstance
(
m
,
nn
.
Embedding
):
# https://pytorch.org/docs/stable/generated/torch.nn.Embedding.html
new_layer
=
nn
.
Embedding
(
num_embeddings
=
m
.
num_embeddings
,
embedding_dim
=
num_kept_channels
,
padding_idx
=
m
.
padding_idx
,
)
kept_weights
=
m
.
weight
.
data
[
kept_filters
][:,
kept_channels
]
if
not
reinit
:
new_layer
.
weight
.
data
.
copy_
(
kept_weights
)
# load weights into the new module
if
bias
:
kept_bias
=
m
.
bias
.
data
[
kept_filters
]
new_layer
.
bias
.
data
.
copy_
(
kept_bias
)
else
:
print
(
f"Layer
{
name
}
is reinited when building the new model!"
)
elif
isinstance
(
m
,
nn
.
BatchNorm2d
):
new_layer
=
nn
.
BatchNorm2d
(
len
(
kept_filter
),
eps
=
m
.
eps
,
momentum
=
m
.
momentum
,
affine
=
m
.
affine
,
track_running_stats
=
m
.
track_running_stats
,
)
# copy bn weight and bias
new_layer
.
weight
.
data
.
copy_
(
m
.
weight
.
data
[
kept_filter
])
new_layer
.
bias
.
data
.
copy_
(
m
.
bias
.
data
[
kept_filter
])
# copy bn running stats
new_layer
.
running_mean
.
data
.
copy_
(
m
.
running_mean
[
kept_filter
])
new_layer
.
running_var
.
data
.
copy_
(
m
.
running_var
[
kept_filter
])
new_layer
.
num_batches_tracked
.
data
.
copy_
(
m
.
num_batches_tracked
)
elif
isinstance
(
m
,
nn
.
GroupNorm
):
kept_filters
=
list
(
range
(
m
.
num_channels
))
new_layer
=
nn
.
GroupNorm
(
num_groups
=
m
.
num_groups
,
num_channels
=
len
(
kept_filters
),
eps
=
m
.
eps
,
affine
=
m
.
affine
,
)
new_layer
.
weight
.
data
.
copy_
(
m
.
weight
.
data
[
kept_filters
])
new_layer
.
bias
.
data
.
copy_
(
m
.
bias
.
data
[
kept_filters
])
elif
isinstance
(
m
,
nn
.
LayerNorm
):
# See https://pytorch.org/docs/stable/generated/torch.nn.LayerNorm.html
new_layer
=
nn
.
LayerNorm
(
len
(
kept_filters
),
eps
=
m
.
eps
,
elementwise_affine
=
m
.
elementwise_affine
,
)
new_layer
.
weight
.
data
.
copy_
(
m
.
weight
.
data
[
kept_filters
])
new_layer
.
bias
.
data
.
copy_
(
m
.
bias
.
data
[
kept_filters
])
elif
isinstance
(
m
,
nn
.
MultiheadAttention
):
raise
NotImplementedError
else
:
raise
NotImplementedError
# Load the new_layer into the new_model
new_layer
=
new_layer
.
cuda
().
half
()
replace_module
(
new_model
,
name
,
new_layer
)
new_params
[
name
]
=
new_layer
.
weight
.
numel
()
# Get model params.
old_total_params
=
sum
(
list
(
old_params
.
values
()))
/
1e9
new_total_params
=
sum
(
list
(
new_params
.
values
()))
/
1e9
old_llm_params
,
old_vt_params
=
0
,
0
new_llm_params
,
new_vt_params
=
0
,
0
for
k
in
old_params
:
if
k
.
startswith
(
layer_prefix
+
'layers'
):
old_llm_params
+=
old_params
[
k
]
/
1e9
new_llm_params
+=
new_params
[
k
]
/
1e9
elif
k
.
startswith
(
layer_prefix
+
'vision_tower'
):
old_vt_params
+=
old_params
[
k
]
/
1e9
new_vt_params
+=
new_params
[
k
]
/
1e9
start_line
=
'-'
*
10
+
' Pruning Summary '
+
'-'
*
10
print
(
'
\n
'
+
start_line
)
print
(
f'Layerwise pruning ratio:
{
pr
}
'
)
print
(
f'Old total params:
{
old_total_params
:.2f
}
G, '
f'New total params:
{
new_total_params
:.2f
}
G '
f'(compression ratio:
{
old_total_params
/
new_total_params
:.2f
}
x) '
)
# print(f'Old llm params: {old_llm_params:.2f}G, '
# f'New llm params: {new_llm_params:.2f}G '
# f'(compression ratio: {old_llm_params / new_llm_params:.2f}x) '
# )
# print(f'Old vt params: {old_vt_params:.2f}G, '
# f'New vt params: {new_vt_params:.2f}G '
# f'(compression ratio: {old_vt_params / new_vt_params:.2f}x) '
# )
print
(
'-'
*
len
(
start_line
)
+
'
\n
'
)
return
new_model
Back
|
FazBrowse Home
|
New Git URL