def test_model_forward_torchscript(model_name, batch_size): """Run a single forward pass with each model""" with set_scriptable(True): model = create_model(model_name, pretrained=False) model.eval() if has_model_default_key(model_name, 'fixed_input_size'): input_size = get_model_default_value(model_name, 'input_size') elif has_model_default_key(model_name, 'min_input_size'): input_size = get_model_default_value(model_name, 'min_input_size') else: input_size = (3, 128, 128) # jit compile is already a bit slow and we've tested normal res already... model = torch.jit.script(model) outputs = model(torch.randn((batch_size, *input_size))) assert outputs.shape[0] == batch_size assert not torch.isnan(outputs).any(), 'Output included NaNs'
def test_model_forward_features(model_name, batch_size): """Run a single forward pass with each model in feature extraction mode""" model = create_model(model_name, pretrained=False, features_only=True) model.eval() expected_channels = model.feature_info.channels() assert len(expected_channels) >= 4 # all models here should have at least 4 feature levels by default, some 5 or 6 if has_model_default_key(model_name, 'fixed_input_size'): input_size = get_model_default_value(model_name, 'input_size') elif has_model_default_key(model_name, 'min_input_size'): input_size = get_model_default_value(model_name, 'min_input_size') else: input_size = (3, 96, 96) # jit compile is already a bit slow and we've tested normal res already... outputs = model(torch.randn((batch_size, *input_size))) assert len(expected_channels) == len(outputs) for e, o in zip(expected_channels, outputs): assert e == o.shape[1] assert o.shape[0] == batch_size assert not torch.isnan(o).any()