def test_pytorch_with_early_stopping_autolog_log_models_configuration_with(log_models, patience):
mlflow.pytorch.autolog(log_models=log_models)
model = IrisClassification()
dm = IrisDataModule()
dm.prepare_data()
dm.setup(stage="fit")
early_stopping = EarlyStopping(monitor="val_loss", mode="min", patience=patience, verbose=True)
with TempDir() as tmp:
keyword = "dirpath" if LooseVersion(pl.__version__) >= LooseVersion("1.2.0") else "filepath"
checkpoint_callback = ModelCheckpoint(
**{keyword: tmp.path()}, save_top_k=1, verbose=True, monitor="val_loss", mode="min",
)
trainer = pl.Trainer(
max_epochs=NUM_EPOCHS * 2, callbacks=[early_stopping, checkpoint_callback],
)
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 ("restored_model_checkpoint" in artifacts) == log_models
def pytorch_model_with_callback(patience):
mlflow.pytorch.autolog()
model = IrisClassification()
dm = IrisDataModule()
dm.prepare_data()
dm.setup(stage="fit")
early_stopping = EarlyStopping(
monitor="val_loss",
mode="min",
min_delta=99999999, # forces early stopping
patience=patience,
verbose=True,
)
with TempDir() as tmp:
keyword = "dirpath" if LooseVersion(pl.__version__) >= LooseVersion("1.2.0") else "filepath"
checkpoint_callback = ModelCheckpoint(
**{keyword: tmp.path()}, save_top_k=1, verbose=True, monitor="val_loss", mode="min",
)
trainer = pl.Trainer(
max_epochs=NUM_EPOCHS * 2, callbacks=[early_stopping, checkpoint_callback],
)
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 pytorch_model_tests():
model = IrisClassification()
dm = IrisDataModule()
dm.prepare_data()
dm.setup(stage="fit")
trainer = pl.Trainer(max_epochs=NUM_EPOCHS)
trainer.fit(model, dm)
trainer.test()
client = mlflow.tracking.MlflowClient()
run = client.get_run(client.list_run_infos(experiment_id="0")[0].run_id)
return trainer, run
def train_evaluate(params, max_epochs=100):
model = IrisClassification(**params)
dm = IrisDataModule()
dm.prepare_data()
dm.setup(stage="fit")
trainer = pl.Trainer(max_epochs=max_epochs)
mlflow.pytorch.autolog()
trainer.fit(model, dm)
trainer.test(datamodule=dm)
test_accuracy = trainer.callback_metrics.get("test_acc")
return test_accuracy
def test_pytorch_autolog_persists_manually_created_run():
with mlflow.start_run() as manual_run:
mlflow.pytorch.autolog()
model = IrisClassification()
dm = IrisDataModule()
dm.prepare_data()
dm.setup(stage="fit")
trainer = pl.Trainer(max_epochs=NUM_EPOCHS)
trainer.fit(model, dm)
trainer.test()
assert mlflow.active_run() is not None
assert mlflow.active_run().info.run_id == manual_run.info.run_id
def pytorch_model_tests():
mlflow.pytorch.autolog()
model = IrisClassification()
dm = IrisDataModule()
dm.prepare_data()
dm.setup(stage="fit")
trainer = pl.Trainer(max_epochs=NUM_EPOCHS)
with mlflow.start_run() as run:
trainer.fit(model, dm)
trainer.test()
client = mlflow.tracking.MlflowClient()
run = client.get_run(run.info.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
"--save-model",
type=bool,
default=True,
help="For Saving the current Model",
)
parser.add_argument(
"--accelerator",
type=lambda x: None if x == "None" else x,
default=None,
help="Accelerator - (default: None)",
)
from iris_data_module import IrisDataModule
parser = IrisClassification.add_model_specific_args(parent_parser=parser)
parser = IrisDataModule.add_model_specific_args(parent_parser=parser)
args = parser.parse_args()
dict_args = vars(args)
dm = IrisDataModule(**dict_args)
dm.prepare_data()
dm.setup(stage="fit")
model = IrisClassification(**dict_args)
trainer = pl.Trainer.from_argparse_args(args)
trainer.fit(model, dm)
trainer.test()
torch.save(model.state_dict(), "iris.pt")
