Skip to content

Instantly share code, notes, and snippets.

@AndBondStyle

AndBondStyle/plot_gradient_flow.py

Created April 15, 2021 14:32
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save AndBondStyle/4cf03e25aa1cac73ddaccb32e117e12d to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save AndBondStyle/4cf03e25aa1cac73ddaccb32e117e12d to your computer and use it in GitHub Desktop.
Download ZIP

Revisions

  1. AndBondStyle created this gist Apr 15, 2021.
    44 changes: 44 additions & 0 deletions plot_gradient_flow.py
    View file
    Original file line number Diff line number Diff line change
    @@ -0,0 +1,44 @@
    import matplotlib.pyplot as plt
    from matplotlib.lines import Line2D
    import numpy as np


    def plot_grad_flow(model):
    ave_grads = []
    max_grads = []
    layers = []
    for name, param in model.named_parameters():
    if param.requires_grad and "bias" not in name:
    if param.grad is None:
    print("[!] No grad:", name)
    continue
    layers.append(name)
    ave_grads.append(param.grad.abs().mean())
    max_grads.append(param.grad.abs().max())

    plt.bar(np.arange(len(max_grads)), max_grads, alpha=0.1, lw=1, color="c")
    plt.bar(np.arange(len(max_grads)), ave_grads, alpha=0.1, lw=1, color="b")
    plt.hlines(0, 0, len(ave_grads) + 1, lw=2, color="k")
    plt.xticks(range(0, len(ave_grads), 1), layers, rotation="vertical")
    plt.xlim(left=0, right=len(ave_grads))
    plt.ylim(bottom=-0.001, top=0.02)
    plt.xlabel("Layers")
    plt.ylabel("average gradient")
    plt.title("Gradient flow")
    plt.grid(True)
    plt.legend(
    [Line2D([0], [0], color=x, lw=4) for x in "cbk"],
    ["max-gradient", "mean-gradient", "zero-gradient"],
    )


    if __name__ == "__main__":
    model = ...
    inputs = ...
    labels = ...
    criterion = ...

    outputs = model.forward(inputs)
    loss = criterion(outputs, labels)
    loss.backward()
    plot_grad_flow(model)