#####################################################
# 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