##################################################### # Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2020.03 # ##################################################### # Reformulate the codes in https://github.com/CalculatedContent/WeightWatcher ##################################################### import numpy as np from typing import List import torch.nn as nn from collections import OrderedDict from sklearn.decomposition import TruncatedSVD def available_module_types(): return (nn.Conv2d, nn.Linear) def get_conv2D_Wmats(tensor: np.ndarray) -> List[np.ndarray]: """ Extract W slices from a 4 index conv2D tensor of shape: (N,M,i,j) or (M,N,i,j). Return ij (N x M) matrices """ mats = [] N, M, imax, jmax = tensor.shape assert N + M >= imax + jmax, 'invalid tensor shape detected: {}x{} (NxM), {}x{} (i,j)'.format(N, M, imax, jmax) for i in range(imax): for j in range(jmax): w = tensor[:, :, i, j] if N < M: w = w.T mats.append(w) return mats def glorot_norm_check(W, N, M, rf_size, lower=0.5, upper=1.5): """Check if this layer needs Glorot Normalization Fix""" kappa = np.sqrt(2 / ((N + M) * rf_size)) norm = np.linalg.norm(W) check1 = norm / np.sqrt(N * M) check2 = norm / (kappa * np.sqrt(N * M)) if (rf_size > 1) and (check2 > lower) and (check2 < upper): return check2, True elif (check1 > lower) & (check1 < upper): return check1, True else: if rf_size > 1: return check2, False else: return check1, False def glorot_norm_fix(w, n, m, rf_size): """Apply Glorot Normalization Fix.""" kappa = np.sqrt(2 / ((n + m) * rf_size)) w = w / kappa return w def analyze_weights(weights, min_size, max_size, alphas, lognorms, spectralnorms, softranks, normalize, glorot_fix): results = OrderedDict() count = len(weights) if count == 0: return results for i, weight in enumerate(weights): M, N = np.min(weight.shape), np.max(weight.shape) Q = N / M results[i] = cur_res = OrderedDict(N=N, M=M, Q=Q) check, checkTF = glorot_norm_check(weight, N, M, count) cur_res['check'] = check cur_res['checkTF'] = checkTF # assume receptive field size is count if glorot_fix: weight = glorot_norm_fix(weight, N, M, count) else: # probably never needed since we always fix for glorot weight = weight * np.sqrt(count / 2.0) if spectralnorms: # spectralnorm is the max eigenvalues svd = TruncatedSVD(n_components=1, n_iter=7, random_state=10) svd.fit(weight) sv = svd.singular_values_ sv_max = np.max(sv) if normalize: evals = sv * sv / N else: evals = sv * sv lambda0 = evals[0] cur_res["spectralnorm"] = lambda0 cur_res["logspectralnorm"] = np.log10(lambda0) else: lambda0 = None if M < min_size: summary = "Weight matrix {}/{} ({},{}): Skipping: too small (<{})".format(i + 1, count, M, N, min_size) cur_res["summary"] = summary continue elif max_size > 0 and M > max_size: summary = "Weight matrix {}/{} ({},{}): Skipping: too big (testing) (>{})".format(i + 1, count, M, N, max_size) cur_res["summary"] = summary continue else: summary = [] if alphas: import powerlaw svd = TruncatedSVD(n_components=M - 1, n_iter=7, random_state=10) svd.fit(weight.astype(float)) sv = svd.singular_values_ if normalize: evals = sv * sv / N else: evals = sv * sv lambda_max = np.max(evals) fit = powerlaw.Fit(evals, xmax=lambda_max, verbose=False) alpha = fit.alpha cur_res["alpha"] = alpha D = fit.D cur_res["D"] = D cur_res["lambda_min"] = np.min(evals) cur_res["lambda_max"] = lambda_max alpha_weighted = alpha * np.log10(lambda_max) cur_res["alpha_weighted"] = alpha_weighted tolerance = lambda_max * M * np.finfo(np.max(sv)).eps cur_res["rank_loss"] = np.count_nonzero(sv > tolerance, axis=-1) logpnorm = np.log10(np.sum([ev ** alpha for ev in evals])) cur_res["logpnorm"] = logpnorm summary.append( "Weight matrix {}/{} ({},{}): Alpha: {}, Alpha Weighted: {}, D: {}, pNorm {}".format(i + 1, count, M, N, alpha, alpha_weighted, D, logpnorm)) if lognorms: norm = np.linalg.norm(weight) # Frobenius Norm cur_res["norm"] = norm lognorm = np.log10(norm) cur_res["lognorm"] = lognorm X = np.dot(weight.T, weight) if normalize: X = X / N normX = np.linalg.norm(X) # Frobenius Norm cur_res["normX"] = normX lognormX = np.log10(normX) cur_res["lognormX"] = lognormX summary.append( "Weight matrix {}/{} ({},{}): LogNorm: {} ; LogNormX: {}".format(i + 1, count, M, N, lognorm, lognormX)) if softranks: softrank = norm ** 2 / sv_max ** 2 softranklog = np.log10(softrank) softranklogratio = lognorm / np.log10(sv_max) cur_res["softrank"] = softrank cur_res["softranklog"] = softranklog cur_res["softranklogratio"] = softranklogratio summary += "{}. Softrank: {}. Softrank log: {}. Softrank log ratio: {}".format(summary, softrank, softranklog, softranklogratio) cur_res["summary"] = "\n".join(summary) return results def compute_details(results): """ Return a pandas data frame. """ final_summary = OrderedDict() metrics = { # key in "results" : pretty print name "check": "Check", "checkTF": "CheckTF", "norm": "Norm", "lognorm": "LogNorm", "normX": "Norm X", "lognormX": "LogNorm X", "alpha": "Alpha", "alpha_weighted": "Alpha Weighted", "spectralnorm": "Spectral Norm", "logspectralnorm": "Log Spectral Norm", "softrank": "Softrank", "softranklog": "Softrank Log", "softranklogratio": "Softrank Log Ratio", "sigma_mp": "Marchenko-Pastur (MP) fit sigma", "numofSpikes": "Number of spikes per MP fit", "ratio_numofSpikes": "aka, percent_mass, Number of spikes / total number of evals", "softrank_mp": "Softrank for MP fit", "logpnorm": "alpha pNorm" } metrics_stats = [] for metric in metrics: metrics_stats.append("{}_min".format(metric)) metrics_stats.append("{}_max".format(metric)) metrics_stats.append("{}_avg".format(metric)) metrics_stats.append("{}_compound_min".format(metric)) metrics_stats.append("{}_compound_max".format(metric)) metrics_stats.append("{}_compound_avg".format(metric)) columns = ["layer_id", "layer_type", "N", "M", "layer_count", "slice", "slice_count", "level", "comment"] + [*metrics] + metrics_stats metrics_values = {} metrics_values_compound = {} for metric in metrics: metrics_values[metric] = [] metrics_values_compound[metric] = [] layer_count = 0 for layer_id, result in results.items(): layer_count += 1 layer_type = np.NAN if "layer_type" in result: layer_type = str(result["layer_type"]).replace("LAYER_TYPE.", "") compounds = {} # temp var for metric in metrics: compounds[metric] = [] slice_count, Ntotal, Mtotal = 0, 0, 0 for slice_id, summary in result.items(): if not str(slice_id).isdigit(): continue slice_count += 1 N = np.NAN if "N" in summary: N = summary["N"] Ntotal += N M = np.NAN if "M" in summary: M = summary["M"] Mtotal += M data = {"layer_id": layer_id, "layer_type": layer_type, "N": N, "M": M, "slice": slice_id, "level": "SLICE", "comment": "Slice level"} for metric in metrics: if metric in summary: value = summary[metric] if value is not None: metrics_values[metric].append(value) compounds[metric].append(value) data[metric] = value data = {"layer_id": layer_id, "layer_type": layer_type, "N": Ntotal, "M": Mtotal, "slice_count": slice_count, "level": "LAYER", "comment": "Layer level"} # Compute the compound value over the slices for metric, value in compounds.items(): count = len(value) if count == 0: continue compound = np.mean(value) metrics_values_compound[metric].append(compound) data[metric] = compound data = {"layer_count": layer_count, "level": "NETWORK", "comment": "Network Level"} for metric, metric_name in metrics.items(): if metric not in metrics_values or len(metrics_values[metric]) == 0: continue values = metrics_values[metric] minimum = min(values) maximum = max(values) avg = np.mean(values) final_summary[metric] = avg # print("{}: min: {}, max: {}, avg: {}".format(metric_name, minimum, maximum, avg)) data["{}_min".format(metric)] = minimum data["{}_max".format(metric)] = maximum data["{}_avg".format(metric)] = avg values = metrics_values_compound[metric] minimum = min(values) maximum = max(values) avg = np.mean(values) final_summary["{}_compound".format(metric)] = avg # print("{} compound: min: {}, max: {}, avg: {}".format(metric_name, minimum, maximum, avg)) data["{}_compound_min".format(metric)] = minimum data["{}_compound_max".format(metric)] = maximum data["{}_compound_avg".format(metric)] = avg return final_summary def analyze(model: nn.Module, min_size=50, max_size=0, alphas: bool = False, lognorms: bool = True, spectralnorms: bool = False, softranks: bool = False, normalize: bool = False, glorot_fix: bool = False): """ Analyze the weight matrices of a model. :param model: A PyTorch model :param min_size: The minimum weight matrix size to analyze. :param max_size: The maximum weight matrix size to analyze (0 = no limit). :param alphas: Compute the power laws (alpha) of the weight matrices. Time consuming so disabled by default (use lognorm if you want speed) :param lognorms: Compute the log norms of the weight matrices. :param spectralnorms: Compute the spectral norm (max eigenvalue) of the weight matrices. :param softranks: Compute the soft norm (i.e. StableRank) of the weight matrices. :param normalize: Normalize or not. :param glorot_fix: :return: (a dict of all layers' results, a dict of the summarized info) """ names, modules = [], [] for name, module in model.named_modules(): if isinstance(module, available_module_types()): names.append(name) modules.append(module) # print('There are {:} layers to be analyzed in this model.'.format(len(modules))) all_results = OrderedDict() for index, module in enumerate(modules): if isinstance(module, nn.Linear): weights = [module.weight.cpu().detach().numpy()] else: weights = get_conv2D_Wmats(module.weight.cpu().detach().numpy()) results = analyze_weights(weights, min_size, max_size, alphas, lognorms, spectralnorms, softranks, normalize, glorot_fix) results['id'] = index results['type'] = type(module) all_results[index] = results summary = compute_details(all_results) return all_results, summary