90 lines
2.6 KiB
Python
90 lines
2.6 KiB
Python
import torch
|
|
import mlflow
|
|
from urllib.parse import urlparse
|
|
import torch.nn as nn
|
|
import numpy as np
|
|
import pandas as pd
|
|
import matplotlib.pyplot as plt
|
|
import torch
|
|
import datetime
|
|
from torch.autograd import Variable
|
|
|
|
|
|
INPUT_DIM = 1
|
|
OUTPUT_DIM = 1
|
|
LEARNING_RATE = 0.01
|
|
EPOCHS = 100
|
|
|
|
|
|
mlflow.set_tracking_uri("http://172.17.0.1:5000")
|
|
dataset = pd.read_csv('datasets/train_set.csv')
|
|
testset = pd.read_csv('datasets/test_set.csv')
|
|
|
|
x_values = [datetime.datetime.strptime(
|
|
item, "%Y-%m-%d").month for item in dataset['date'].values]
|
|
x_train = np.array(x_values, dtype=np.float32)
|
|
x_train = x_train.reshape(-1, 1)
|
|
|
|
y_values = [min(dataset['result_1'].values[i]/dataset['result_2'].values[i], dataset['result_2'].values[i] /
|
|
dataset['result_1'].values[i]) for i in range(len(dataset['result_1'].values))]
|
|
y_train = np.array(y_values, dtype=np.float32)
|
|
y_train = y_train.reshape(-1, 1)
|
|
|
|
|
|
class LinearRegression(torch.nn.Module):
|
|
def __init__(self, inputSize, outputSize):
|
|
super(LinearRegression, self).__init__()
|
|
self.linear = torch.nn.Linear(inputSize, outputSize)
|
|
|
|
def forward(self, x):
|
|
out = self.linear(x)
|
|
return out
|
|
|
|
|
|
model = LinearRegression(INPUT_DIM, OUTPUT_DIM)
|
|
|
|
|
|
criterion = torch.nn.MSELoss()
|
|
optimizer = torch.optim.SGD(model.parameters(), lr=LEARNING_RATE)
|
|
|
|
for epoch in range(EPOCHS):
|
|
inputs = Variable(torch.from_numpy(x_train))
|
|
labels = Variable(torch.from_numpy(y_train))
|
|
|
|
optimizer.zero_grad()
|
|
|
|
outputs = model(inputs)
|
|
|
|
loss = criterion(outputs, labels)
|
|
print(loss)
|
|
loss.backward()
|
|
|
|
optimizer.step()
|
|
|
|
print('epoch {}, loss {}'.format(epoch, loss.item()))
|
|
|
|
torch.save(model.state_dict(), 'model.pt')
|
|
|
|
with mlflow.start_run():
|
|
test_input = x_train[0]
|
|
|
|
mlflow.log_param("train size", dataset.size)
|
|
mlflow.log_param("test size", testset.size)
|
|
mlflow.log_param("epochs", EPOCHS)
|
|
|
|
predicted = model(Variable(torch.from_numpy(x_train))).data.numpy()
|
|
|
|
signature = mlflow.models.signature.infer_signature(
|
|
x_train, predicted)
|
|
|
|
mlflow.set_experiment("s434700")
|
|
tracking_url_type_store = urlparse(mlflow.get_tracking_uri()).scheme
|
|
if tracking_url_type_store != "file":
|
|
mlflow.pytorch.log_model(model, "cs_model", registered_model_name="s434700", signature=signature,
|
|
input_example=test_input)
|
|
else:
|
|
mlflow.pytorch.log_model(model, "cs_model",
|
|
signature=signature, input_example=test_input)
|
|
mlflow.pytorch.save_model(
|
|
model, "cs_model", signature=signature, input_example=test_input)
|