Skip to content

Commit

Permalink
Add migrate_from_folder in readme.
Browse files Browse the repository at this point in the history
  • Loading branch information
bdvllrs committed Mar 20, 2024
1 parent a443ed7 commit 95fc2cc
Showing 1 changed file with 36 additions and 2 deletions.
38 changes: 36 additions & 2 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -2,10 +2,10 @@

```python
import torch
from migrate_ckpt import Migration, migrate_ckpt
from migrate_ckpt import CkptType, Migration, migrate_ckpt


def update_some_keys_callback(ckpt):
def update_some_keys_callback(ckpt: CkptType) -> CkptType:
"""
Define a callback that takes a checkpoints and updates it.
"""
Expand Down Expand Up @@ -36,3 +36,37 @@ Note: the list of migration to perform is determined by the last done migration.
Missed migration in between will never be done.
For example, if migrations to do are ["0", "1", "2"] and model has already had migration
"1", only "2" will be done, but not "0".

## Store migrations in a folder

It migth be convenient to store all your migrations in a specific folder, and execute
migrations from this folder. You can do this with `migrate_from_folder` function.

For example, create a `migrations` folder with this files (loaded in alphabetical order):
```
0_initial_migration.py
1_second_migration.py
2_this_comes_next.py
```

In each migration file, you need to define a function:
```python
from migrate_ckpt import CkptType


def handle(ckpt: CkptType) -> CkptType:
# do stuff with ckpt
return ckpt
```


You can then execute all migrations in the folder with:
```python
import torch

from migrate_ckpt import migrate_from_folder

ckpt, done_migrations = migrate_from_folder(
torch.load("/path/to/some/checkpoint.ckpt"), "path/to/migration/folder"
)
```

0 comments on commit 95fc2cc

Please sign in to comment.