43 lines
1.6 KiB
Python
43 lines
1.6 KiB
Python
import torch
|
|
import numpy as np
|
|
import torch.nn as nn
|
|
import torch.optim as optim
|
|
from torch.utils.data import Dataset, TensorDataset, DataLoader
|
|
import mlflow
|
|
import mlflow.pytorch
|
|
from urllib.parse import urlparse
|
|
|
|
from mlflow.models.signature import infer_signature
|
|
|
|
class LayerLinearRegression(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
# Instead of our custom parameters, we use a Linear layer with single input and single output
|
|
self.linear = nn.Linear(1, 1)
|
|
|
|
def forward(self, x):
|
|
# Now it only takes a call to the layer to make predictions
|
|
return self.linear(x)
|
|
|
|
checkpoint = torch.load('model.pt')
|
|
|
|
model = LayerLinearRegression()
|
|
#optimizer = optim.SGD(model.parameters(), lr=checkpoint['loss'])
|
|
|
|
model.load_state_dict(checkpoint['model_state_dict'])
|
|
|
|
train_dataset = torch.load('train_dataset.pt')
|
|
x_train = np.array(train_dataset)[:,0] #(Sales Sum row)
|
|
input_example = np.reshape(x_train, (-1,1))
|
|
|
|
with torch.no_grad():
|
|
model.eval()
|
|
siganture = infer_signature(x_train, model(torch.tensor(np.reshape(x_train, (-1,1))).float()).numpy())
|
|
mlflow.set_tracking_uri("http://172.17.0.1:5000")
|
|
tracking_url_type_store = urlparse(mlflow.get_tracking_uri()).scheme
|
|
# print(tracking_url_type_store)
|
|
# Model registry does not work with file store
|
|
if tracking_url_type_store != "file":
|
|
mlflow.sklearn.log_model(model, "model", registered_model_name="s426206", signature=siganture, input_example=input_example)
|
|
else:
|
|
mlflow.sklearn.log_model(model, "model", signature=siganture, input_example=input_example) |