@@ -146,21 +146,21 @@ def test_model_inference(model_name, batch_size):
146146 rand_output = model (rand_tensors ['input' ])
147147 rand_features = model .forward_features (rand_tensors ['input' ])
148148 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'
152152
153- def _test_owl (owl_input ):
153+ def _test_owl (owl_input , tol = ( 1e-3 , 1e-4 ) ):
154154 owl_output = model (owl_input )
155155 owl_features = model .forward_features (owl_input )
156156 owl_pre_logits = model .forward_head (owl_features .clone (), pre_logits = True )
157157 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'
161161
162162 _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
164164
165165
166166@pytest .mark .base
0 commit comments