test_utils.py

import unittest
import numpy as np
import torch
from torch.nn.modules.distance import CosineSimilarity, PairwiseDistance

from few_shot.utils import *
from config import PATH


class TestDistance(unittest.TestCase):
    def test_query_support_distances(self):
        # Create some dummy data with easily verifiable distances
        q = 1  # 1 query per class
        k = 3  # 3 way classification
        d = 2  # embedding dimension of two

        query = torch.zeros([q * k, d], dtype=torch.double)
        query[0] = torch.Tensor([0, 0])
        query[1] = torch.Tensor([0, 1])
        query[2] = torch.Tensor([1, 0])
        support = torch.zeros([k, d], dtype=torch.double)
        support[0] = torch.Tensor([1, 1])
        support[1] = torch.Tensor([1, 2])
        support[2] = torch.Tensor([2, 2])

        distances = pairwise_distances(query, support, 'l2')
        self.assertEqual(
            distances.shape, (q * k, k),
            'Output should have shape (q * k, k).'
        )

        # Calculate squared distances by iterating through all query-support pairs
        for i, q_ in enumerate(query):
            for j, s_ in enumerate(support):
                self.assertEqual(
                    (q_ - s_).pow(2).sum(),
                    distances[i, j].item(),
                    'The jth column of the ith row should be the squared distance between the '
                    'ith query sample and the kth query sample'
                )

        # Create some dummy data with easily verifiable distances
        q = 1  # 1 query per class
        k = 3  # 3 way classification
        d = 2  # embedding dimension of two
        query = torch.zeros([q * k, d], dtype=torch.double)
        query[0] = torch.Tensor([1, 0])
        query[1] = torch.Tensor([0, 1])
        query[2] = torch.Tensor([1, 1])
        support = torch.zeros([k, d], dtype=torch.double)
        support[0] = torch.Tensor([1, 1])
        support[1] = torch.Tensor([-1, -1])
        support[2] = torch.Tensor([0, 2])

        distances = pairwise_distances(query, support, 'cosine')

        # Calculate distances by iterating through all query-support pairs
        for i, q_ in enumerate(query):
            for j, s_ in enumerate(support):
                self.assertTrue(
                    torch.isclose(1-CosineSimilarity(dim=0)(q_, s_), distances[i, j], atol=2e-8),
                    'The jth column of the ith row should be the squared distance between the '
                    'ith query sample and the kth query sample'
                )

    def test_no_nans_on_zero_vectors(self):
        """Cosine distance calculation involves a divide-through by vector magnitude which
        can divide by zeros to occur.
        """
        # Create some dummy data with easily verifiable distances
        q = 1  # 1 query per class
        k = 3  # 3 way classification
        d = 2  # embedding dimension of two
        query = torch.zeros([q * k, d], dtype=torch.double)
        query[0] = torch.Tensor([0, 0])  # First query sample is all zeros
        query[1] = torch.Tensor([0, 1])
        query[2] = torch.Tensor([1, 1])
        support = torch.zeros([k, d], dtype=torch.double)
        support[0] = torch.Tensor([1, 1])
        support[1] = torch.Tensor([-1, -1])
        support[2] = torch.Tensor([0, 0])  # Third support sample is all zeros

        distances = pairwise_distances(query, support, 'cosine')

        self.assertTrue(torch.isnan(distances).sum() == 0, 'Cosine distances between 0-vectors should not be nan')


class TestAutogradGraphRetrieval(unittest.TestCase):
    def test_retrieval(self):
        """Create a simple autograd graph and check that the output is what is expected"""
        x = torch.ones(2, 2, requires_grad=True)
        y = x + 2
        # The operation on the next line will create two edges because the y variable is
        # y variable is used twice.
        z = y * y
        out = z.mean()

        nodes, edges = autograd_graph(out)

        # This is quite a brittle test as it will break if the names of the autograd Functions change
        # TODO: Less brittle test

        expected_nodes = [
            'MeanBackward1',
            'MulBackward0',
            'AddBackward0',
            'AccumulateGrad',
        ]

        self.assertEqual(
            set(expected_nodes),
            set(n.__class__.__name__ for n in nodes),
            'autograd_graph() must return all nodes in the autograd graph.'
        )

        # Check for the existence of the expected edges
        expected_edges = [
            ('MulBackward0', 'MeanBackward1'),  # z = y * y, out = z.mean()
            ('AddBackward0', 'MulBackward0'),    # y = x + 2, z = y * y
            ('AccumulateGrad', 'AddBackward0'),   # x = torch.ones(2, 2, requires_grad=True), y = x + 2
        ]
        for e in edges:
            self.assertIn(
                (e[0].__class__.__name__, e[1].__class__.__name__),
                expected_edges,
                'autograd_graph() must return all edges in the autograd graph.'
            )

        # Check for two edges between the AddBackward node and the ThMulBackward node
        num_y_squared_edges = 0
        for e in edges:
            if e[0].__class__.__name__ == 'AddBackward0' and e[1].__class__.__name__ == 'MulBackward0':
                num_y_squared_edges += 1

        self.assertEqual(
            num_y_squared_edges,
            2,
            'autograd_graph() must return multiple edges between nodes if they exist.'
        )