How to use the mlflow.log_metric function in mlflow

def test_start_and_end_run(tracking_uri_mock):
    # Use the start_run() and end_run() APIs without a `with` block, verify they work.

    with start_run() as active_run:
        mlflow.log_metric("name_1", 25)
    finished_run = tracking.MlflowClient().get_run(active_run.info.run_id)
    # Validate metrics
    assert len(finished_run.data.metrics) == 1
    assert finished_run.data.metrics["name_1"] == 25

def test_search_runs(tracking_uri_mock, reset_active_experiment):
    mlflow.set_experiment("exp-for-search")
    # Create a run and verify that the current active experiment is the one we just set
    logged_runs = {}
    with mlflow.start_run() as active_run:
        logged_runs["first"] = active_run.info.run_id
        mlflow.log_metric("m1", 0.001)
        mlflow.log_metric("m2", 0.002)
        mlflow.log_metric("m1", 0.002)
        mlflow.log_param("p1", "a")
        mlflow.set_tag("t1", "first-tag-val")
    with mlflow.start_run() as active_run:
        logged_runs["second"] = active_run.info.run_id
        mlflow.log_metric("m1", 0.008)
        mlflow.log_param("p2", "aa")
        mlflow.set_tag("t2", "second-tag-val")

    def verify_runs(runs, expected_set):
        assert set([r.info.run_id for r in runs]) == set([logged_runs[r] for r in expected_set])

    experiment_id = MlflowClient().get_experiment_by_name("exp-for-search").experiment_id

    # 2 runs in this experiment

with mlflow.start_run():

        # lr = LinearRegression()
        er = ElasticNet()

        er.fit(train_x, train_y)
        predicted_qualities = er.predict(test_x)

        (rmse, mae, r2) = eval_metrics(test_y, predicted_qualities)

        print("  RMSE: %s" % rmse)
        print("  MAE: %s" % mae)
        print("  R2: %s" % r2)

        mlflow.log_param("coefficients", er.coef_)
        mlflow.log_metric("rmse", rmse)
        mlflow.log_metric("r2", r2)
        mlflow.log_metric("mae", mae)

        tracking_url_type_store = urlparse(mlflow.get_tracking_uri()).scheme
        joblib.dump(er, output_dir + "/" + "ElasticNet.pkl")
        # Model registry does not work with file store
        if tracking_url_type_store != "file":
            # Register the model
            # There are other ways to use the Model Registry, which depends on the use case,
            # please refer to the doc for more information:
            # https://mlflow.org/docs/latest/model-registry.html#api-workflow
            # mlflow.sklearn.log_model(lr, "ElasticnetModel", registered_model_name="ElasticnetWineModel")

            mlflow.log_artifact(output_dir, "ElasticNet", registered_model_name="ElasticNetWineModel")
        else:
            mlflow.log_artifact(output_dir, "ElasticNet")

#Train model
    results = weathernet.fit_generator(datagen_train(), steps_per_epoch=steps_per_epoch, workers=10, max_queue_size=100, epochs=epochs, verbose=2, validation_steps=val_steps, validation_data=datagen_val())
    #Save trained model
    weathernet.save(filepath=model_filepath)
    # Run prediction on trained model
    predict_weather(weathernet)
    # Plot the metrics of the trained model
    plot_metrics(results)
    # Log metrics, parameters, artifacts and log the model
    with mlflow.start_run():
        run_uuid = mlflow.active_run().info.run_uuid
        print("MLflow Run ID: %s" % run_uuid)
        mlflow.keras.log_model(weathernet, "models")
        mlflow.log_artifact(image_dir + str(now) + '_' + city + '_Loss_Diag.png', "images")
        mlflow.log_artifact(image_dir + str(now) + '_' + city + '_Daily_Temp_Predicted.png', "images")
        mlflow.log_metric('Loss', loss(results))
        mlflow.log_metric('Validation Loss', val_loss(results))
        mlflow.log_param('City_Name', city)
        mlflow.log_param('Training_Epochs', epochs)
        mlflow.log_param('Steps_per_epoch', steps_per_epoch)
        mlflow.log_param('Validations_steps', val_steps)
        mlflow.log_param('Prediction_steps', predict_steps)

def eval_and_log_metrics(prefix, actual, pred, epoch):
    rmse = np.sqrt(mean_squared_error(actual, pred))
    mlflow.log_metric("{}_rmse".format(prefix), rmse, step=epoch)
    return rmse

run_settings = CustomSettings(data=args.data, neg=32, seed=int(time.time()), pretrain=args.pretrain,
                                          loss_mode=loss_mode, lr_match=lr_match, l2=l2, dropout_p=dropout_p,
                                          tim_emb_size=tim_emb_size, loc_emb_size=loc_emb_size,
                                          hidden_size=hidden_size, epoch=args.epoch, threshold=args.threshold,
                                          rnn_mod=rnn_unit, attn_mod=attn_unit, save_path=archive_path,
                                          noise=0 if USE_POI else args.noise_level, poi_type=args.poi_type)
            model, rank, hit, rank_pre, hit_pre = run_experiments(run_settings, model_type=args.model,
                                                                  run_id=run_id,
                                                                  device=device, USE_POI=USE_POI,
                                                                  unit=args.pretrain_unit)
            if run_id == 0:
                rank_pre2, hit_pre2 = rank_pre, hit_pre
            mlflow.log_metric("rank_32", rank)
            mlflow.log_metric("hit_32", hit)
            mlflow.log_metric("rank_pre", rank_pre2)
            mlflow.log_metric("hit_pre", hit_pre2)

def _mlflow_log_metrics(metrics, metric_name):
    """Record metric value during each epoch using the step parameter in
    mlflow.log_metric.

    :param metrics:
    :param metric_name:
    :return:
    """
    for epoch, metric in enumerate(metrics[metric_name], 1): mlflow.log_metric(
        metric_name, metric,
        step=epoch)

def on_epoch_end(self, epoch, logs=None):
        """
        Log Keras metrics with MLflow. Update the best model if the model improved on the validation
        data.
        """
        if not logs:
            return
        for name, value in logs.items():
            if name.startswith("val_"):
                name = "valid_" + name[4:]
            else:
                name = "train_" + name
            mlflow.log_metric(name, value)
        val_loss = logs["val_loss"]
        if val_loss < self._best_val_loss:
            # Save the "best" weights
            self._best_val_loss = val_loss
            self._best_weights = [x.copy() for x in self._model.get_weights()]

)

                loss = outputs[0]  # model outputs are always tuple in pytorch-transformers

                loss.sum().backward()
                torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.max_grad_norm)

                tr_loss += loss.sum().item()
                optimizer.step()
                # Update learning rate schedule
                scheduler.step()
                optimizer.zero_grad()
                global_step += 1
                # logging of learning rate and loss
                if logging_steps > 0 and global_step % logging_steps == 0:
                    mlflow.log_metric("learning rate", scheduler.get_lr()[0], step=global_step)
                    mlflow.log_metric(
                        "training loss",
                        (tr_loss - logging_loss) / (logging_steps * self.batch_size),
                        step=global_step,
                    )
                    logging_loss = tr_loss
                # model checkpointing
                if save_steps > 0 and global_step % save_steps == 0:
                    checkpoint_dir = os.path.join(os.getcwd(), "checkpoints")
                    if not os.path.isdir(checkpoint_dir):
                        os.makedirs(checkpoint_dir)
                    checkpoint_path = checkpoint_dir + "/" + str(global_step) + ".pth"
                    torch.save(self.model.state_dict(), checkpoint_path)
                    mlflow.log_artifact(checkpoint_path)
                # model validation
                if val_steps > 0 and global_step % val_steps == 0:

def train(args, model, device, train_loader, optimizer, epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()
        if batch_idx % args.log_interval == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), loss.item()))
            # Use MLflow logging
            mlflow.log_metric("epoch_loss", loss.item())