Cockpit is a visual and statistical debugger specifically designed for deep learning.

No support yet for custom optimizers?

Error(log below) while trying cockpit with custom optimizer class(that subclasses torch.optim.Optimizer). Although seems to work as expected with in-built optimizers (torch.optim.SGD).

Log

Traceback (most recent call last): File "laplace_primer.py", line 137, in sgld.train(prefinal_train, train_loader, model) File "../SGLD_laplace.py", line 123, in train loss.backward(create_graph=cockpit.create_graph(epochs)) File "./miniconda3/envs/backpack/lib/python3.7/site-packages/cockpit/context.py", line 137, in exit self.cp.track(self.global_step, protected_savefields=self.protected_savefields) File "../miniconda3/envs/backpack/lib/python3.7/site-packages/cockpit/cockpit.py", line 178, in track q.track(global_step, self.params, batch_loss) File "../miniconda3/envs/backpack/lib/python3.7/site-packages/cockpit/quantities/quantity.py", line 87, in track iteration, result = self.compute(global_step, params, batch_loss) File "./miniconda3/envs/backpack/lib/python3.7/site-packages/cockpit/quantities/quantity.py", line 516, in compute save_result = self._compute(global_step, params, batch_loss) File "../miniconda3/envs/backpack/lib/python3.7/site-packages/cockpit/quantities/quantity.py", line 538, in _compute self._compute_start(global_step, params, batch_loss) File "../miniconda3/envs/backpack/lib/python3.7/site-packages/cockpit/quantities/alpha.py", line 280, in _compute_start self._save_1st_order_info(global_step, params, batch_loss, point, until) File "../miniconda3/envs/backpack/lib/python3.7/site-packages/cockpit/quantities/alpha.py", line 332, in _save_1st_order_info id(p): batch_size * p.grad_batch.data.clone().detach() for p in params File "../miniconda3/envs/backpack/lib/python3.7/site-packages/cockpit/quantities/alpha.py", line 332, in id(p): batch_size * p.grad_batch.data.clone().detach() for p in params AttributeError: 'Parameter' object has no attribute 'grad_batch'

Description

Plotting to a file crashes (see stacktrace below). Calling plot with show_plot=True, save_plot=True works and results in the file being saved correctly. The issue is reproducible over multiple runs and thus probably not caused by weird gradients etc.

Steps to Reproduce

1. Setup

config = [
                quantities.Distance(schedules.linear(interval=1)),
                quantities.GradHist1d(schedules.linear(interval=1)),
                quantities.GradNorm(schedules.linear(interval=1)),
                quantities.InnerTest(schedules.linear(interval=1)),
                quantities.Loss(schedules.linear(interval=1)),
                quantities.NormTest(schedules.linear(interval=1)),
                quantities.OrthoTest(schedules.linear(interval=1)),
                quantities.UpdateSize(schedules.linear(interval=1)),
                quantities.GradHist2d(schedules.linear(interval=1))
            ]
            self._cockpit = Cockpit(model.parameters(), quantities=config)
            self._cockpit_plotter = CockpitPlotter()

2. Track a few backward steps

3. Call plot with show_plot=False and save_plot=True

self._cockpit_plotter.plot(self._cockpit,
                                        savedir=cockpit_log_dir,
                                        show_plot=False,
                                        save_plot=True)

Traceback (most recent call last):
  File "code/train.py", line 215, in <module>
    metrics = main(opts)
  File "code/train.py", line 185, in main
    training_info = trainer.train(model, dataset, eval_data, logger)
  File "./code/trainers/mse_trainer.py", line 186, in train
    self.log(logger.root_dir)
  File "./code/trainers/base_trainer.py", line 145, in log
    save_plot=True)
  File "./envs/pytorch/lib/python3.7/site-packages/cockpit/plotter.py", line 142, in plot
    self._plot_gradients(self.grid_spec[0, 1])
  File "./envs/pytorch/lib/python3.7/site-packages/cockpit/plotter.py", line 223, in _plot_gradients
    instruments.gradient_tests_gauge(self, self.fig, self.gs_gradients[1, 1])
  File "./envs/pytorch/lib/python3.7/site-packages/cockpit/instruments/gradient_tests_gauge.py", line 69, in gradient_tests_gauge
    _format(self, ax_all, ax_norm, ax_inner, ax_ortho)
  File "./envs/pytorch/lib/python3.7/site-packages/cockpit/instruments/gradient_tests_gauge.py", line 91, in _format
    ax_norm.set_yscale("log")
  File "./envs/pytorch/lib/python3.7/site-packages/matplotlib/axes/_base.py", line 3708, in set_yscale
    ax.yaxis._set_scale(value, **kwargs)
  File "./envs/pytorch/lib/python3.7/site-packages/matplotlib/axis.py", line 800, in _set_scale
    self._scale.set_default_locators_and_formatters(self)
  File "./envs/pytorch/lib/python3.7/site-packages/matplotlib/scale.py", line 406, in set_default_locators_and_formatters
    axis.set_major_locator(LogLocator(self.base))
  File "./envs/pytorch/lib/python3.7/site-packages/matplotlib/axis.py", line 1651, in set_major_locator
    self.stale = True
  File "./envs/pytorch/lib/python3.7/site-packages/matplotlib/artist.py", line 230, in stale
    self.stale_callback(self, val)
  File "./envs/pytorch/lib/python3.7/site-packages/matplotlib/artist.py", line 51, in _stale_axes_callback
    self.axes.stale = val
  File "./envs/pytorch/lib/python3.7/site-packages/matplotlib/artist.py", line 230, in stale
    self.stale_callback(self, val)
  File "./envs/pytorch/lib/python3.7/site-packages/matplotlib/figure.py", line 51, in _stale_figure_callback
    self.figure.stale = val
  File "./envs/pytorch/lib/python3.7/site-packages/matplotlib/artist.py", line 230, in stale
    self.stale_callback(self, val)
  File "./envs/pytorch/lib/python3.7/site-packages/matplotlib/pyplot.py", line 589, in _auto_draw_if_interactive
    fig.canvas.draw_idle()
  File "./envs/pytorch/lib/python3.7/site-packages/matplotlib/backend_bases.py", line 1907, in draw_idle
    self.draw(*args, **kwargs)
  File "./envs/pytorch/lib/python3.7/site-packages/matplotlib/backends/backend_agg.py", line 388, in draw
    self.figure.draw(self.renderer)
  File "./envs/pytorch/lib/python3.7/site-packages/matplotlib/artist.py", line 38, in draw_wrapper
    return draw(artist, renderer, *args, **kwargs)
  File "./envs/pytorch/lib/python3.7/site-packages/matplotlib/figure.py", line 1709, in draw
    renderer, self, artists, self.suppressComposite)
  File "./envs/pytorch/lib/python3.7/site-packages/matplotlib/image.py", line 135, in _draw_list_compositing_images
    a.draw(renderer)
  File "./envs/pytorch/lib/python3.7/site-packages/matplotlib/artist.py", line 38, in draw_wrapper
    return draw(artist, renderer, *args, **kwargs)
  File "./envs/pytorch/lib/python3.7/site-packages/matplotlib/axes/_base.py", line 2647, in draw
    mimage._draw_list_compositing_images(renderer, self, artists)
  File "./envs/pytorch/lib/python3.7/site-packages/matplotlib/image.py", line 135, in _draw_list_compositing_images
    a.draw(renderer)
  File "./envs/pytorch/lib/python3.7/site-packages/matplotlib/artist.py", line 38, in draw_wrapper
    return draw(artist, renderer, *args, **kwargs)
  File "./envs/pytorch/lib/python3.7/site-packages/matplotlib/axis.py", line 1203, in draw
    ticks_to_draw = self._update_ticks()
  File "./envs/pytorch/lib/python3.7/site-packages/matplotlib/axis.py", line 1079, in _update_ticks
    major_locs = self.get_majorticklocs()
  File "./envs/pytorch/lib/python3.7/site-packages/matplotlib/axis.py", line 1324, in get_majorticklocs
    return self.major.locator()
  File "./envs/pytorch/lib/python3.7/site-packages/matplotlib/ticker.py", line 2271, in __call__
    return self.tick_values(vmin, vmax)
  File "./envs/pytorch/lib/python3.7/site-packages/matplotlib/ticker.py", line 2297, in tick_values
    "Data has no positive values, and therefore can not be "
ValueError: Data has no positive values, and therefore can not be log-scaled.

Cockpit crashes with non descriptive error.

Description

I am trying to run cockpit on a simple network trained with MSE loss (although implemented through custom modules, not nn.Sequential). As far as I can see, no unsupported operations are involved (Linear, Sequential, Tanh, Identity).

Cockpit crashes on computing the BatchGrad extension. Using backpack to compute the batch gradient works without crashes.

Steps to Reproduce

1. Setup

self.loss_fn = lambda pred, target: ((target - pred) ** 2).mean(dim=-1)

self._cockpit = Cockpit(model.parameters(),
                                    quantities=configuration("economy"))
model = MyCustomModel()
model = extend(model)
optimizer = torch.optim.Adam(params=model.parameters(), lr=1e-3)

2. Training Step

pred = model(inp)
loss = self.loss_fn(pred, target)
mean_loss = loss.mean()

optimizer.zero_grad()
info = {
    "batch_size": len(loss),
    "individual_losses": loss,
    "loss": mean_loss,
    "optimizer": optimizer,
}
with self._cockpit(step, info=info, debug=True):
    create_graph = self._cockpit.create_graph(step)
    mean_loss.backward(create_graph=create_graph)  # CRASH HERE
optimizer.step()

Source or Possible Fix

Stacktrace

[DEBUG, step 0]
 ↪Quantities  : [<cockpit.quantities.alpha.Alpha object at 0x7fa751af5f10>, <cockpit.quantities.distance.Distance object at 0x7fa74e7c0cd0>, <cockpit.quantities.grad_hist.GradHist1d object at 0x7fa74e7c0d10>, <cockpit.quantities.grad_norm.GradNorm object at 0x7fa74e7c0d90>, <cockpit.quantities.inner_test.InnerTest object at 0x7fa74e7c0dd0>, <cockpit.quantities.loss.Loss object at 0x7fa74e7cd050>, <cockpit.quantities.norm_test.NormTest object at 0x7fa74e7cd090>, <cockpit.quantities.ortho_test.OrthoTest object at 0x7fa74e7cd0d0>, <cockpit.quantities.update_size.UpdateSize object at 0x7fa74e7cd110>]
 ↪Extensions  : [<backpack.extensions.firstorder.batch_grad.BatchGrad object at 0x7fa74e7cd750>]
 ↪Hooks       : <cockpit.quantities.utils_transforms.BatchGradTransformsHook object at 0x7fa74e7931d0>
 ↪Create graph: False
 ↪Save memory : True
[DEBUG] Running extension <backpack.extensions.firstorder.batch_grad.BatchGrad object at 0x7fa74e7cd750> on Linear(in_features=128, out_features=1, bias=True)
[DEBUG] Running extension <backpack.extensions.firstorder.batch_grad.BatchGrad object at 0x7fa74e7cd750> on LinearHead(
  (l): Linear(in_features=128, out_features=1, bias=True)
)
[DEBUG] Running extension <backpack.extensions.firstorder.batch_grad.BatchGrad object at 0x7fa74e7cd750> on Identity()
[DEBUG] Running extension <backpack.extensions.firstorder.batch_grad.BatchGrad object at 0x7fa74e7cd750> on Sequential(
  (0): Linear(in_features=1, out_features=128, bias=True)
  (1): Tanh()
  (2): Linear(in_features=128, out_features=128, bias=True)
  (3): Tanh()
  (4): LinearHead(
    (l): Linear(in_features=128, out_features=1, bias=True)
  )
  (5): Identity()
)
[DEBUG] Running extension <backpack.extensions.firstorder.batch_grad.BatchGrad object at 0x7fa74e7cd750> on MLP(
  (layers): Sequential(
    (0): Linear(in_features=1, out_features=128, bias=True)
    (1): Tanh()
    (2): Linear(in_features=128, out_features=128, bias=True)
    (3): Tanh()
    (4): LinearHead(
      (l): Linear(in_features=128, out_features=1, bias=True)
    )
    (5): Identity()
  )
)
Traceback (most recent call last):
  File "./envs/pytorch/lib/python3.7/site-packages/backpack/__init__.py", line 169, in run_extension_hook
    CTX.get_extension_hook()(module)
  File "./envs/pytorch/lib/python3.7/site-packages/cockpit/quantities/hooks/base.py", line 53, in __call__
    self.run_hook(param, module)
  File "./envs/pytorch/lib/python3.7/site-packages/cockpit/quantities/hooks/base.py", line 80, in run_hook
    value = self.module_hook(param, module)
  File "./envs/pytorch/lib/python3.7/site-packages/cockpit/quantities/hooks/base.py", line 139, in module_hook
    return self.param_hook(param)
  File "./envs/pytorch/lib/python3.7/site-packages/cockpit/quantities/utils_transforms.py", line 78, in param_hook
    param.grad_batch._param_weakref = weakref.ref(param)
AttributeError: 'Parameter' object has no attribute 'grad_batch'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "./code/trainers/base_trainer.py", line 112, in backward
    mean_loss.backward(create_graph=create_graph)
  File "./envs/pytorch/lib/python3.7/site-packages/torch/tensor.py", line 221, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph)
  File "./envs/pytorch/lib/python3.7/site-packages/torch/autograd/__init__.py", line 132, in backward
    allow_unreachable=True)  # allow_unreachable flag
  File "./envs/pytorch/lib/python3.7/site-packages/backpack/__init__.py", line 151, in hook_run_extensions
    run_extension_hook(module)
  File "./envs/pytorch/lib/python3.7/site-packages/backpack/__init__.py", line 172, in run_extension_hook
    raise RuntimeError(f"Post extensions hook failed: {message}")
RuntimeError: Post extensions hook failed: AttributeError("'Parameter' object has no attribute 'grad_batch'")

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "code/train.py", line 215, in <module>
    metrics = main(opts)
  File "code/train.py", line 185, in main
    training_info = trainer.train(model, dataset, eval_data, logger)
  File "./code/trainers/mse_trainer.py", line 128, in train
    self.backward(global_step, loss, mean_loss, optimizer)
  File "./code/trainers/base_trainer.py", line 112, in backward
    mean_loss.backward(create_graph=create_graph)
  File "./envs/pytorch/lib/python3.7/site-packages/cockpit/context.py", line 137, in __exit__
    self.cp.track(self.global_step, protected_savefields=self.protected_savefields)
  File "./envs/pytorch/lib/python3.7/site-packages/cockpit/cockpit.py", line 178, in track
    q.track(global_step, self.params, batch_loss)
  File "./envs/pytorch/lib/python3.7/site-packages/cockpit/quantities/quantity.py", line 87, in track
    iteration, result = self.compute(global_step, params, batch_loss)
  File "./envs/pytorch/lib/python3.7/site-packages/cockpit/quantities/quantity.py", line 516, in compute
    save_result = self._compute(global_step, params, batch_loss)
  File "./envs/pytorch/lib/python3.7/site-packages/cockpit/quantities/quantity.py", line 538, in _compute
    self._compute_start(global_step, params, batch_loss)
  File "./envs/pytorch/lib/python3.7/site-packages/cockpit/quantities/alpha.py", line 280, in _compute_start
    self._save_1st_order_info(global_step, params, batch_loss, point, until)
  File "./envs/pytorch/lib/python3.7/site-packages/cockpit/quantities/alpha.py", line 326, in _save_1st_order_info
    grad_dict = {id(p): p.grad.data.clone().detach() for p in params}
  File "./envs/pytorch/lib/python3.7/site-packages/cockpit/quantities/alpha.py", line 326, in <dictcomp>
    grad_dict = {id(p): p.grad.data.clone().detach() for p in params}
AttributeError: 'NoneType' object has no attribute 'data'

I have a question regarding the scaling of the GSNR computation as provided in Cockpit

Description

In the following gist, I've written 3 different ways of computing the GSNR as defined in this paper, the same one cited in the Cockpit paper:

1. Extracting the variance directly from the sample-wise gradients, via BackPACK BatchGrad() extension (bad, memory inefficient)
2. Using the Variance() BackPACK extension (efficient & straightforward)
3. Using a replacement formula for the variance to avoid sample-wise grad squaring: Var(X) = E[X²] - E[X]² (likely faster, less straightforward)

https://gist.github.com/andres-fr/db9d0ba31d1502df62a09d382e504d1e

Specifically, the first and second reflect the mathematical definition from the paper:

GSNR_paper = grad_squared / Var(gradients)

And this is the implementation of the replacement formula that I found to be equivalent:

GSNR_gist = grad_squared / ( (sum_grad_squared / batch_size) - (grad_squared / (batch_size**2)) + epsilon)

My issue/question is when I compare the last expression to the one in the Cockpit source code, a clear difference can be noticed in the batch scaling:

GSNR_cockpit = grad_squared / (batch_size * sum_grad_squared - grad_squared + epsilon)

When I used this last expression in my gist, it led to results clearly different from the other approaches. Am I missing something here? Is cockpit doing some further scaling under the hood? Or are these results scaled differently from the definition for some reason?

Steps to Reproduce

The gist works on CPU and CUDA. For easier reproduction, the gist depends on PyTorch and BackPACK only. The script docstring includes steps to install a minimal working conda environment. Then, simply run the script.

Source or Possible Fix

The BackPACK docs already warned about scaling issues with the Variance, which may even involve the loss function, so I wonder if this is a bug/feature.

Intuitively, I'd say it is good if the quantity matches the definition from the paper, but again, implicit scaling may be getting in the way?

I am running some experiments on a remote machine, and I can't see any cockpit plot. I have the same problem when running the basic example. Is there a way of solving this? Am I missing something?

This may be a dumb question, sorry I am quite new to this.

Thank you!

Hi!

Thank you for this project, indeed it is quite interesting work.

I am curious to know that if we can use it to interpret/observe the training of a 3D resnet model?

Regards

Description

What should I do if my model returns the tuple of loss and individual losses and I have no explicit loss function class?

At the moment if I just extend the model, I get:

cockpit/quantities/alpha.py in <dictcomp>(.0)
    324         self.save_to_cache(global_step, f"params_{point}", params_dict, block_fn)
    325 
--> 326         grad_dict = {id(p): p.grad.data.clone().detach() for p in params}
    327         self.save_to_cache(global_step, f"grad_{point}", grad_dict, block_fn)
    328 

AttributeError: 'NoneType' object has no attribute 'data'

Thanks for any suggestions!

Description

I suggest to extend CockpitPlotter.plot such that a keyword dictionary can be passed to it specifying parameters to fig.savefig. This is to be able to choose 1) a DPI value or 2) a different file format (e.g. PDF).

Currently, the stored image files have a poor resolution and thus are not very readable. Being able to choose a DPI value or saving to PDF would fix this issue.

Description

Plotter has a savedir parameter for which the documentation says it specifies the directory where to save the plot to. However, the _save function interprets this parameter as a file path: https://github.com/f-dangel/cockpit/blob/937b3eac8d9fef7e6e6fbc9c91863dd08b1362b1/cockpit/plotter.py#L400-L405

This means plotting with savedir='/example/path' would result in a file /example/path__primary.png, where /example/path/__primary.png would be expected.

Besides, when choosing to not display the plot, I don't think __primary or __secondary should be added to the filename. From an API standpoint, it would probably be best to let the user flexibly choose the path and filename to save to (potentially adding the image extension).

Hi!

I am trying to replicate the plots you show, but I'm not able to find Learning Rates and HyperParameters on the json file created. Where can I find them?

Also, I was wondering if I can find the descriptions of the quantities somewhere in the documentation. I found a table with names and descriptions, but for example the dimensions are not explained and for some quantities I am not able to understand why they have such dimensions on the json file.

Thanks a lot!

Bests, Cristian

I have cockpit mostly working except I get this warnings like this because I have not implemented backpack's batch grad for all my modules:

/usr/local/lib/python3.9/site-packages/backpack/extensions/backprop_extension.py:106: UserWarning: Extension saving to grad_batch does not have an extension for Module <class 'pycompress.model_parts.PredHist'> although the module has parameters 

So as a result I am not able to get all the quantities without error. I understand backpack is central to the design of cockpit but I am fine with some things not working when I just want to prototype quickly.

My questions/requests are:

1. I still see the gradient histogram. Is that histogram only including the gradients for modules where the backpack extension worked or is it all of them?
2. Have you thought about integrating the much smaller project https://github.com/alwynmathew/gradflow-check. They provide a chart that lets you check for vanishing gradient and they do it without backpack based on just having the named_parameters
3. Unrelated but btw: Have you thought about including the ESD of the weight matrices as done in? https://github.com/CalculatedContent/WeightWatcher

Thank you for your excellent work on this project!

Requesting a new feature (I can't seem to add a label myself).

Description

Do you currently support Huggingface Transformers? In particular, I'd like to debug T5.

I tried an editable install of Transformers, and modifying the PyTorch implementation underneath, but I can't quite get the two different loss functions to play together properly. For example, I'm trying to steal this implementation from the Basic Examples, and add it to Transformers:

loss_fn = extend(torch.nn.CrossEntropyLoss(reduction="mean"))
individual_loss_fn = torch.nn.CrossEntropyLoss(reduction="none")

After 30 minutes of tinkering, I can't get arround this error. Thanks for any assistance.

AssertionError: BackPACK extension expects a backpropagation quantity but it is None. Module: Linear(in_features=768, out_features=32128, bias=False), Extension: <backpack.extensions.secondorder.diag_hessian.DiagHessian object at 0x7f6a72f102b0>.

During handling of the above exception, another exception occurred:

AttributeError                            Traceback (most recent call last)
Cell In[6], line 80
     69 # loss.sum().backward()
     70 # backward pass
     71 with cockpit(
     72     global_step,
     73     info={
   (...)
     78     },
     79 ):
---> 80     loss.sum().backward(create_graph=cockpit.create_graph(global_step))
     82 # optimizer step
     83 optimizer.step()

File ~/githubs/cockpit/cockpit/context.py:155, in BackwardCTX.__exit__(self, type, value, traceback)
    152 for ctx in self.contexts:
    153     ctx.__exit__(type, value, traceback)
--> 155 self.cp.track(self.global_step, protected_savefields=self.protected_savefields)
    157 CockpitCTX.erase()

File ~/githubs/cockpit/cockpit/cockpit.py:195, in Cockpit.track(self, global_step, protected_savefields)
    190 before_cleanup = [
    191     q for q in self.quantities if not isinstance(q, quantities.HessMaxEV)
    192 ]
    194 for q in before_cleanup:
--> 195     q.track(global_step, self.params, batch_loss)
    197 self._free_backpack_buffers(global_step, protected_savefields)
    199 after_cleanup = [
    200     q for q in self.quantities if isinstance(q, quantities.HessMaxEV)
    201 ]

File ~/githubs/cockpit/cockpit/quantities/quantity.py:101, in Quantity.track(self, global_step, params, batch_loss)
     92 """Perform scheduled computations and store result.
     93 
     94 Args:
   (...)
...
--> 335 grad_dict = {id(p): p.grad.data.clone().detach() for p in params}
    336 self.save_to_cache(global_step, f"grad_{point}", grad_dict, block_fn)
    338 # L = ¹/ₙ ∑ᵢ ℓᵢ, BackPACK's BatchGrad computes ¹/ₙ ∇ℓᵢ, we have to rescale

AttributeError: 'NoneType' object has no attribute 'data'

Update: after 2 more hours of tinkering, I'm pretty sure I am exactly matching your examples, but I'm still facing the exact same issue.

BTW, I think I found this project at NeurIPS. Great work on it!