https://github.com/nateraw/hf-hub-lightning
A PyTorch Lightning Callback for pushing models to the Hugging Face Hub 🤗⚡️
https://github.com/nateraw/hf-hub-lightning
huggingface machine-learning pytorch-lightning
Last synced: about 1 year ago
JSON representation
A PyTorch Lightning Callback for pushing models to the Hugging Face Hub 🤗⚡️
- Host: GitHub
- URL: https://github.com/nateraw/hf-hub-lightning
- Owner: nateraw
- License: mit
- Created: 2021-10-27T00:26:17.000Z (over 4 years ago)
- Default Branch: main
- Last Pushed: 2022-05-13T03:39:21.000Z (about 4 years ago)
- Last Synced: 2025-03-30T00:05:39.581Z (over 1 year ago)
- Topics: huggingface, machine-learning, pytorch-lightning
- Language: Python
- Homepage:
- Size: 18.6 KB
- Stars: 36
- Watchers: 2
- Forks: 3
- Open Issues: 1
-
Metadata Files:
- Readme: README.md
- License: LICENSE
Awesome Lists containing this project
README
# hf-hub-lightning
A callback for pushing lightning models to the Hugging Face Hub.
**Note:** I made this package for myself, mostly...if folks seem to be interested in it, we'll move this into `huggingface_hub` or an official PyTorch Lightning repo directly and archive this one. Since there are a lot of options/considerations, I decided to just post it separately for now as a proof of concept.
## Setup
```
pip install hf-hub-lightning
```
## Usage
To periodically upload your model ckpt and TensorBoard logs while training, you can do the following...
```python
import pytorch_lightning as pl
from hf_hub_lightning import HuggingFaceHubCallback
train_dataloader = ...
model = YourCoolModel()
trainer = pl.Trainer(callbacks=[HuggingFaceHubCallback('your_username/model_id')])
trainer.fit(model, train_dataloader)
```
To load your model back from the huggingface hub, just do...
```python
from huggingface_hub import hf_hub_download
ckpt_path = hf_hub_download('your_username/model_id', 'lit_model.ckpt')
model = YourCoolModel.load_from_checkpoint(ckpt_path)
```