1818from transducer_np import RNNTLoss as rnntloss
1919
2020parser = argparse .ArgumentParser (description = 'MXNet RNN Transducer Test.' )
21- parser .add_argument ('B' , type = int , default = 1 , help = 'batch size' )
22- parser .add_argument ('T' , type = int , default = 300 , help = 'time step' )
23- parser .add_argument ('U' , type = int , default = 100 , help = 'prediction step' )
24- parser .add_argument ('V' , type = int , default = 60 , help = 'vocab size' )
2521parser .add_argument ('--np' , default = False , action = 'store_true' , help = 'numpy loss' )
2622args = parser .parse_args ()
2723
28- fn = rnntloss () if args .np else RNNTLoss ()
24+ fn = rnntloss () if args .np else RNNTLoss (size_average = False )
2925
3026def wrap_and_call (acts , labels ):
3127 acts = torch .FloatTensor (acts )
@@ -38,48 +34,47 @@ def wrap_and_call(acts, labels):
3834 labels = autograd .Variable (torch .IntTensor (labels ))
3935 lengths = autograd .Variable (torch .IntTensor (lengths ))
4036 label_lengths = autograd .Variable (torch .IntTensor (label_lengths ))
37+ if use_cuda :
38+ labels = labels .cuda ()
39+ lengths = lengths .cuda ()
40+ label_lengths = label_lengths .cuda ()
4141
4242 log_probs = nn .functional .log_softmax (acts , dim = 3 )
43- def grad_hook (grad ):
44- log_probs .saved_grad = grad .clone ()
45- log_probs .register_hook (grad_hook )
4643
4744 costs = fn (log_probs , labels , lengths , label_lengths )
4845 cost = torch .sum (costs )
4946 cost .backward ()
50- grads = log_probs .saved_grad
51- if use_cuda :
52- costs = costs .cpu ()
53- grads = grads .cpu ()
54- for i , a in enumerate (acts .grad .data .numpy ().reshape (- 1 )):
55- if i % 6 == 0 : print (end = '\n ' )
56- print ('{:.6f}, ' .format (a ), end = '' )
57- print ()
58- # print(acts.grad.data.numpy())
59- return costs .data .numpy (), grads .data .numpy ()
47+ print (repr (acts .grad .data .cpu ().numpy ()))
48+ return costs .data .cpu ().numpy (), acts .grad .data .cpu ().numpy ()
6049
6150
6251def small_test ():
63- acts = np .array ([[[0.1 , 0.6 , 0.1 , 0.1 , 0.1 ],
52+ acts = np .array ([[[[ 0.1 , 0.6 , 0.1 , 0.1 , 0.1 ],
6453 [0.1 , 0.1 , 0.6 , 0.1 , 0.1 ],
6554 [0.1 , 0.1 , 0.2 , 0.8 , 0.1 ]],
6655 [[0.1 , 0.6 , 0.1 , 0.1 , 0.1 ],
6756 [0.1 , 0.1 , 0.2 , 0.1 , 0.1 ],
68- [0.7 , 0.1 , 0.2 , 0.1 , 0.1 ]]])
57+ [0.7 , 0.1 , 0.2 , 0.1 , 0.1 ]]]] )
6958 labels = [[1 , 2 ]]
7059
71- acts = acts [None , ...]
72-
7360 cost , grads = wrap_and_call (acts , labels )
7461 expected_cost = 4.495666
75- expected_grads = np .array ([[[- 0.308198071906 , - 0.6918019280939998 , 0.0 , 0.0 , 0.0 ],
76- [- 0.308198071906 , 0.0 , - 0.3836038561880001 , 0.0 , 0.0 ],
77- [- 0.3836038561880001 , 0.0 , 0.0 , 0.0 , 0.0 ]],
78- [[0.0 , - 0.308198071906 , 0.0 , 0.0 , 0.0 ],
79- [0.0 , 0.0 , - 0.6163961438119995 , 0.0 , 0.0 ],
80- [- 0.9999999999999991 , 0.0 , 0.0 , 0.0 , 0.0 ]]])
62+ expected_grads = np .array ([[[[- 0.13116688 , - 0.3999269 , 0.17703125 , 0.17703125 ,
63+ 0.17703125 ],
64+ [- 0.18572757 , 0.12247056 , - 0.18168412 , 0.12247056 ,
65+ 0.12247056 ],
66+ [- 0.32091254 , 0.06269141 , 0.06928472 , 0.12624499 ,
67+ 0.06269141 ]],
68+
69+ [[ 0.05456069 , - 0.21824276 , 0.05456069 , 0.05456069 ,
70+ 0.05456069 ],
71+ [ 0.12073959 , 0.12073959 , - 0.48295835 , 0.12073959 ,
72+ 0.12073959 ],
73+ [- 0.6925882 , 0.16871116 , 0.18645467 , 0.16871116 ,
74+ 0.16871116 ]]]])
8175 assert np .allclose (cost , expected_cost , rtol = 1e-6 ), \
8276 "small_test costs mismatch."
77+ print (grads ) # TODO change this exptected grad to actis.
8378 assert np .allclose (grads , expected_grads ), \
8479 "small_test gradient mismatch."
8580
@@ -121,38 +116,38 @@ def big_test():
121116 [0.6607698886038497 , 0.3771277082495921 , 0.3580209022231813 ]]]]
122117
123118 expected_costs = [4.2806528590890736 , 3.9384369822503591 ]
124- expected_grads = [
125- [[[- 0.4322264564338117 , - 0.5677735435661883 , 0.0 ],
126- [- 0.36565009313836844 , 0.0 , - 0.20212345042782007 ],
127- [- 0.20212345042782007 , 0.0 , 0.0 ]],
119+ expected_grads = [[[[- 1.86843902e-01 , - 6.25548810e-02 , 2.49398798e-01 ],
120+ [- 2.03376666e-01 , 2.02399328e-01 , 9.77333169e-04 ],
121+ [- 1.41016081e-01 , 7.91234672e-02 , 6.18926100e-02 ]],
128122
129- [[- 0.16521672442463506 , - 0.2670097320091765 , 0.0 ],
130- [ - 0.3943653886107811 , 0.0 , - 0.2382944365367636 ],
131- [ - 0.44041788696458367 , 0.0 , 0.0 ]],
123+ [[- 1.15517676e-02 , - 8.12802389e-02 , 9.28319991e-02 ],
124+ [ - 1.54257029e-01 , 2.29432687e-01 , - 7.51756504e-02 ],
125+ [ - 2.46593088e-01 , 1.46404594e-01 , 1.00188486e-01 ]],
132126
133- [[- 0.052129794015740985 , - 0.11308693040889405 , 0.0 ],
134- [ - 0.18313786985332664 , 0.0 , - 0.3243144491663483 ],
135- [ - 0.7647323361309323 , 0.0 , 0.0 ]],
127+ [[- 1.29182907e-02 , - 6.15932420e-02 , 7.45115355e-02 ],
128+ [ - 5.59857301e-02 , 2.19830811e-01 , - 1.63845062e-01 ],
129+ [ - 4.97626871e-01 , 2.09239945e-01 , 2.88386941e-01 ]],
136130
137- [[ 0.0 , - 0.052129794015740985 , 0.0 ],
138- [ 0.0 , 0.0 , - 0.23526766386906767 ],
139- [ - 1.0 , 0.0 , 0.0 ]]],
131+ [[ 1.36048580e-02 , - 3.02196294e-02 , 1.66147724e-02 ],
132+ [ 1.13924511e-01 , 6.27811998e-02 , - 1.76705718e-01 ],
133+ [ - 6.67078257e-01 , 3.67658824e-01 , 2.99419403e-01 ]]],
140134
141- [[[- 0.7161424128232795 , - 0.2838575871767207 , 0.0 ],
142- [- 0.18382932237365335 , - 0.10002826480306751 , 0.0 ],
143- [- 0.10002826480306751 , 0.0 , 0.0 ]],
144135
145- [[ - 0.41121794618117213 , - 0.3049244666421072 , 0.0 ],
146- [ - 0.3295759402552584 , - 0.15917784876050195 , 0.0 ],
147- [ - 0.2592061135635692 , 0.0 , 0.0 ]],
136+ [[[ - 3.56343776e-01 , - 5.53474613e-02 , 4.11691219e-01 ],
137+ [ - 9.69219357e-02 , 2.94591039e-02 , 6.74628317e-02 ],
138+ [ - 6.35175705e-02 , 2.76544970e-02 , 3.58630717e-02 ]],
148139
149- [[- 0.11607642141651396 , - 0.29514152476465827 , 0.0 ],
150- [ - 0.2865333615432337 , - 0.3381841034766833 , 0.0 ],
151- [ - 0.5973902170402529 , 0.0 , 0.0 ]],
140+ [[- 1.54499024e-01 , - 7.39420280e-02 , 2.28441030e-01 ],
141+ [ - 1.66789949e-01 , - 8.78955179e-05 , 1.66877866e-01 ],
142+ [ - 1.72369644e-01 , 1.05565332e-01 , 6.68043196e-02 ]],
152143
153- [[0.0 , - 0.11607642141651396 , 0.0 ],
154- [0.0 , - 0.4026097829597475 , 0.0 ],
155- [- 1.0 , 0.0 , 0.0 ]]]]
144+ [[ 2.38748826e-02 , - 1.18255816e-01 , 9.43809375e-02 ],
145+ [- 1.04707085e-01 , - 1.08934477e-01 , 2.13641584e-01 ],
146+ [- 3.69844258e-01 , 1.80118099e-01 , 1.89726159e-01 ]],
147+
148+ [[ 2.57137045e-02 , - 7.94617534e-02 , 5.37480488e-02 ],
149+ [ 1.22328237e-01 , - 2.38788679e-01 , 1.16460443e-01 ],
150+ [- 5.98686993e-01 , 3.02203178e-01 , 2.96483815e-01 ]]]]
156151
157152 activations = np .array (activations )
158153 labels = [[1 , 2 ],
@@ -163,33 +158,16 @@ def big_test():
163158 assert np .allclose (costs , sum (expected_costs )), \
164159 "big_test average costs mismatch."
165160
166- assert np .allclose (grads , expected_grads ), \
161+ assert np .allclose (grads , expected_grads , rtol = 1e-3 ), \
167162 "big_test grads for average cost mismatch."
168163
169- def time_test (blank = 0 ):
170- batch_size = args .B
171- vocab_size = args .V
172- input_len = args .T
173- output_len = args .U
174- acts = np .random .rand (batch_size , input_len , output_len + 1 , vocab_size )
175- labels = np .random .randint (1 , vocab_size , (batch_size , output_len ))
176-
177- start = time .time ()
178- iters = 10
179- for _ in range (iters ):
180- wrap_and_call (acts , labels )
181- end = time .time ()
182-
183- print ("Time per iteration: {:.3f}(s)" .format ((end - start )/ iters ))
184-
185-
186164if __name__ == "__main__" :
187165 use_cuda = False
188- # small_test()
166+ small_test ()
189167 big_test ()
190- # print("CPU Tests passed!")
191- # if torch.cuda.is_available():
192- # use_cuda = True
193- # small_test()
194- # print("GPU Tests passed!" )
195- # time_test( )
168+ print ("CPU Tests passed!" )
169+ if torch .cuda .is_available ():
170+ use_cuda = True
171+ small_test ()
172+ big_test ( )
173+ print ( "GPU Tests passed!" )
0 commit comments