Exemplo n.º 1
0
def test_pytorch_model_save():
    artifacts_folder = os.path.join(TESTS_PATH, "serve", "data",
                                    "mlflow_pytorch", "model")
    pytorch_mnist_model = Model(
        name="test-pytorch-mnist",
        platform=ModelFramework.MLFlow,
        local_folder=artifacts_folder,
        description="A pytorch MNIST model - python 3.7",
    )
    pytorch_mnist_model.save(save_env=True)

    remote_model = deploy_local(pytorch_mnist_model)
    data = np.random.randn(1, 28 * 28).astype(np.float32)

    remote_model.predict(data)
    remote_model.undeploy()
Exemplo n.º 2
0
def test_model_save(custom_model: Model):
    custom_model.save(save_env=False)
    loaded = Model.load(custom_model.details.local_folder)

    assert len(custom_model.context.__dict__) > 0
    assert len(loaded.context.__dict__) == 0