Let's consider an example in Pytorch where we define our own new Model with a class.
import torchimport dstack as dsfrom dstack.torch.handlers import TorchModelEncoder# define a new modelclass LinearRegression(torch.nn.Module):def __init__(self, input_size, output_size):super(LinearRegression, self).__init__()self.linear = torch.nn.Linear(input_size, output_size)def forward(self, x):out = self.linear(x)return outmodel = LinearRegression(1, 1)# here you are training the modelfor epoch in range(100):...# to avoid compatibility issues we will store only model weightsTorchModelEncoder.STORE_WHOLE_MODEL = False
# finally push the modelds.push("my_torch_model", model, "My first PyTorch model")
We stored only model weights, so to pull it we should provide model class to decoder, because
pull method is not smart enough to guess which particular class to use. The following example shows a common pattern how to use pull in this case:
You can also import the Decoder from
dstack.torch.handlers like this -
from dstack.torch.handlers import TorchModelWeightsDecoder
my_model = ds.pull("my_torch_model", decoder=TorchModelWeightsDecoder(LinearRegression(1, 1)))
That's it! Now we can use this pulled model for any task we'd like and also share it with the rest of the team to collaborate.