# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python [conda env:py39]
#     language: python
#     name: conda-env-py39-py
# ---

# +
import torch
import torch.nn as nn
import git
import os

from torch.optim import SGD
# -

CHECKPOINTS_PATH = os.path.join(os.getcwd(), 'checkpoints')
os.makedirs(CHECKPOINTS_PATH, exist_ok=True)
CHECKPOINTS_PATH

# +
X = torch.rand(10, 10)
y = torch.add(X, 2.)
cfg = dict(d_in=10, d_hidden=8)

model = torch.nn.Sequential(nn.Linear(cfg['d_in'], cfg['d_hidden']),
                            nn.ReLU(),
                            nn.Linear(cfg['d_hidden'], 1),
                            nn.ReLU())

optimizer = SGD(model.parameters(), lr=1e-2)


# -

def get_training_state(cfg, model):
    training_state = {
        "commit_hash": git.Repo(search_parent_directories=True).head.object.hexsha,

        # Model structure
        "d_in": d_in,
        "d_out": d_hidden,

        # Model state
        "state_dict": model.state_dict()
    }

    return training_state


torch.save(get_training_state(cfg, model),
           os.path.join(CHECKPOINTS_PATH, 'my_checkpoint.pth'))

model.load_state_dict(torch.load(os.path.join(CHECKPOINTS_PATH, 'my_checkpoint.pth'))['state_dict'])



StarcoderdataJupyterScriptsDedupFiltered