SAG-ViT / tests /test_model_components.py

Added model files and updated config.json

039647a almost 2 years ago

2.08 kB

import unittest
import torch
from model_components import EfficientNetV2FeatureExtractor, GATGNN, TransformerEncoder, MLPBlock
from torch_geometric.data import Data

class TestModelComponents(unittest.TestCase):
    def test_efficientnetv2_extractor_output_shape(self):
        model = EfficientNetV2FeatureExtractor()
        model.eval()
        x = torch.randn(2, 3, 224, 224)
        with torch.no_grad():
            features = model(x)
        # Check output shape - depends on inception intermediate layer
        # Example: shape could be (2, 768, 8, 8) depending on the chosen layer
        self.assertEqual(features.size(0), 2)
        self.assertTrue(features.size(1) > 0)
        self.assertTrue(features.size(2) > 0)
        self.assertTrue(features.size(3) > 0)

    def test_gatgnn_forward(self):
        # Graph with 4 nodes, each node feature dim=256
        x = torch.randn(4, 256)
        edge_index = torch.tensor([[0,1,1,2],[1,0,2,3]], dtype=torch.long)
        batch = torch.tensor([0,0,0,0])
        data = Data(x=x, edge_index=edge_index, batch=batch)
        
        gnn = GATGNN(in_channels=256, hidden_channels=64, out_channels=32)
        output = gnn(data)
        # After pooling: should be (batch_size, out_channels) = (1,32)
        self.assertEqual(output.shape, (1, 32))

    def test_transformer_encoder(self):
        # (B, N, D) = (2, 10, 64)
        x = torch.randn(2, 10, 64)
        encoder = TransformerEncoder(d_model=64, nhead=4, num_layers=2, dim_feedforward=64)
        out = encoder(x)
        # same shape as input
        self.assertEqual(out.shape, (2, 10, 64))

    def test_mlp_block(self):
        mlp = MLPBlock(in_features=64, hidden_features=128, out_features=10)
        x = torch.randn(2, 64)
        out = mlp(x)
        self.assertEqual(out.shape, (2,10))

    def test_efficientnetv2_freeze(self):
        # Ensure params are frozen
        model = EfficientNetV2FeatureExtractor()
        for param in model.parameters():
            self.assertFalse(param.requires_grad)

if __name__ == '__main__':
    unittest.main()

	import unittest
	import torch
	from model_components import EfficientNetV2FeatureExtractor, GATGNN, TransformerEncoder, MLPBlock
	from torch_geometric.data import Data

	class TestModelComponents(unittest.TestCase):
	def test_efficientnetv2_extractor_output_shape(self):
	model = EfficientNetV2FeatureExtractor()
	model.eval()
	x = torch.randn(2, 3, 224, 224)
	with torch.no_grad():
	features = model(x)
	# Check output shape - depends on inception intermediate layer
	# Example: shape could be (2, 768, 8, 8) depending on the chosen layer
	self.assertEqual(features.size(0), 2)
	self.assertTrue(features.size(1) > 0)
	self.assertTrue(features.size(2) > 0)
	self.assertTrue(features.size(3) > 0)

	def test_gatgnn_forward(self):
	# Graph with 4 nodes, each node feature dim=256
	x = torch.randn(4, 256)
	edge_index = torch.tensor([[0,1,1,2],[1,0,2,3]], dtype=torch.long)
	batch = torch.tensor([0,0,0,0])
	data = Data(x=x, edge_index=edge_index, batch=batch)

	gnn = GATGNN(in_channels=256, hidden_channels=64, out_channels=32)
	output = gnn(data)
	# After pooling: should be (batch_size, out_channels) = (1,32)
	self.assertEqual(output.shape, (1, 32))

	def test_transformer_encoder(self):
	# (B, N, D) = (2, 10, 64)
	x = torch.randn(2, 10, 64)
	encoder = TransformerEncoder(d_model=64, nhead=4, num_layers=2, dim_feedforward=64)
	out = encoder(x)
	# same shape as input
	self.assertEqual(out.shape, (2, 10, 64))

	def test_mlp_block(self):
	mlp = MLPBlock(in_features=64, hidden_features=128, out_features=10)
	x = torch.randn(2, 64)
	out = mlp(x)
	self.assertEqual(out.shape, (2,10))

	def test_efficientnetv2_freeze(self):
	# Ensure params are frozen
	model = EfficientNetV2FeatureExtractor()
	for param in model.parameters():
	self.assertFalse(param.requires_grad)

	if __name__ == '__main__':
	unittest.main()