hypernet_similarity.py

import torch
import torch.nn.functional as F
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
import os
import itertools

# Function to get hypernetwork-generated weights based on input parameters
def get_hypernetwork_weights(model, input_params):
    # Ensure the model is in evaluation mode
    model.hypernet.eval()

    # Pass the input parameters through the hypernetwork
    with torch.no_grad():
        weights = model.hypernet(input_params.to(model.device))

    return weights

# Function to compute and save the similarity matrix of weights generated by the hypernetwork
def hypernet_similarity_matrix(model, device, dir_model):

    first_param_range = [0.12, 0.1375, 0.155]
    second_param_range = [0.0215, 0.0225, 0.0235]
    third_param = 0.55  # Fixed third parameter

    # Generate all combinations
    param_combinations = list(itertools.product(first_param_range, second_param_range))

    # Add the third parameter (fixed) to each combination and convert to tensors
    params = [torch.tensor([p1, p2, third_param], device=device).unsqueeze(0) for p1, p2 in param_combinations]

    # Get weights for each parameter set
    weights_dict = {}
    for i, param in enumerate(params):
        weight_dict = get_hypernetwork_weights(model, param)
        # Create a name for each parameter combination for the dictionary key
        param_name = f'OmM_{param[0, 0]:.3f}_OmB_{param[0, 1]:.4f}'
        weights_dict[param_name] = weight_dict

    # Flatten and store the weights for distance calculation
    flattened_weights = []
    for name, weights in weights_dict.items():
        flat_weights = torch.cat([w.flatten() for w in weights.values()], dim=0)
        flattened_weights.append(flat_weights)

    # Stack the flattened weights into a tensor
    flattened_weights_tensor = torch.stack(flattened_weights)

    # Compute pairwise Euclidean distance matrix
    distance_matrix = torch.cdist(flattened_weights_tensor, flattened_weights_tensor)

    similarity_matrix = 1 / (1 + distance_matrix)

    # Convert the distance matrix to numpy for plotting
    # distance_matrix_np = distance_matrix.cpu().detach().numpy()
    similarity_matrix_np = similarity_matrix.cpu().detach().numpy()

    # Plot the similarity matrix using a heatmap
    # plt.figure(figsize=(20, 16))
    plt.figure(figsize=(10, 8))
    # sns.heatmap(distance_matrix_np, annot=True, cmap='viridis', 
    #             xticklabels=weights_dict.keys(), yticklabels=weights_dict.keys())
    sns.heatmap(similarity_matrix_np, annot=True, cmap='viridis', 
            xticklabels=weights_dict.keys(), yticklabels=weights_dict.keys())
    
    plt.xticks(rotation=45, ha="right")  # Rotate for x-axis labels
    plt.yticks(rotation=0)  # No rotation for y-axis labels

    plt.title('Similarity Matrix Based on Hypernetwork Weights')
    plt.xlabel('Parameters')
    plt.ylabel('Parameters')

    # Save the plot to the specified directory
    if not os.path.exists(dir_model):
        os.makedirs(dir_model)
    plot_filename = os.path.join(dir_model, "similarity_matrix.png")
    plt.savefig(plot_filename, bbox_inches="tight")
    plt.close()

    print(f"Similarity matrix saved at: {plot_filename}")