@@ -41,7 +41,7 @@ def test_rpn(self):
41
41
"loss_rpn_loc" : torch .tensor (0.0990132466 ),
42
42
}
43
43
for name in expected_losses .keys ():
44
- assert torch .allclose (proposal_losses [name ], expected_losses [name ])
44
+ self . assertTrue ( torch .allclose (proposal_losses [name ], expected_losses [name ]) )
45
45
46
46
expected_proposal_boxes = [
47
47
Boxes (torch .tensor ([[0 , 0 , 10 , 10 ], [7.3365392685 , 0 , 10 , 10 ]])),
@@ -63,13 +63,15 @@ def test_rpn(self):
63
63
torch .tensor ([0.1415634006 , 0.0989848152 , 0.0565387346 , - 0.0072308783 , - 0.0428492837 ]),
64
64
]
65
65
66
- for i in range (len (image_sizes )):
67
- assert len (proposals [i ]) == len (expected_proposal_boxes [i ])
68
- assert proposals [i ].image_size == (image_sizes [i ][0 ], image_sizes [i ][1 ])
69
- assert torch .allclose (
70
- proposals [i ].proposal_boxes .tensor , expected_proposal_boxes [i ].tensor
66
+ for proposal , expected_proposal_box , im_size , expected_objectness_logit in zip (
67
+ proposals , expected_proposal_boxes , image_sizes , expected_objectness_logits
68
+ ):
69
+ self .assertEqual (len (proposal ), len (expected_proposal_box ))
70
+ self .assertEqual (proposal .image_size , im_size )
71
+ self .assertTrue (
72
+ torch .allclose (proposal .proposal_boxes .tensor , expected_proposal_box .tensor )
71
73
)
72
- assert torch .allclose (proposals [ i ] .objectness_logits , expected_objectness_logits [ i ] )
74
+ self . assertTrue ( torch .allclose (proposal .objectness_logits , expected_objectness_logit ) )
73
75
74
76
def test_rrpn (self ):
75
77
torch .manual_seed (121 )
@@ -103,7 +105,7 @@ def test_rrpn(self):
103
105
"loss_rpn_loc" : torch .tensor (0.1552739739 ),
104
106
}
105
107
for name in expected_losses .keys ():
106
- assert torch .allclose (proposal_losses [name ], expected_losses [name ])
108
+ self . assertTrue ( torch .allclose (proposal_losses [name ], expected_losses [name ]) )
107
109
108
110
expected_proposal_boxes = [
109
111
RotatedBoxes (
@@ -185,25 +187,32 @@ def test_rrpn(self):
185
187
186
188
torch .set_printoptions (precision = 8 , sci_mode = False )
187
189
188
- for i in range (len (image_sizes )):
189
- assert len (proposals [i ]) == len (expected_proposal_boxes [i ])
190
- assert proposals [i ].image_size == (image_sizes [i ][0 ], image_sizes [i ][1 ])
190
+ for proposal , expected_proposal_box , im_size , expected_objectness_logit in zip (
191
+ proposals , expected_proposal_boxes , image_sizes , expected_objectness_logits
192
+ ):
193
+ self .assertEqual (len (proposal ), len (expected_proposal_box ))
194
+ self .assertEqual (proposal .image_size , im_size )
191
195
# It seems that there's some randomness in the result across different machines:
192
196
# This test can be run on a local machine for 100 times with exactly the same result,
193
197
# However, a different machine might produce slightly different results,
194
198
# thus the atol here.
195
199
err_msg = "computed proposal boxes = {}, expected {}" .format (
196
- proposals [ i ] .proposal_boxes .tensor , expected_proposal_boxes [ i ] .tensor
200
+ proposal .proposal_boxes .tensor , expected_proposal_box .tensor
197
201
)
198
- assert torch .allclose (
199
- proposals [i ].proposal_boxes .tensor , expected_proposal_boxes [i ].tensor , atol = 1e-5
200
- ), err_msg
202
+ self .assertTrue (
203
+ torch .allclose (
204
+ proposal .proposal_boxes .tensor , expected_proposal_box .tensor , atol = 1e-5
205
+ ),
206
+ err_msg ,
207
+ )
208
+
201
209
err_msg = "computed objectness logits = {}, expected {}" .format (
202
- proposals [i ].objectness_logits , expected_objectness_logits [i ]
210
+ proposal .objectness_logits , expected_objectness_logit
211
+ )
212
+ self .assertTrue (
213
+ torch .allclose (proposal .objectness_logits , expected_objectness_logit , atol = 1e-5 ),
214
+ err_msg ,
203
215
)
204
- assert torch .allclose (
205
- proposals [i ].objectness_logits , expected_objectness_logits [i ], atol = 1e-5
206
- ), err_msg
207
216
208
217
209
218
if __name__ == "__main__" :
0 commit comments