muellerzr's picture
muellerzr HF staff
Include better doc links that align with the TOC
7576dfc
raw
history blame
No virus
1.56 kB
##
<pre>
from accelerate import Accelerator
accelerator = Accelerator()
dataloader, model, optimizer scheduler = accelerator.prepare(
dataloader, model, optimizer, scheduler
)
for batch in dataloader:
optimizer.zero_grad()
inputs, targets = batch
outputs = model(inputs)
loss = loss_function(outputs, targets)
accelerator.backward(loss)
optimizer.step()
scheduler.step()
+accelerator.save_state("checkpoint_dir")
+accelerator.load_state("checkpoint_dir")</pre>
##
To save or load a checkpoint in, `Accelerator` provides the `save_state` and `load_state` methods.
These methods will save or load the state of the model, optimizer, scheduler, as well as random states and
any custom registered objects from the main process on each device to a passed in folder.
**This API is designed to save and resume training states only from within the same python script or training setup.**
##
To learn more checkout the related documentation:
- <a href="https://huggingface.co/docs/accelerate/usage_guides/checkpoint" target="_blank">Saving and loading training states</a>
- <a href="https://huggingface.co/docs/accelerate/package_reference/accelerator#accelerate.Accelerator.save_state" target="_blank">`save_state` API reference</a>
- <a href="https://huggingface.co/docs/accelerate/package_reference/accelerator#accelerate.Accelerator.load_state" target="_blank">`load_state` API reference</a>
- <a href="https://github.com/huggingface/accelerate/blob/main/examples/by_feature/checkpointing.py" target="_blank">Example script</a>