JellyML is an open-source tool (python API and command line) for effortlessly embedding a snapshot of your code into a checkpoint of a pytorch (and pytorch lightning) model.
What is JellyML?
$ pip install jellyml
JellyML comes with a getting started guide built-in, just run:
$ jellyml -h
What can JellyML do?
With JellyML you can go back to any of your checkpoints, and reproduce your code exactly as it was when you trained it.
All of this without having to worry about git commits, branches, configuration files, or any of that other stuff. It just works.
It only takes a few lines of code to set up a snapshot of your code:
import jellyml
import torch
# At the start of your training script
# create a snapshot of your code
snapshot = jellyml.create_jelly_filling()
# or use its alias "create_snapshot"
# if you are not in a whimsical mood
snapshot = jellyml.create_snapshot()
# do some stuff, like actually
# train your model...
...
# then, add one extra line when
# you save your model
torch.save({
"model_state_dict":
model.state_dict(),
# this will add the
# snapshot to the model file
**snapshot,
}, "my_model.pth")
And loading a snapshot is just one command away:
$ jellyml load-snapshot-from my_model.pth
But wait, that's not all!
As an added bonus, you can easily share your model with others, and they can reproduce it with a single command. The entire model, code included, is in one file. You can upload it, put it on a flash drive, email it, print it out, or put it on a QR code.
No more messing with git lfs.
You can read more about why I built JellyML in this blog post.
Why is it called JellyML?
Think of your model checkpoint as a donut, and your code as the delicious jelly filling. See Figure 1:
How does JellyML work?
JellyML is built on top of the world's most popular version control system, git.
It's not a replacement for git, rather it works with git. JellyML handles the git commands for you, so you don't have to read the git book again.If you want to learn more about how it works, each command has a detailed explanation in the Documentation.