<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Overview
This notebook calculate 
* How often each unit is on, averages accross the whole datset
* The level of sparsity seen for each output to a layer
* The entropy of each unit $$-\sum p_i \log p_i $$ where $p_i$ is the fraction of times the unit is on throughout the whole dataset.
Load the experiment
This will load the desired experiment config and related Dataloader.
Step1:  The the model from a checkpoint.
Step6:  Helper class and method.
Step7:  Add forward hooks to all layers of the model
Step8:  A Brief Overview of the Model
Step9:  Compute statistics of the output representations
How often each unit is on
The level of sparsity for each output from a layer
Step11:  Print summary of statitics
How often each unit is on
The level of sparsity for each output to a layer
The entropy of each unit (excluding dead units; otherwise we take the lof of zero)
Step12:  Visualize the distribution of how often each unit is on or off
<ASSISTANT_TASK:>
Python Code:
BATCH_SIZE = 16
PROJECT_DIR = os.path.expanduser("~/nta/nupic.research/projects/gsc")
EXPERIMENT_NAME = "default_sparse_cnn"
CHECKPOINT_FILE = "/home/ec2-user/nta/results/experiments/gsc/default_sparse_cnn/RemoteProcessTrainable_0_2021-01-07_15-54-473wnu2slv/checkpoint_30/checkpoint"
USE_CUDA = torch.cuda.is_available()
# Load the experiment
import os
import sys
sys.path.insert(0, PROJECT_DIR)
from experiments import CONFIGS
config = CONFIGS[EXPERIMENT_NAME]
exp_cls = config["experiment_class"]
exp = exp_cls()
# Load dataloader with specified BATCH_SIZE.
dataset = exp_cls.load_dataset(config, train=True)
train_loader = DataLoader(
    dataset=dataset,
    batch_size=BATCH_SIZE,
    shuffle=True,
    num_workers=0,
    sampler=None,
    pin_memory=USE_CUDA,
    drop_last=config.get("train_loader_drop_last", False),
)
Explanation: Overview
This notebook calculate 
* How often each unit is on, averages accross the whole datset
* The level of sparsity seen for each output to a layer
* The entropy of each unit $$-\sum p_i \log p_i $$ where $p_i$ is the fraction of times the unit is on throughout the whole dataset.
Load the experiment
This will load the desired experiment config and related Dataloader.
End of explanation
# Path to checkpoint.
assert CHECKPOINT_FILE is not None, "You should specify a path to a checkpoint."
# Load model
config = CONFIGS[EXPERIMENT_NAME]
exp_cls = config["experiment_class"]
exp = exp_cls()
model = exp_cls.create_model(config, "cuda")
load_state_from_checkpoint(model=model, checkpoint_path=CHECKPOINT_FILE, device="cuda");
Explanation: The the model from a checkpoint.
End of explanation
class SaveForwardOutput:
    
    Forward hook for atorch.nn.Module. This will store the output from a module during the forward pass.
    :param name: name of the module (e.g. "classifier")
    :param verbose: whether to print something every forward pass
    
    def __init__(self, name=None, verbose=False):
        self.name = name
        self.verbose = verbose
        self.output = None
    
    def __call__(self, module, x, y):
        
        Forward hook on torch.nn.Module
        :param module: module
        :param x: tuple of inputs
        :param y: ouput of module
        
        if self.verbose:
            print(f"Forward hook of {self.name} ({module.__class__.__name__})")
        # Save the output of the layer.
        self.output = y
def calculate_on_counts(rep):
    
    Calculates the number of times a unit is on for the given batch.
    :param rep: tensor of an intermediate representation
    :return: a 1D tensor of with len equal to the total number of units
    
    rep = rep.flatten(start_dim=1)
    on_mask = rep != 0
    on_counts = on_mask.sum(axis=0)
    return on_counts
def calculate_sparsities(rep):
    
    Calculates the fraction of units off for each sample in the batch.
    :param rep: tensor of an intermediate representation
    :return: tensor of sparsities, one for each sample in the batch
    
    rep = rep.flatten(start_dim=1)
    off_mask = rep == 0
    sparsities = off_mask.sum(axis=1).float()
    sparsities = sparsities / rep.shape[1]
    return sparsities
Explanation: Helper class and method.
End of explanation
forward_hooks = {}
tracked_modules = []
for n, m in model.named_modules():
    # Skip over modules that have sub-modules.
    # These sub-modules will be found recursively by 'named_modules'. 
    if len(list(m.children())) > 0:
        continue
    # TODO: We can filter which modules get tracked if needed.
    forward_hook = SaveForwardOutput(name=n)
    m.register_forward_hook(forward_hook)
    forward_hooks[n] = forward_hook
    tracked_modules.append(m)
Explanation: Add forward hooks to all layers of the model
End of explanation
total_params, on_params = count_nonzero_params(model)
print(f"Model density (on_params / total_params): {on_params:,} / {total_params:,} = {on_params/total_params:.2f}")
print(model)
Explanation: A Brief Overview of the Model
End of explanation
# This will be a tensor of length equal to the total number of output units to the given layer.
# Each value will represent the frequency ("freq") of how often that unit is on throughout the whol dataset.
layer_on_freqs = defaultdict(int)
# This will be a list of sparsities measure for each output passed to the layer.
layer_sparsities = defaultdict(list)
# Loop through the whole dataset. See "Load the experiment" section
# above to use the validation dataloader if desired.
total_images = 0
model.eval()
for x, y in train_loader:
    x, y = x.to("cuda"), y.to("cuda")
    out = model(x)
    for name, hook in forward_hooks.items():
        # Retrieve ouput saved in the forward hook.
        representation = hook.output
        # Tally the on-counts for each unit.
        on_counts = calculate_on_counts(representation)
        layer_on_freqs[name] += on_counts
        # Compute the level of sparsity for each output.
        sparsities = calculate_sparsities(representation)
        layer_sparsities[name].extend(sparsities.tolist())
    total_images += len(x)
# Divide through by number of samples to find on-frequency for eash unit.
for layer, on_counts in layer_on_freqs.items():
    layer_on_freqs[layer] = on_counts.float() / total_images
Explanation: Compute statistics of the output representations
How often each unit is on
The level of sparsity for each output from a layer
End of explanation
df = DataFrame(columns=["module type", "mu output sparsity", "sd output sparsity", "entropy", "dead neuron count", "total neurons", "fraction_dead"])
def log2_of_zeros(arr):
    
    Intentionally leave zero's as they are when taking the log.
    
    arr = arr.copy()
    arr[arr != 0] = np.log(arr[arr != 0])
    return arr
for (layer, sparsities), freqs, module in zip(layer_sparsities.items(), layer_on_freqs.values(), tracked_modules):
    mu_sparsity = np.mean(sparsities)
    sd_sparsity = np.std(sparsities)
    module_type = module.__class__.__name__
    print(f"Layer {layer} ({module_type})")
    print(f"    Average sparsity of output = {mu_sparsity:.3f} ± {sd_sparsity:.3f}")
    freqs = freqs.cpu().numpy()
    entropy = np.sum(-freqs * log2_of_zeros(freqs) / len(freqs))
    print(f"    Entropy of layer = {entropy}")
    dead_count = (freqs == 0).sum()
    total_neurons = len(freqs)
    fraction_dead = dead_count / total_neurons
    print(f"    Fractions of dead neurons in layer = {fraction_dead:.3f} ({dead_count} out of {total_neurons})")
    print()
    df.at[layer] = (module_type, mu_sparsity, sd_sparsity, entropy, dead_count, total_neurons, fraction_dead)
df
Explanation: Print summary of statitics
How often each unit is on
The level of sparsity for each output to a layer
The entropy of each unit (excluding dead units; otherwise we take the lof of zero)
End of explanation
print("Names of layers observed via hooks: (helpful for specifying 'layer' in the cell below.\n")
pprint(list(layer_on_freqs.keys()))
# Choose one layer to visualize the frequency of it's output units being on or off.
layer = "linear_kwinner"
on_freqs = layer_on_freqs[layer].cpu().numpy()
fig = sns.histplot(data=on_freqs, kde=True, stat="probability")
fig.set_xlabel("on-frequency")
fig.set_title(f"How often are units on for layer '{layer}'");
Explanation: Visualize the distribution of how often each unit is on or off
End of explanation
<END_TASK>