An open API service indexing awesome lists of open source software.

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 🤗⚡️

Awesome Lists containing this project

README

          

# hf-hub-lightning

Open In Colab

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)
```