$ sudo apt update
$ sudo apt-get install -y postgresql postgresql-contrib postgresql-server-dev-all
$ sudo apt install gcc
$ pip install psycopg2
$ sudo -u postgres -i
# Create new user: mlflow_user
[postgres@user_name~]$ createuser --interactive -P
Enter name of role to add: mlflow_user
Enter password for new role: mlflow
Enter it again: mlflow
Shall the new role be a superuser? (y/n) n
Shall the new role be allowed to create databases? (y/n) n
Shall the new role be allowed to create more new roles? (y/n) n
# Create database mlflow_bd owned by mlflow_user
$ createdb -O mlflow_user mlflow_db

pip install mlflow
pip install pysftp

$ mlflow server  \
                 --backend-store-uri postgresql://mlflow_user:mlflow@localhost/mlflow_db \
                 --default-artifact-root sftp://mlflowsftp:mlflow@sftp_host/upload  \
                --host server_host \
                --port server_port

def run(self, epochs, lr, experiment_name):
        # getting the id of the experiment, creating an experiment in its absence
        remote_experiment_id = self.remote_server.get_experiment_id(name=experiment_name)
        # creating a "run" and getting its id
        remote_run_id = self.remote_server.get_run_id(remote_experiment_id)
        # indicate that we want to save the results on a remote server
        mlflow.set_tracking_uri(self.tracking_uri)
        mlflow.set_experiment(experiment_name)
        with mlflow.start_run(run_id=remote_run_id, nested=False):
            mlflow.keras.autolog()
            self.train_pipeline.train(lr=lr, epochs=epochs)
        try:
            self.log_tags_and_params(remote_run_id)
        except mlflow.exceptions.RestException as e:
            print(e)

name: flow_segmentation
conda_env: conda.yaml
entry_points:
  main:
    parameters:
        categories: {help: 'list of categories from coco dataset'}
        epochs: {type: int, help: 'number of epochs in training'}
        lr: {type: float, default: 0.001, help: 'learning rate'}
        batch_size: {type: int, default: 8}
        model_name: {type: str, default: 'Unet', help: 'Unet, PSPNet, Linknet, FPN'}
        backbone_name: {type: str, default: 'resnet18', help: 'exampe resnet18, resnet50, mobilenetv2 ...'}
        tracking_uri: {type: str, help: 'the server address'}
        experiment_name: {type: str, default: 'My_experiment', help: 'remote and local experiment name'}
    command: "python mlflow_training.py \
            --epochs={epochs}
            --categories={categories}
            --lr={lr}
            --tracking_uri={tracking_uri}
            --model_name={model_name}
            --backbone_name={backbone_name}
            --batch_size={batch_size}
            --experiment_name={experiment_name}"

name: flow_segmentation
channels:
  - defaults
  - anaconda
dependencies:
  - python==3.7
  - pip:
    - mlflow==1.8.0
    - pysftp==0.2.9
    - Cython==0.29.19
    - numpy==1.18.4
    - pycocotools==2.0.0
    - requests==2.23.0
    - matplotlib==3.2.1
    - segmentation-models==1.0.1
    - Keras==2.3.1
    - imgaug==0.4.0
    - tqdm==4.46.0
    - tensorflow-gpu==1.14.0

git clone https://github.com/simbakot/mlflow_example.git
cd mlflow_example/

pip install mlflow
pip install pysftp

$ mlflow run -P epochs=10 -P categories=cat,dog -P tracking_uri=http://server_host:server_port .

class SegmentationModel:
    def __init__(self, tracking_uri, model_name):
        self.registry = RemoteRegistry(tracking_uri=tracking_uri)
        self.model_name = model_name
        self.model = self.build_model(model_name)
    def get_latest_model(self, model_name):
        registered_models = self.registry.get_registered_model(model_name)
        last_model = self.registry.get_last_model(registered_models)
        local_path = self.registry.download_artifact(last_model.run_id, 'model', './')
        return local_path
    def build_model(self, model_name):
        local_path = self.get_latest_model(model_name)
        return mlflow.keras.load_model(local_path)
    def predict(self, image):
        image = self.preprocess(image)
        result = self.model.predict(image)
        return self.postprocess(result)
    def preprocess(self, image):
        image = cv2.resize(image, (256, 256))
        image = image / 255.
        image = np.expand_dims(image, 0)
        return image
    def postprocess(self, result):
        return result