Jovian
⭐️
Sign In

Stockroom API design

In [ ]:
from stockroom import StockRoom
stock = StockRoom()

modelstore

In [ ]:
import torch

net = torch.nn.Sequential(
    torch.nn.Linear(10, 6),
    torch.nn.ReLU(),
    torch.nn.Linear(6, 3))

Saving

In [ ]:
# storing state dict
stock.model['SeqNet'] = net.state_dict()

# or letting stockroom figure out how to fetch the state dict
stock.model.save('SeqNet', net)

loading

In [ ]:
# fetching state dict
state_dict = stock.model['SeqNet']
net.load_state_dict(state_dict)

# Letting stockroom does the state dict loading
stock.model.load('SeqNet', net)

datastore

In [ ]:
import numpy as np

stock.data['arraysetname', 'samplename'] = np.array([1, 2])

# or

aset = stock.data['arraysetname']
aset['samplename'] = np.array([1, 2])

But if we are managing the checkout internally, we need another way to let the user control the checkout and triggering the CM to optimize the transaction

In [ ]:
# on the CM open we'll make the checkout open and open the CM of the checkout
# for making the read and write optimized
with stock.optimize():
    for i, data in enumerage(get_data()):
        stock.data['arraysetname', i] = data

genericstore

In [ ]:
stock.tag['epochs'] = epochs
stock.tag['optimizer'] = 'SGD'
for e in range(epochs):
    for i, data in enumerate(get_batch_from_stockroom()):
        optimizer.zero_grad)_
        out = model(data)
        loss = loss_fn(out, target)
        loss.backward()
        optimizer.step()
        stock.tag[f'loss_{e}_{i}'] = loss.item()
In [ ]:
import jovian
jovian.commit()
[jovian] Saving notebook..
In [ ]: