|
|
import unittest
|
|
|
|
|
|
import torch
|
|
|
|
|
|
from gat.models.surrogate import build_surrogate, list_surrogates
|
|
|
from gat.models.surrogate.hooks import feat_col
|
|
|
|
|
|
|
|
|
class TestTVModels(unittest.TestCase):
|
|
|
|
|
|
@torch.no_grad()
|
|
|
def test_outputs_shape(self):
|
|
|
inputs = torch.rand(1, 3, 224, 224)
|
|
|
for surrogate_id in list_surrogates():
|
|
|
if surrogate_id not in ['inception_v3']:
|
|
|
outputs = build_surrogate(surrogate_id, pretrain=False)(inputs)
|
|
|
self.assertEqual(outputs.shape, (1, 1000))
|
|
|
|
|
|
|
|
|
class TestFeatCol(unittest.TestCase):
|
|
|
|
|
|
@torch.no_grad()
|
|
|
def test_feat_col(self):
|
|
|
testcases = [{
|
|
|
'surrogate_id': 'vgg16',
|
|
|
'feat_layer': 'features.16',
|
|
|
'input_shape': (1, 3, 224, 224),
|
|
|
'feat_shape': (1, 256, 28, 28)
|
|
|
}, {
|
|
|
'surrogate_id': 'vgg19',
|
|
|
'feat_layer': 'features.18',
|
|
|
'input_shape': (1, 3, 224, 224),
|
|
|
'feat_shape': (1, 256, 28, 28)
|
|
|
}, {
|
|
|
'surrogate_id': 'resnet152',
|
|
|
'feat_layer': 'layer2',
|
|
|
'input_shape': (1, 3, 224, 224),
|
|
|
'feat_shape': (1, 512, 28, 28)
|
|
|
}, {
|
|
|
'surrogate_id': 'densenet169',
|
|
|
'feat_layer': 'features.denseblock2',
|
|
|
'input_shape': (1, 3, 224, 224),
|
|
|
'feat_shape': (1, 512, 28, 28)
|
|
|
}]
|
|
|
for testcase in testcases:
|
|
|
model = build_surrogate(testcase['surrogate_id'], pretrain=False)
|
|
|
with feat_col(model, testcase['feat_layer']) as _feat_collecter:
|
|
|
inputs = torch.rand(testcase['input_shape'])
|
|
|
model(inputs)
|
|
|
model(inputs)
|
|
|
self.assertEqual(len(_feat_collecter), 2)
|
|
|
self.assertEqual(_feat_collecter.pop().shape,
|
|
|
testcase['feat_shape'])
|
|
|
self.assertEqual(_feat_collecter.pop().shape,
|
|
|
testcase['feat_shape'])
|
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
unittest.main()
|
|
|
|
