Test ability to modify dropout and fc modules of network
(net, img_size)
| 60 | |
| 61 | @pytest.mark.parametrize('img_size', [224, 256, 512]) |
| 62 | def test_modify_dropout(net, img_size): |
| 63 | """Test ability to modify dropout and fc modules of network""" |
| 64 | dropout = nn.Sequential(OrderedDict([ |
| 65 | ('_bn2', nn.BatchNorm1d(net._bn1.num_features)), |
| 66 | ('_drop1', nn.Dropout(p=net._global_params.dropout_rate)), |
| 67 | ('_linear1', nn.Linear(net._bn1.num_features, 512)), |
| 68 | ('_relu', nn.ReLU()), |
| 69 | ('_bn3', nn.BatchNorm1d(512)), |
| 70 | ('_drop2', nn.Dropout(p=net._global_params.dropout_rate / 2)) |
| 71 | ])) |
| 72 | fc = nn.Linear(512, net._global_params.num_classes) |
| 73 | |
| 74 | net._dropout = dropout |
| 75 | net._fc = fc |
| 76 | |
| 77 | data = torch.zeros((2, 3, img_size, img_size)) |
| 78 | output = net(data) |
| 79 | assert not torch.isnan(output).any() |
| 80 | |
| 81 | |
| 82 | @pytest.mark.parametrize('img_size', [224, 256, 512]) |