A PyTorch helper library designed to save only the changes in a fine-tuned base model.
Project description
pytorch-diff-checkpoint
pytorch-diff-checkpoint
is a simple library designed to efficiently save only the modified parameters of a fine-tuned base model. This tool is particularly advantageous in scenarios where minimizing storage usage is crucial, as it ensures that only the altered parameters are stored.
Installation
poetry add pytorch-diff-checkpoint
Usage
import torch
from torch.nn import Module
from diff_checkpoint import DiffCheckpoint
class SimpleModel(Module):
def __init__(self):
super(SimpleModel, self).__init__()
self.fc1 = torch.nn.Linear(10, 10)
self.bn1 = torch.nn.BatchNorm1d(10)
self.fc2 = torch.nn.Linear(10, 1)
def forward(self, x):
x = torch.relu(self.fc1(x))
x = self.bn1(x)
x = self.fc2(x)
return x
model = SimpleModel()
# Create a DiffCheckpoint from the base model
diff_checkpoint = DiffCheckpoint.from_base_model(model)
# Train
# ...
# Save the differential checkpoint
diff_checkpoint.save(model, 'diff_checkpoint.pth')
Project details
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distributions
No source distribution files available for this release.See tutorial on generating distribution archives.
Built Distribution
Close
Hashes for pytorch_diff_checkpoint-1.0.0-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 489b7bf5a6c5e6208b3be61537f8dfd8adb253f38b234a6db1bb5fb4582a0509 |
|
MD5 | 0d712201b3f32066161b20bd46ca0ab4 |
|
BLAKE2b-256 | b6e51916430da0934b52827579561a17553b34552c5a9335e36434eb65637aaa |