Add MLproject for MLFlow

This commit is contained in:
Marcin Kostrzewski 2022-05-09 17:05:50 +02:00
parent d64d3c1d5a
commit 60a565098c
3 changed files with 16 additions and 1 deletions

1
.gitignore vendored
View File

@ -20,3 +20,4 @@ evaluation_results.txt
model_out
trend.png
sacred_runs
mlruns

14
MLproject Normal file
View File

@ -0,0 +1,14 @@
name: ium_s444409
docker_env:
image: s444409-create-dataset
volumes: ["/tmp/mlruns:/mlruns:rw"]
entry_points:
main:
parameters:
epochs: {type: float, default: 5}
batch_size: {type: float, default: 64}
command: "python train_model.py -e {epochs} -b {batch_size}"
eval:
command: "python eval_model.py"

View File

@ -12,7 +12,7 @@ default_epochs = 5
device = "cuda" if torch.cuda.is_available() else "cpu"
mlflow.set_tracking_uri("http://172.17.0.1:5000")
#mlflow.set_tracking_uri("http://172.17.0.1:5000")
mlflow.set_experiment("s444409")