1
1
# This code references https://huggingface.co/JosephusCheung/ASimilarityCalculatior/blob/main/qwerty.py
2
2
# Fill in the path of the model to be queried and the root directory of the reference models, and this script will return the similarity between the model to be queried and all reference models.
3
- import sys ,os
3
+ import sys , os
4
4
import torch
5
5
import torch .nn as nn
6
6
import torch .nn .functional as F
7
7
8
+
8
9
def cal_cross_attn (to_q , to_k , to_v , rand_input ):
9
10
hidden_dim , embed_dim = to_q .shape
10
11
attn_to_q = nn .Linear (hidden_dim , embed_dim , bias = False )
@@ -16,41 +17,50 @@ def cal_cross_attn(to_q, to_k, to_v, rand_input):
16
17
17
18
return torch .einsum (
18
19
"ik, jk -> ik" ,
19
- F .softmax (torch .einsum ("ij, kj -> ik" , attn_to_q (rand_input ), attn_to_k (rand_input )), dim = - 1 ),
20
- attn_to_v (rand_input )
20
+ F .softmax (
21
+ torch .einsum ("ij, kj -> ik" , attn_to_q (rand_input ), attn_to_k (rand_input )),
22
+ dim = - 1 ,
23
+ ),
24
+ attn_to_v (rand_input ),
21
25
)
22
26
27
+
23
28
def model_hash (filename ):
24
29
try :
25
30
with open (filename , "rb" ) as file :
26
31
import hashlib
32
+
27
33
m = hashlib .sha256 ()
28
34
29
35
file .seek (0x100000 )
30
36
m .update (file .read (0x10000 ))
31
37
return m .hexdigest ()[0 :8 ]
32
38
except FileNotFoundError :
33
- return 'NOFILE'
39
+ return "NOFILE"
40
+
34
41
35
42
def eval (model , n , input ):
36
43
qk = f"enc_p.encoder.attn_layers.{ n } .conv_q.weight"
37
44
uk = f"enc_p.encoder.attn_layers.{ n } .conv_k.weight"
38
45
vk = f"enc_p.encoder.attn_layers.{ n } .conv_v.weight"
39
- atoq , atok , atov = model [qk ][:,:, 0 ], model [uk ][:,:, 0 ], model [vk ][:,:, 0 ]
46
+ atoq , atok , atov = model [qk ][:, :, 0 ], model [uk ][:, :, 0 ], model [vk ][:, :, 0 ]
40
47
41
48
attn = cal_cross_attn (atoq , atok , atov , input )
42
49
return attn
43
50
44
- def main (path ,root ):
51
+
52
+ def main (path , root ):
45
53
torch .manual_seed (114514 )
46
54
model_a = torch .load (path , map_location = "cpu" )["weight" ]
47
55
48
- print ("query:\t \t %s\t %s" % (path ,model_hash (path )))
56
+ print ("query:\t \t %s\t %s" % (path , model_hash (path )))
49
57
50
58
map_attn_a = {}
51
59
map_rand_input = {}
52
60
for n in range (6 ):
53
- hidden_dim , embed_dim ,_ = model_a [f"enc_p.encoder.attn_layers.{ n } .conv_v.weight" ].shape
61
+ hidden_dim , embed_dim , _ = model_a [
62
+ f"enc_p.encoder.attn_layers.{ n } .conv_v.weight"
63
+ ].shape
54
64
rand_input = torch .randn ([embed_dim , hidden_dim ])
55
65
56
66
map_attn_a [n ] = eval (model_a , n , rand_input )
@@ -59,7 +69,7 @@ def main(path,root):
59
69
del model_a
60
70
61
71
for name in sorted (list (os .listdir (root ))):
62
- path = "%s/%s" % (root ,name )
72
+ path = "%s/%s" % (root , name )
63
73
model_b = torch .load (path , map_location = "cpu" )["weight" ]
64
74
65
75
sims = []
@@ -70,9 +80,13 @@ def main(path,root):
70
80
sim = torch .mean (torch .cosine_similarity (attn_a , attn_b ))
71
81
sims .append (sim )
72
82
73
- print ("reference:\t %s\t %s\t %s" % (path ,model_hash (path ),f"{ torch .mean (torch .stack (sims )) * 1e2 :.2f} %" ))
83
+ print (
84
+ "reference:\t %s\t %s\t %s"
85
+ % (path , model_hash (path ), f"{ torch .mean (torch .stack (sims )) * 1e2 :.2f} %" )
86
+ )
87
+
74
88
75
89
if __name__ == "__main__" :
76
- query_path = r"weights\mi v3.pth"
77
- reference_root = r"weights"
78
- main (query_path ,reference_root )
90
+ query_path = r"weights\mi v3.pth"
91
+ reference_root = r"weights"
92
+ main (query_path , reference_root )
0 commit comments