Coverage for tests/tests_advertrain/test_dependencies/test_dropblock.py: 100%

1import pytest 

2import torch 

3 

4from robustML.advertrain.dependencies.dropblock import (DropBlock2d, 

5 drop_block_2d, 

6 drop_block_fast_2d) 

7 

8torch.manual_seed(0) 

9 

10 

11@pytest.fixture 

12def sample_input_tensor(): 

13 return torch.randn(1, 3, 24, 24) 

14 

15 

16def test_drop_block_2d(sample_input_tensor): 

17 drop_prob = 0.1 

18 block_size = 7 

19 output = drop_block_2d(sample_input_tensor, drop_prob=drop_prob, block_size=block_size) 

20 

21 assert output.shape == sample_input_tensor.shape 

22 

23 

24def test_drop_block_fast_2d(sample_input_tensor): 

25 drop_prob = 0.1 

26 block_size = 7 

27 output = drop_block_fast_2d(sample_input_tensor, drop_prob=drop_prob, block_size=block_size) 

28 

29 assert output.shape == sample_input_tensor.shape 

30 

31 

32def test_DropBlock2d_initialization(): 

33 drop_block = DropBlock2d(drop_prob=0.1, block_size=7, gamma_scale=1.0) 

34 assert drop_block.drop_prob == 0.1 

35 assert drop_block.block_size == 7 

36 assert drop_block.gamma_scale == 1.0 

37 

38 

39def test_DropBlock2d_forward(sample_input_tensor): 

40 drop_block = DropBlock2d(drop_prob=0.1, block_size=7, gamma_scale=1.0) 

41 drop_block.train() # Set to training mode 

42 

43 output_train = drop_block(sample_input_tensor) 

44 assert output_train.shape == sample_input_tensor.shape 

45 dropped_elements_train = (sample_input_tensor != output_train).float().mean().item() 

46 assert dropped_elements_train > 0 

47 

48 drop_block.eval() 

49 output_eval = drop_block(sample_input_tensor) 

50 assert output_eval.shape == sample_input_tensor.shape 

51 assert torch.equal(output_eval, sample_input_tensor)