@@ -146,21 +146,21 @@ def test_model_inference(model_name, batch_size):
146
146
rand_output = model (rand_tensors ['input' ])
147
147
rand_features = model .forward_features (rand_tensors ['input' ])
148
148
rand_pre_logits = model .forward_head (rand_features , pre_logits = True )
149
- assert torch .allclose (rand_output , rand_tensors ['output' ], rtol = 1e-3 , atol = 1e-4 )
150
- assert torch .allclose (rand_features , rand_tensors ['features' ], rtol = 1e-3 , atol = 1e-4 )
151
- assert torch .allclose (rand_pre_logits , rand_tensors ['pre_logits' ], rtol = 1e-3 , atol = 1e-4 )
149
+ assert torch .allclose (rand_output , rand_tensors ['output' ], rtol = 1e-3 , atol = 1e-4 ), 'rand output does not match'
150
+ assert torch .allclose (rand_features , rand_tensors ['features' ], rtol = 1e-3 , atol = 1e-4 ), 'rand features do not match'
151
+ assert torch .allclose (rand_pre_logits , rand_tensors ['pre_logits' ], rtol = 1e-3 , atol = 1e-4 ), 'rand pre_logits do not match'
152
152
153
- def _test_owl (owl_input ):
153
+ def _test_owl (owl_input , tol = ( 1e-3 , 1e-4 ) ):
154
154
owl_output = model (owl_input )
155
155
owl_features = model .forward_features (owl_input )
156
156
owl_pre_logits = model .forward_head (owl_features .clone (), pre_logits = True )
157
157
assert owl_output .softmax (1 ).argmax (1 ) == 24 # owl
158
- assert torch .allclose (owl_output , owl_tensors ['output' ], rtol = 1e-3 , atol = 1e-4 )
159
- assert torch .allclose (owl_features , owl_tensors ['features' ], rtol = 1e-3 , atol = 1e-4 )
160
- assert torch .allclose (owl_pre_logits , owl_tensors ['pre_logits' ], rtol = 1e-3 , atol = 1e-4 )
158
+ assert torch .allclose (owl_output , owl_tensors ['output' ], rtol = tol [ 0 ] , atol = tol [ 1 ]), 'owl output does not match'
159
+ assert torch .allclose (owl_features , owl_tensors ['features' ], rtol = tol [ 0 ] , atol = tol [ 1 ]), 'owl output does not match'
160
+ assert torch .allclose (owl_pre_logits , owl_tensors ['pre_logits' ], rtol = tol [ 0 ] , atol = tol [ 1 ]), 'owl output does not match'
161
161
162
162
_test_owl (owl_tensors ['input' ]) # test with original pp owl tensor
163
- _test_owl (pp (test_owl ).unsqueeze (0 )) # re-process from original jpg
163
+ _test_owl (pp (test_owl ).unsqueeze (0 ), tol = ( 1e-1 , 1e-1 )) # re-process from original jpg, Pillow output can change a lot btw ver
164
164
165
165
166
166
@pytest .mark .base
0 commit comments