Train a ready to use TensorFlow model with a simple pipeline


In [1]:
import os
import sys
import warnings
warnings.filterwarnings("ignore")

import numpy as np
import matplotlib.pyplot as plt

# the following line is not required if BatchFlow is installed as a python package.
sys.path.append("../..")
from batchflow import Pipeline, B, C, D, F, V
from batchflow.opensets import MNIST, CIFAR10, CIFAR100
from batchflow.models.tf import ResNet18

BATCH_SIZE might be increased for modern GPUs with lots of memory (4GB and higher).


In [2]:
BATCH_SIZE = 64

Create a dataset

MNIST is a dataset of handwritten digits frequently used as a baseline for machine learning tasks.

Downloading MNIST database might take a few minutes to complete.


In [3]:
dataset = MNIST(bar=True)


100%|██████████| 8/8 [00:08<00:00,  2.57s/it]

There are also predefined CIFAR10 and CIFAR100 datasets.

Define a pipeline config

Config allows to create flexible pipelines which take parameters.

For instance, if you put a model type into config, you can run a pipeline against different models.

See a list of available models to choose the one which fits you best.


In [4]:
config = dict(model=ResNet18)

Create a template pipeline

A template pipeline is not linked to any dataset. It's just an abstract sequence of actions, so it cannot be executed, but it serves as a convenient building block.


In [5]:
train_template = (Pipeline()
                .init_variable('loss_history', [])
                .init_model('dynamic', C('model'), 'conv_nn',
                            config={'inputs/images/shape': B.image_shape,
                                    'inputs/labels/classes': D.num_classes,
                                    'initial_block/inputs': 'images'})
                .to_array()
                .train_model('conv_nn', fetches='loss', images=B.images, labels=B.labels,
                             save_to=V('loss_history', mode='a'))
)

Train the model

Apply a dataset and a config to a template pipeline to create a runnable pipeline:


In [6]:
train_pipeline = (train_template << dataset.train) << config

Run the pipeline (it might take from a few minutes to a few hours depending on your hardware)


In [7]:
train_pipeline.run(BATCH_SIZE, shuffle=True, n_epochs=1, drop_last=True, bar=True, prefetch=1)


100%|██████████| 937/937 [16:51<00:00,  1.08s/it]
Out[7]:
<batchflow.pipeline.Pipeline at 0x7f5d0419a2e8>

Note that the progress bar often increments by 2 at a time - that's prefetch in action.

It does not give much here, though, since almost all time is spent in model training which is performed under a thread-lock one batch after another without any parallelism (otherwise the model would not learn anything as different batches would rewrite one another's model weights updates).


In [8]:
plt.figure(figsize=(15, 5))
plt.plot(train_pipeline.v('loss_history'))
plt.xlabel("Iterations"), plt.ylabel("Loss")
plt.show()


Test the model

It is much faster than training, but if you don't have GPU it would take some patience.


In [9]:
test_pipeline = (dataset.test.p
                .import_model('conv_nn', train_pipeline)
                .init_variable('predictions') 
                .init_variable('metrics') 
                .to_array()
                .predict_model('conv_nn', fetches='predictions', images=B.images, save_to=V('predictions'))
                .gather_metrics('class', targets=B.labels, predictions=V('predictions'),
                                fmt='logits', axis=-1, save_to=V('metrics', mode='a'))
                .run(BATCH_SIZE, shuffle=True, n_epochs=1, drop_last=False, bar=True)
)


 99%|█████████▉| 156/157 [00:08<00:00, 19.16it/s]

Let's get the accumulated metrics information


In [10]:
metrics = test_pipeline.get_variable('metrics')

Or a shorter version: metrics = test_pipeline.v('metrics')

Now we can easiliy calculate any metrics we need


In [11]:
metrics.evaluate('accuracy')


Out[11]:
0.9773089171974523

In [12]:
metrics.evaluate(['false_positive_rate', 'false_negative_rate'], multiclass=None)


Out[12]:
{'false_negative_rate': array([0.01185421, 0.02188146, 0.00161763, 0.01733417, 0.08088979,
        0.01340107, 0.01051461, 0.03377313, 0.01964688, 0.01724571]),
 'false_positive_rate': array([0.00088685, 0.00022548, 0.00677109, 0.00087906, 0.        ,
        0.00064096, 0.00234634, 0.00165906, 0.00443706, 0.00756822])}

Save the model

After learning the model, you may need to save it. It's easy to do this.


In [13]:
train_pipeline.save_model_now('conv_nn', path='path/to/save')

What's next?