def pytorch_model_without_validation():
mlflow.pytorch.autolog()
model = IrisClassificationWithoutValidation()
dm = IrisDataModuleWithoutValidation()
dm.setup(stage="fit")
trainer = pl.Trainer(max_epochs=NUM_EPOCHS)
trainer.fit(model, dm)
client = mlflow.tracking.MlflowClient()
run = client.get_run(client.list_run_infos(experiment_id="0")[0].run_id)
return trainer, run
def test_pytorch_autolog_log_models_configuration(log_models):
mlflow.pytorch.autolog(log_models=log_models)
model = IrisClassificationWithoutValidation()
dm = IrisDataModule()
dm.prepare_data()
dm.setup(stage="fit")
trainer = pl.Trainer(max_epochs=NUM_EPOCHS)
trainer.fit(model, dm)
client = mlflow.tracking.MlflowClient()
run = client.get_run(client.list_run_infos(experiment_id="0")[0].run_id)
run_id = run.info.run_id
client = mlflow.tracking.MlflowClient()
artifacts = [f.path for f in client.list_artifacts(run_id)]
assert ("model" in artifacts) == log_models