Skip to content

Commit 5e16b3d

Browse files
committed
test for both cpu and gpu
1 parent b7a5b48 commit 5e16b3d

1 file changed

Lines changed: 56 additions & 78 deletions

File tree

pytorch_binding/test/test.py

Lines changed: 56 additions & 78 deletions
Original file line numberDiff line numberDiff line change
@@ -18,14 +18,10 @@
1818
from transducer_np import RNNTLoss as rnntloss
1919

2020
parser = 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')
2521
parser.add_argument('--np', default=False, action='store_true', help='numpy loss')
2622
args = parser.parse_args()
2723

28-
fn = rnntloss() if args.np else RNNTLoss()
24+
fn = rnntloss() if args.np else RNNTLoss(size_average=False)
2925

3026
def 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

6251
def 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-
186164
if __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

Comments
 (0)