""" Script to generate gaussian fit pdf and statistics for the visualization results. Tony Fu, June 29, 2022 """ import os import sys import numpy as np import torch.nn as nn from torchvision import models # from torchvision.models import AlexNet_Weights, VGG16_Weights import matplotlib.pyplot as plt from matplotlib.backends.backend_pdf import PdfPages from tqdm import tqdm sys.path.append('../../..') from src.rf_mapping.gaussian_fit import (gaussian_fit, calc_f_explained_var, theta_to_ori) from src.rf_mapping.gaussian_fit import GaussianFitParamFormat as ParamFormat from src.rf_mapping.hook import ConvUnitCounter from src.rf_mapping.spatial import get_rf_sizes from src.rf_mapping.reproducibility import set_seeds import src.rf_mapping.constants as c # Please specify some details here: set_seeds() model = models.alexnet(pretrained=False).to(c.DEVICE) model_name = 'alexnet' # model = models.vgg16(pretrained=True).to(c.DEVICE) # model_name = "vgg16" # model = models.resnet18(pretrained=True).to(c.DEVICE) # model_name = "resnet18" sum_modes = ['abs'] this_is_a_test_run = False is_random = True # Please double-check the directories: if is_random: backprop_sum_dir = os.path.join(c.REPO_DIR, 'results', 'ground_truth', 'backprop_sum_random', model_name) result_dir = os.path.join(c.REPO_DIR, 'results', 'ground_truth', 'gaussian_fit_random', model_name) else: backprop_sum_dir = os.path.join(c.REPO_DIR, 'results', 'ground_truth', 'backprop_sum', model_name) result_dir = os.path.join(c.REPO_DIR, 'results', 'ground_truth', 'gaussian_fit', model_name) ############################################################################### # Script guard if __name__ == "__main__": print("Look for a prompt.") user_input = input("This code may take time to run. Are you sure? [y/n] ") if user_input == 'y': pass else: raise KeyboardInterrupt("Interrupted by user") # Get info of conv layers. unit_counter = ConvUnitCounter(model) layer_indices, nums_units = unit_counter.count() _, rf_sizes = get_rf_sizes(model, (227, 227), layer_type=nn.Conv2d) # Helper functions for txt files: def write_txt(f, layer_name, unit_i, raw_params, explained_variance, map_size): # Unpack params amp = raw_params[ParamFormat.A_IDX] mu_x = raw_params[ParamFormat.MU_X_IDX] mu_y = raw_params[ParamFormat.MU_Y_IDX] sigma_1 = raw_params[ParamFormat.SIGMA_1_IDX] sigma_2 = raw_params[ParamFormat.SIGMA_2_IDX] theta = raw_params[ParamFormat.THETA_IDX] offset = raw_params[ParamFormat.OFFSET_IDX] # Some primitive processings: # (1) move original from top-left to map center. mu_x = mu_x - (map_size/2) mu_y = mu_y - (map_size/2) # (2) take the abs value of sigma values. sigma_1 = abs(sigma_1) sigma_2 = abs(sigma_2) # (3) convert theta to orientation. orientation = theta_to_ori(sigma_1, sigma_2, theta) f.write(f"{layer_name} {unit_i} ") f.write(f"{mu_x:.2f} {mu_y:.2f} ") f.write(f"{sigma_1:.2f} {sigma_2:.2f} ") f.write(f"{orientation:.2f} ") f.write(f"{amp:.3f} {offset:.3f} ") f.write(f"{explained_variance:.4f}\n") for sum_mode in sum_modes: backprop_sum_dir_with_mode = os.path.join(backprop_sum_dir, sum_mode) if this_is_a_test_run: result_dir_with_mode = os.path.join(result_dir, 'test') else: result_dir_with_mode = os.path.join(result_dir, sum_mode) # Delete previous files. top_file_path = os.path.join(result_dir_with_mode, f"{model_name}_gt_gaussian_top.txt") bot_file_path = os.path.join(result_dir_with_mode, f"{model_name}_gt_gaussian_bot.txt") if os.path.exists(top_file_path): os.remove(top_file_path) if os.path.exists(bot_file_path): os.remove(bot_file_path) for conv_i in range(len(layer_indices)): layer_name = f"conv{conv_i + 1}" print(f"Fitting elliptical Gaussian for {layer_name}...") # Load backprop sums: max_file_path = os.path.join(backprop_sum_dir_with_mode, f"{layer_name}_max.npy") min_file_path = os.path.join(backprop_sum_dir_with_mode, f"{layer_name}_min.npy") max_maps = np.load(max_file_path) # [unit, y, x] min_maps = np.load(min_file_path) # [unit, y, x] # Initialize arrays for parameters and standard error (SEM) values: num_units = nums_units[conv_i] # For param_cleaner to check if Gaussian is inside in RF or not. rf_size = rf_sizes[conv_i][0] box = (0, 0, rf_size, rf_size) pdf_path = os.path.join(result_dir_with_mode, f"{layer_name}.pdf") with PdfPages(pdf_path) as pdf: for unit_i, (max_map, min_map) in enumerate(tqdm(zip(max_maps, min_maps))): # Do only the first 5 unit during testing phase if this_is_a_test_run and unit_i >= 5: break # Fit 2D Gaussian, and plot them. plt.figure(figsize=(20, 10)) plt.suptitle(f"Elliptical Gaussian fit ({layer_name} no.{unit_i}, " f"sum mode: {sum_mode})", fontsize=20) plt.subplot(1, 2, 1) params, sems = gaussian_fit(max_map, plot=True, show=False) fxvar = calc_f_explained_var(max_map, params) with open(top_file_path, 'a') as top_f: write_txt(top_f, layer_name, unit_i, params, fxvar, rf_size) plt.title(f"max (fxvar = {fxvar:.4f})\n" f"A={params[ParamFormat.A_IDX]:.2f}(err={sems[ParamFormat.A_IDX]:.2f}), " f"mu_x={params[ParamFormat.MU_X_IDX]:.2f}(err={sems[ParamFormat.MU_X_IDX]:.2f}), " f"mu_y={params[ParamFormat.MU_Y_IDX]:.2f}(err={sems[ParamFormat.MU_Y_IDX]:.2f}),\n" f"sigma_1={params[ParamFormat.SIGMA_1_IDX]:.2f}(err={sems[ParamFormat.SIGMA_1_IDX]:.2f}), " f"sigma_2={params[ParamFormat.SIGMA_2_IDX]:.2f}(err={sems[ParamFormat.SIGMA_2_IDX]:.2f}),\n" f"theta={params[ParamFormat.THETA_IDX]:.2f}(err={sems[ParamFormat.THETA_IDX]:.2f}), " f"offset={params[ParamFormat.OFFSET_IDX]:.2f}(err={sems[ParamFormat.OFFSET_IDX]:.2f})", fontsize=14) plt.subplot(1, 2, 2) params, sems = gaussian_fit(min_map, plot=True, show=False) fxvar = calc_f_explained_var(min_map, params) with open(bot_file_path, 'a') as bot_f: write_txt(bot_f, layer_name, unit_i, params, fxvar, rf_size) plt.title(f"min (fxvar = {fxvar:.4f})\n" f"A={params[ParamFormat.A_IDX]:.2f}(err={sems[ParamFormat.A_IDX]:.2f}), " f"mu_x={params[ParamFormat.MU_X_IDX]:.2f}(err={sems[ParamFormat.MU_X_IDX]:.2f}), " f"mu_y={params[ParamFormat.MU_Y_IDX]:.2f}(err={sems[ParamFormat.MU_Y_IDX]:.2f}),\n" f"sigma_1={params[ParamFormat.SIGMA_1_IDX]:.2f}(err={sems[ParamFormat.SIGMA_1_IDX]:.2f}), " f"sigma_2={params[ParamFormat.SIGMA_2_IDX]:.2f}(err={sems[ParamFormat.SIGMA_2_IDX]:.2f}),\n" f"theta={params[ParamFormat.THETA_IDX]:.2f}(err={sems[ParamFormat.THETA_IDX]:.2f}), " f"offset={params[ParamFormat.OFFSET_IDX]:.2f}(err={sems[ParamFormat.OFFSET_IDX]:.2f})", fontsize=14) pdf.savefig() plt.close()