Repository: mlflow/mlflow

Issue:
[FR] Autologging functionality for PyTorch/fastai
## Describe the proposal

It'd be nice to add an mlflow.pytorch.autolog() API for automatically logging metrics, params & models generated via PyTorch. In particular, I'd like this to be incorporated with the fastai library (built on PyTorch), which also generates artifacts such as learning rate finder plots, and autolog those as well. 

### Motivation
PyTorch is heavily used in research for Deep Learning and is becoming very prevalent in industry as well with the fastai library. The fastai community has over 450 contributors, and it incorporates practical deep learning best practices + cutting edge algorithms in the library. In the fast.ai course that is run by USF, they use Excel spreadsheets to keep track of their runs. If we could just autolog everything once they import mlflow, it will help bring the adoption of MLflow to the fastai community. 

### Proposed Changes
See above - we propose adding a new mlflow.pytorch.autolog API


Discussion:
@shalomma @LarsHill would either of you be interested in working on this feature? If so I'm happy to write up / help collaborate on an API proposal :)
@smurching sorry, but in the near future I won't have time working on it. I just thought it would be a useful feature.
@smurching are you working on it ? I will be happy to help :)
@pierrelandry I haven't worked on it, that'd be awesome! Before writing code & opening a PR, a good first step would be to write up a description of the proposed API & its behavior. Happy to review a proposal if you're able to write one & share it here.

I'm new to Pytorch, but I'd recommend focusing on adding autologging for the most common Pytorch APIs first (will spend some time tomorrow learning more about Pytorch). Also, I'd suggest `mlflow.pytorch.autolog()` as the API name, to match the existing Keras & TF autologging APIs, but open to alternatives. Thanks for being willing to contribute :)!

@pierrelandry hope you enjoyed the holidays! Thanks again for your interest in helping with this - are you still interested in collaborating?
Hi, I'd love to work on the fastai integration, is there any proposed API?
I thought of adding support for logging metrics during training and saving artifacts generated from
[SaveModelCallback](https://docs.fast.ai/callbacks.html#SaveModelCallback) or [Learner.save](https://docs.fast.ai/basic_train.html#Learner.save) method.
@antoniomdk sorry for the delayed response, that'd be awesome - one strawman API proposal would be to add an `mlflow.pytorch.autolog()` API, and an `mlflow.fastai.autolog()` if appropriate - AFAICT fastai & PyTorch are analogous to Keras and TensorFlow (fastai seems to be a higher-level API implemented using PyTorch), so seems reasonable to have separate APIs for the two libraries. 

That being said, it's probably worth considering whether implementing PyTorch autologging could indirectly solve fastai autologging (i.e. could we reuse the implementation of PyTorch autologging for fastai autologging?). My guess is that it might be possible, but the right way to implement fastai autologging would be instead be to use public `fastai` APIs, e.g. by adding callbacks to save & log the user's model as you described - I don't know too much about either PyTorch or fastai so curious to hear your opinion.


Yes, you're right, Fastai is a higher-level API on top of Pytorch as Keras is on top of Tensorflow. Adding model saving autologging for Pytorch would solve it also for Fastai, as the latter relies on torch.save
https://github.com/fastai/fastai/blob/master/fastai/basic_train.py#L246
So, monkey patching [torch.save](https://pytorch.org/docs/stable/torch.html?highlight=save#torch.save) would work for both libraries. It'd also work for model checkpoints along epochs. 

On the other hand, I think we should go for adding a callback to handle metric logging on Fastai.
There is also a proposal from the Fastai community to add a callback for this: https://forums.fast.ai/t/mlflow-tracking-support/38893

For Pytorch I'm not sure how we should proceed because the training cycle is done manually and the evaluation depends on the user implementing it.
Adding autologging to PyTorch is much harder inherently because PyTorch's training cycle can be very flexible and can take many forms, like @antoniomdk mentioned. I took a look at what MLFlow autologging does for other frameworks (https://www.mlflow.org/docs/latest/tracking.html#automatic-logging-from-tensorflow-and-keras-experimental), and I think we might be able to do the following:

- Metrics

| Logging Item | Possible Implementation |
| ------------- | ------------- |
| training loss  | Monkey-patch [Tensor.backward](https://pytorch.org/docs/stable/autograd.html#torch.Tensor.backward) method and log the tensor |
| validation loss  | (This is difficult to track because we can't know programmatically which variable is the validation loss. We might have to resort to manual logging.) |
| user-specified metrics | (Not sure how we should map this concept from Keras to PyTorch.) |

- Parameters

| Logging Item | Possible Implementation |
| ------------- | ------------- |
| fit() parameters  | (Not sure how we should map this concept from Keras to PyTorch.) |
| optimizer name  | Find instantiations of [torch.optim.Optimizer](https://pytorch.org/docs/stable/optim.html#torch.optim.Optimizer) subclasses in the code, and log the optimizer name |
| learning rate | Log the `lr` value of the optimizer found above |
| epsilon | Log the `eps` value of the optimizer found above |

- Artifacts

| Logging Item | Possible Implementation |
| ------------- | ------------- |
| model summary on training start | Monkey-patch `Tensor.backward` or `Optimizer.step` method and log the model summary on the first method call |
| MLFlow model (PyTorch model) on training end | The end of training is hard to detect, but maybe we can monkey-patch `Optimizer.step` or `torch.nn.Module.eval` method to trigger MLFlow model saving after every epoch |

> AFAICT fastai & PyTorch are analogous to Keras and TensorFlow (fastai seems to be a higher-level API implemented using PyTorch), so seems reasonable to have separate APIs for the two libraries.

+1 I think it's easier to add autologging to fastai as a separate API because callbacks are a cleaner way to implement autologging.

I think it'd be hard to detect the last epoch using Optimizer.step or torch.nn.Module.eval.
For model persistence it's easier to rely on [torch.save](https://pytorch.org/docs/stable/torch.html?highlight=save#torch.save).

For optimizer tracking we can monkey-patch torch.optim.Optimizer class to keep track of all instances created. We can use those instances to log optimizer parameters, as @jacobhess118 suggested.

For tracking loss I'm not sure if using Tensor.backward would be a good idea (could it make Tensor.backward subcalls?).
@jacobhess118 @antoniomdk thanks for the detailed investigation! Here are my impressions:
* Makes sense to implement fastai autologging via a separate `mlflow.fastai.autolog` API, where calling `mlflow.fastai.autolog` patches [`Learner.fit`](https://docs.fast.ai/basic_train.html#Learner.fit) & modifies the `callbacks` argument, injecting a callback that logs model & training information (see [analogous implementation](https://github.com/mlflow/mlflow/blob/1a762d56564f5f708bd7f9cbc5a48530fb354ae3/mlflow/keras.py#L528) in Keras autologging)
* PyTorch autologging: PyTorch is a popular framework (2x the GitHub stars of fastai) so would be nice to have autologging support, but if it seems likely to be brittle/complex (as it currently seems from @jacobhess118's description), we can punt on it for now - it's possible it'll become easier down the line

Thus @antoniomdk if you're interested in adding `fastai` autologging for now, would be happy to review a PR that does so :)

 > For tracking loss I'm not sure if using Tensor.backward would be a good idea (could it make Tensor.backward subcalls?).

Tensor.backward doesn’t make subcalls of itself at the moment (and likely won’t be in the future) because calling Python functions is slow. However I do think that it smells more like a hack, and we probably shouldn’t do it this way unless absolutely necessary.
> Thus @antoniomdk if you're interested in adding fastai autologging for now, would be happy to review a PR that does so :)

@smurching I'll work on that as soon as I've reviewed the contributing guideline and the Keras implementation. Thanks!
Hi @antoniomdk, just wanted to check in (no rush of course!) - did you get a chance to get started with this and/or have any questions?
Hi @smurching, I started to work on this recently. I've forked the repo and create a new branch but my advances are slow. I'm quite busy at this moment at work, but I hope to make a PR by the end of this month.

Resolving patch:
diff --git a/mlflow/fastai.py b/mlflow/fastai.py
--- a/mlflow/fastai.py
+++ b/mlflow/fastai.py
@@ -14,6 +14,9 @@
 """
 import os
 import yaml
+import gorilla
+import tempfile
+import shutil
 import pandas as pd
 import numpy as np
 
@@ -25,6 +28,8 @@
 from mlflow.tracking.artifact_utils import _download_artifact_from_uri
 from mlflow.utils.environment import _mlflow_conda_env
 from mlflow.utils.model_utils import _get_flavor_configuration
+from mlflow.utils.annotations import experimental
+from mlflow.utils.autologging_utils import try_mlflow_log, log_fn_args_as_params
 
 from fastai.tabular import TabularList
 from fastai.basic_data import DatasetType
@@ -243,3 +248,149 @@ def load_model(model_uri):
     flavor_conf = _get_flavor_configuration(model_path=local_model_path, flavor_name=FLAVOR_NAME)
     model_file_path = os.path.join(local_model_path, flavor_conf.get("data", "model.fastai"))
     return _load_model(path=model_file_path)
+
+
+@experimental
+def autolog():
+    """
+    Enable automatic logging from Fastai to MLflow.
+    Logs loss and any other metrics specified in the fit
+    function, and optimizer data as parameters. Model checkpoints
+    are logged as artifacts to a 'models' directory.
+
+    MLflow will also log the parameters of the EarlyStopping and OneCycleScheduler callbacks
+    """
+    from fastai.basic_train import LearnerCallback, Learner
+    from fastai.callbacks.hooks import model_summary, layers_info
+    from fastai.callbacks import EarlyStoppingCallback, OneCycleScheduler
+
+    class __MLflowFastaiCallback(LearnerCallback):
+        """
+        Callback for auto-logging metrics and parameters.
+        Records model structural information as params when training begins
+        """
+        def __init__(self, learner, ):
+            super().__init__(learner)
+            self.learner = learner
+            self.opt = self.learn.opt
+            self.metrics_names = ['train_loss', 'valid_loss'] + \
+                                 [o.__name__ for o in learner.metrics]
+
+        def on_epoch_end(self, **kwargs):
+            """
+            Log loss and other metrics values after each epoch
+            """
+            if kwargs['smooth_loss'] is None or kwargs["last_metrics"] is None:
+                return
+            epoch = kwargs['epoch']
+            metrics = [kwargs['smooth_loss']] + kwargs["last_metrics"]
+            metrics = map(float, metrics)
+            metrics = dict(zip(self.metrics_names, metrics))
+            try_mlflow_log(mlflow.log_metrics, metrics, step=epoch)
+
+        def on_train_begin(self, **kwargs):
+            info = layers_info(self.learner)
+            try_mlflow_log(mlflow.log_param, 'num_layers', len(info))
+            try_mlflow_log(mlflow.log_param, 'opt_func', self.opt_func.func.__name__)
+
+            if hasattr(self.opt, 'true_wd'):
+                try_mlflow_log(mlflow.log_param, 'true_wd', self.opt.true_wd)
+
+            if hasattr(self.opt, 'bn_wd'):
+                try_mlflow_log(mlflow.log_param, 'bn_wd', self.opt.bn_wd)
+
+            if hasattr(self.opt, 'train_bn'):
+                try_mlflow_log(mlflow.log_param, 'train_bn', self.train_bn)
+
+            summary = model_summary(self.learner)
+            try_mlflow_log(mlflow.set_tag, 'model_summary', summary)
+
+            tempdir = tempfile.mkdtemp()
+            try:
+                summary_file = os.path.join(tempdir, "model_summary.txt")
+                with open(summary_file, 'w') as f:
+                    f.write(summary)
+                try_mlflow_log(mlflow.log_artifact, local_path=summary_file)
+            finally:
+                shutil.rmtree(tempdir)
+
+        def on_train_end(self, **kwargs):
+            try_mlflow_log(log_model, self.learner, artifact_path='model')
+
+    def _find_callback_of_type(callback_type, callbacks):
+        for callback in callbacks:
+            if isinstance(callback, callback_type):
+                return callback
+        return None
+
+    def _log_early_stop_callback_params(callback):
+        if callback:
+            try:
+                earlystopping_params = {'early_stop_monitor': callback.monitor,
+                                        'early_stop_min_delta': callback.min_delta,
+                                        'early_stop_patience': callback.patience,
+                                        'early_stop_mode': callback.mode}
+                try_mlflow_log(mlflow.log_params, earlystopping_params)
+            except Exception:  # pylint: disable=W0703
+                return
+
+    def _log_one_cycle_callback_params(callback):
+        if callback:
+            try:
+                params = {
+                    'lr_max': callback.lr_max,
+                    'div_factor': callback.div_factor,
+                    'pct_start': callback.pct_start,
+                    'final_div': callback.final_div,
+                    'tot_epochs': callback.tot_epochs,
+                    'start_epoch': callback.start_epoch,
+                    'moms': callback.moms,
+                }
+                try_mlflow_log(mlflow.log_params, params)
+            except Exception:  # pylint: disable=W0703
+                return
+
+    def _run_and_log_function(self, original, args, kwargs, unlogged_params, callback_arg_index):
+        if not mlflow.active_run():
+            try_mlflow_log(mlflow.start_run)
+            auto_end_run = True
+        else:
+            auto_end_run = False
+
+        log_fn_args_as_params(original, [self] + list(args), kwargs, unlogged_params)
+
+        callbacks = [cb(self) for cb in self.callback_fns] + (self.callbacks or [])
+
+        # Checking if the 'callback' argument of the function is set
+        if len(args) > callback_arg_index:
+            tmp_list = list(args)
+            callbacks += list(args[callback_arg_index])
+            tmp_list[callback_arg_index] += [__MLflowFastaiCallback(self)]
+            args = tuple(tmp_list)
+        elif 'callbacks' in kwargs:
+            callbacks += list(kwargs['callbacks'])
+            kwargs['callbacks'] += [__MLflowFastaiCallback(self)]
+        else:
+            kwargs['callbacks'] = [__MLflowFastaiCallback(self)]
+
+        early_stop_callback = _find_callback_of_type(EarlyStoppingCallback, callbacks)
+        one_cycle_callback = _find_callback_of_type(OneCycleScheduler, callbacks)
+
+        _log_early_stop_callback_params(early_stop_callback)
+        _log_one_cycle_callback_params(one_cycle_callback)
+
+        result = original(self, *args, **kwargs)
+
+        if auto_end_run:
+            try_mlflow_log(mlflow.end_run)
+
+        return result
+
+    @gorilla.patch(Learner)
+    def fit(self, *args, **kwargs):
+        original = gorilla.get_original_attribute(Learner, 'fit')
+        unlogged_params = ['self', 'callbacks', 'learner']
+        return _run_and_log_function(self, original, args, kwargs, unlogged_params, 3)
+
+    settings = gorilla.Settings(allow_hit=True, store_hit=True)
+    gorilla.apply(gorilla.Patch(Learner, 'fit', fit, settings=settings))


Test patch:
diff --git a/tests/fastai/test_fastai_autolog.py b/tests/fastai/test_fastai_autolog.py
new file mode 100644
--- /dev/null
+++ b/tests/fastai/test_fastai_autolog.py
@@ -0,0 +1,248 @@
+import pytest
+import numpy as np
+from tests.conftest import tracking_uri_mock  # pylint: disable=W0611
+
+import pandas as pd
+import sklearn.datasets as datasets
+from fastai.tabular import tabular_learner, TabularList
+from fastai.metrics import accuracy
+import mlflow
+import mlflow.fastai
+from fastai.callbacks import EarlyStoppingCallback
+
+np.random.seed(1337)
+
+NUM_EPOCHS = 3
+MIN_DELTA = 99999999  # Forces earlystopping
+
+
+@pytest.fixture(params=[True, False])
+def manual_run(request, tracking_uri_mock):
+    if request.param:
+        mlflow.start_run()
+    yield
+    mlflow.end_run()
+
+
+def iris_dataframe():
+    iris = datasets.load_iris()
+    return pd.DataFrame(iris.data[:, :2], columns=iris.feature_names[:2])
+
+
+@pytest.fixture(scope="session")
+def iris_data():
+    iris = datasets.load_iris()
+    X = pd.DataFrame(iris.data[:, :2], columns=iris.feature_names[:2])
+    y = pd.Series(iris.target, name='label')
+    return (TabularList.from_df(pd.concat([X, y], axis=1), cont_names=list(X.columns))
+            .split_by_rand_pct(valid_pct=0.1, seed=42)
+            .label_from_df(cols='label')
+            .databunch())
+
+
+def fastai_model(data, **kwargs):
+    return tabular_learner(data, metrics=accuracy, layers=[5, 3, 2], **kwargs)
+
+
+@pytest.mark.large
+@pytest.mark.parametrize('fit_variant', ['fit', 'fit_one_cycle'])
+def test_fastai_autolog_ends_auto_created_run(iris_data, fit_variant):
+    mlflow.fastai.autolog()
+    model = fastai_model(iris_data)
+    if fit_variant == 'fit_one_cycle':
+        model.fit_one_cycle(1)
+    else:
+        model.fit(1)
+    assert mlflow.active_run() is None
+
+
+@pytest.mark.large
+@pytest.mark.parametrize('fit_variant', ['fit', 'fit_one_cycle'])
+def test_fastai_autolog_persists_manually_created_run(iris_data, fit_variant):
+    mlflow.fastai.autolog()
+
+    with mlflow.start_run() as run:
+        model = fastai_model(iris_data)
+
+        if fit_variant == 'fit_one_cycle':
+            model.fit_one_cycle(NUM_EPOCHS)
+        else:
+            model.fit(NUM_EPOCHS)
+
+        assert mlflow.active_run()
+        assert mlflow.active_run().info.run_id == run.info.run_id
+
+
+@pytest.fixture
+def fastai_random_data_run(iris_data, fit_variant, manual_run):
+    mlflow.fastai.autolog()
+
+    model = fastai_model(iris_data)
+
+    if fit_variant == 'fit_one_cycle':
+        model.fit_one_cycle(NUM_EPOCHS)
+    else:
+        model.fit(NUM_EPOCHS)
+
+    client = mlflow.tracking.MlflowClient()
+    return model, client.get_run(client.list_run_infos(experiment_id='0')[0].run_id)
+
+
+@pytest.mark.large
+@pytest.mark.parametrize('fit_variant', ['fit', 'fit_one_cycle'])
+def test_fastai_autolog_logs_expected_data(fastai_random_data_run, fit_variant):
+    model, run = fastai_random_data_run
+    data = run.data
+
+    # Testing metrics are logged
+    assert 'train_loss' in data.metrics
+    assert 'valid_loss' in data.metrics
+
+    for o in model.metrics:
+        assert o.__name__ in data.metrics
+
+    client = mlflow.tracking.MlflowClient()
+    metric_history = client.get_metric_history(run.info.run_id, 'valid_loss')
+    assert np.array_equal([m.value for m in metric_history], model.recorder.val_losses)
+
+    # Testing explicitly passed parameters are logged correctly
+    assert 'epochs' in data.params
+    assert data.params['epochs'] == str(NUM_EPOCHS)
+
+    # Testing implicitly passed parameters are logged correctly
+    assert 'wd' in data.params
+
+    # Testing unwanted parameters are not logged
+    assert 'callbacks' not in data.params
+    assert 'learn' not in data.params
+
+    # Testing optimizer parameters are logged
+    assert 'opt_func' in data.params
+    assert data.params['opt_func'] == 'Adam'
+    assert 'model_summary' in data.tags
+
+    # Testing model_summary.txt is saved
+    client = mlflow.tracking.MlflowClient()
+    artifacts = client.list_artifacts(run.info.run_id)
+    artifacts = map(lambda x: x.path, artifacts)
+    assert 'model_summary.txt' in artifacts
+
+
+@pytest.mark.large
+@pytest.mark.parametrize('fit_variant', ['fit', 'fit_one_cycle'])
+def test_fastai_autolog_logs_default_params(fastai_random_data_run, fit_variant):
+    _, run = fastai_random_data_run
+    if fit_variant == 'fit':
+        assert 'lr' in run.data.params
+        assert run.data.params['lr'] == 'slice(None, 0.003, None)'
+    else:
+        assert 'pct_start' in run.data.params
+        assert run.data.params['pct_start'] == '0.3'
+
+
+@pytest.mark.large
+@pytest.mark.parametrize('fit_variant', ['fit', 'fit_one_cycle'])
+def test_fastai_autolog_model_can_load_from_artifact(fastai_random_data_run):
+    run_id = fastai_random_data_run[1].info.run_id
+    client = mlflow.tracking.MlflowClient()
+    artifacts = client.list_artifacts(run_id)
+    artifacts = map(lambda x: x.path, artifacts)
+    assert 'model' in artifacts
+    model = mlflow.fastai.load_model("runs:/" + run_id + "/model")
+    model_wrapper = mlflow.fastai._FastaiModelWrapper(model)
+    model_wrapper.predict(iris_dataframe())
+
+
+@pytest.fixture
+def fastai_random_data_run_with_callback(iris_data, fit_variant, manual_run, callback, patience):
+    mlflow.fastai.autolog()
+    callbacks = []
+
+    if callback == 'early':
+        # min_delta is set as such to guarantee early stopping
+        callbacks.append(lambda learn: EarlyStoppingCallback(
+            learn, patience=patience, min_delta=MIN_DELTA))
+
+    model = fastai_model(iris_data, callback_fns=callbacks)
+
+    if fit_variant == 'fit_one_cycle':
+        model.fit_one_cycle(NUM_EPOCHS)
+    else:
+        model.fit(NUM_EPOCHS)
+
+    client = mlflow.tracking.MlflowClient()
+    return model, client.get_run(client.list_run_infos(experiment_id='0')[0].run_id)
+
+
+@pytest.mark.large
+@pytest.mark.parametrize('fit_variant', ['fit', 'fit_one_cycle'])
+@pytest.mark.parametrize('callback', ['early'])
+@pytest.mark.parametrize('patience', [0, 1, 5])
+def test_fastai_autolog_early_stop_logs(fastai_random_data_run_with_callback, patience):
+    model, run = fastai_random_data_run_with_callback
+    params = run.data.params
+    assert 'early_stop_patience' in params
+    assert params['early_stop_patience'] == str(patience)
+    assert 'early_stop_monitor' in params
+    assert params['early_stop_monitor'] == 'valid_loss'
+    assert 'early_stop_mode' in params
+    assert params['early_stop_mode'] == 'auto'
+    assert 'early_stop_min_delta' in params
+    assert params['early_stop_min_delta'] == f'-{MIN_DELTA}'
+
+    client = mlflow.tracking.MlflowClient()
+    metric_history = client.get_metric_history(run.info.run_id, 'valid_loss')
+    num_of_epochs = len(model.recorder.val_losses)
+
+    assert len(metric_history) == num_of_epochs
+
+
+@pytest.mark.large
+@pytest.mark.parametrize('fit_variant', ['fit', 'fit_one_cycle'])
+@pytest.mark.parametrize('callback', ['early'])
+@pytest.mark.parametrize('patience', [11])
+def test_fastai_autolog_early_stop_no_stop_does_not_log(
+        fastai_random_data_run_with_callback, patience):
+    model, run, = fastai_random_data_run_with_callback
+    metrics = run.data.metrics
+    params = run.data.params
+    assert 'early_stop_patience' in params
+    assert params['early_stop_patience'] == str(patience)
+    assert 'early_stop_monitor' in params
+    assert params['early_stop_monitor'] == 'valid_loss'
+    assert 'early_stop_mode' in params
+    assert 'early_stop_min_delta' in params
+    assert params['early_stop_min_delta'] == f'-{99999999}'
+    """
+    assert 'stopped_epoch' in metrics
+    assert metrics['stopped_epoch'] == 0
+    assert 'restored_epoch' not in metrics
+    """
+    num_of_epochs = len(model.recorder.val_losses)
+    client = mlflow.tracking.MlflowClient()
+    metric_history = client.get_metric_history(run.info.run_id, 'valid_loss')
+    # Check the test epoch numbers are correct
+    assert num_of_epochs == NUM_EPOCHS
+    assert len(metric_history) == num_of_epochs
+
+
+@pytest.mark.large
+@pytest.mark.parametrize('fit_variant', ['fit', 'fit_one_cycle'])
+@pytest.mark.parametrize('callback', ['not-early'])
+@pytest.mark.parametrize('patience', [5])
+def test_fastai_autolog_non_early_stop_callback_does_not_log(fastai_random_data_run_with_callback):
+    model, run, = fastai_random_data_run_with_callback
+    metrics = run.data.metrics
+    params = run.data.params
+    assert 'early_stop_patience' not in params
+    assert 'early_stop_monitor' not in params
+    assert 'early_stop_mode' not in params
+    assert 'stopped_epoch' not in metrics
+    assert 'restored_epoch' not in metrics
+    assert 'early_stop_min_delta' not in params
+    num_of_epochs = len(model.recorder.val_losses)
+    client = mlflow.tracking.MlflowClient()
+    metric_history = client.get_metric_history(run.info.run_id, 'valid_loss')
+    # Check the test epoch numbers are correct
+    assert num_of_epochs == NUM_EPOCHS
+    assert len(metric_history) == num_of_epochs
diff --git a/tests/fastai/test_fastai_model_export.py b/tests/fastai/test_fastai_model_export.py
--- a/tests/fastai/test_fastai_model_export.py
+++ b/tests/fastai/test_fastai_model_export.py
@@ -26,6 +26,7 @@
 from tests.helper_functions import mock_s3_bucket  # pylint: disable=unused-import
 from tests.helper_functions import score_model_in_sagemaker_docker_container
 
+
 ModelWithData = namedtuple("ModelWithData", ["model", "inference_dataframe"])