main.py
# -*- encoding: utf-8 -*-
import torch
from torch import nn
import mlflow
import mlflow.pytorch
x = torch.tensor([[1.0], [2.0], [3.0]])
y = torch.tensor([[2.0], [4.0], [6.0]])
class Model(nn.Module):
def __init__(self):
super(Model, self).__init__()
self.liner = nn.Linear(1, 1)
def forward(self, x):
y_pred = self.liner(x)
return y_pred
model = Model()
loss_fn = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
num_epochs = 100
for epoch in range(num_epochs):
# forward
y_predict = model(x)
# compute loss
loss = loss_fn(y_predict, y)
print(epoch, loss.data.item())
# update parameters
optimizer.zero_grad()
loss.backward()
optimizer.step()
# After training
for i in [4.0, 5.0, 6.0]:
hour_var = torch.tensor([i])
print('predict (after training): {},{}'.format(i, model(hour_var).data.item()))
# log model to service
with mlflow.start_run() as run:
mlflow.log_param('epoch', num_epochs)
mlflow.pytorch.log_model(model, 'model_example')
conda.yaml
name: pytorch-example
channels:
- https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/free/
- https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main/
dependencies:
- numpy==1.18.1
- tqdm==4.32.1
- scikit-learn==0.21.1
- pytorch==1.4.0
- pip:
- torchtext==0.4.0
- torchvision==0.5.0
- mlflow
MLproject
name: pytorch-example
conda_env: conda.yaml
entry_points:
main:
command: |
python main.py
训练模型并保存:
mlflow run pytorch_example
启动服务
mlflow models serve -m runs:/937f655626b54bdeb905d326e0bd4c47/model_example -h 172.16.1.120 -p 1234
MLmodel
artifact_path: model_example
flavors:
python_function:
data: data
env: conda.yaml
loader_module: mlflow.pytorch
pickle_module_name: mlflow.pytorch.pickle_module
python_version: 3.6.2
pytorch:
model_data: data
pytorch_version: 1.4.0
run_id: 937f655626b54bdeb905d326e0bd4c47
utc_time_created: '2020-04-20 13:31:54.677497'
http请求POST:
{
"columns": ["x1"],
"data": [2.0, 3.0]
}