graftr

An interactive shell to view and edit PyTorch checkpoints.

graftr can be used to remove, rename, and move around the layers and parameters of your saved model. It's also a handy tool to peek into the structure of pre-trained PyTorch models that you can find online (e.g. Transformer, DCGAN, etc.).

The screencast above shows an example of taking a pre-trained Densenet and preparing it for integration into a larger model. We remove the final classification layer and move the feature extractor into its own densenet module.

Install

pip install graftr

Documentation

graftr presents a hierarchical directory structure for state_dicts and parameters in your checkpoint. You can list (ls), move/rename (mv), and print (cat) parameters. And, of course, you can navigate (cd) through the hierarchy. It also supports standard shell beahvior like command history, up-arrow, tab-completion, etc.

All changes are kept in-memory until you're ready to write them back to your checkpoint with save.

Supported commands

• cd - change working directory.
• pwd - print working directory.
• ls - list directory contents.
• cat - print the contents of a value or directory.
• cp - copy value or directory.
• mv - move/rename value or directory.
• rm - remove value or directory.
• parameters - print the number of model parameters under a directory.
• shape - print tensor shape.
• device - get or set the device of a tensor or group of tensors.
• save - write back changes to disk.
• where - print the location on disk where changes will be saved.
• exit - exits the shell.

FAQ

Should this be a mountable user-mode filesystem instead?

Maybe? Some operations (e.g. shape, parameters, device) don't map easily onto standard filesystem operations. On the other hand, it would be interesting to insert/extract tensors by copying NumPy files in and out of the virtual filesystem.