Train Baseline Model - 1.4.1 English

Vitis AI Optimizer User Guide (UG1333)

Document ID
UG1333
Release Date
2021-10-29
Version
1.4.1 English

To train the model by creating an Estimator and calling train() on it, create a file named train.py and add following codes:

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from est_cnn import model_fn, train_input_fn, eval_input_fn

# Imports
import numpy as np
import tensorflow as tf

tf.logging.set_verbosity(tf.logging.INFO)

def main(unused_argv):
  mnist_classifier = model_fn()

  mnist_classifier.train(
      input_fn=train_input_fn(),
      max_steps=20000)

  eval_results = mnist_classifier.evaluate(input_fn=eval_input_fn())
  print(eval_results)

if __name__ == "__main__":
  tf.app.run()

Run train.py.

$ python train.py

As the model trains, an output similar to the following is displayed:

INFO:tensorflow:Running local_init_op.
INFO:tensorflow:Done running local_init_op.
INFO:tensorflow:Saving checkpoints for 0 into ./models/train/model.ckpt.
INFO:tensorflow:loss = 2.294087, step = 0
INFO:tensorflow:global_step/sec: 201.741
INFO:tensorflow:loss = 2.2876544, step = 100 (0.496 sec)
INFO:tensorflow:global_step/sec: 228.126
INFO:tensorflow:loss = 2.2656975, step = 200 (0.439 sec)
INFO:tensorflow:global_step/sec: 225.094
INFO:tensorflow:loss = 2.2483034, step = 300 (0.444 sec)
INFO:tensorflow:global_step/sec: 234.019
…
INFO:tensorflow:Saving dict for global step 20000: accuracy = 0.9684, global_step = 20000, loss = 0.10172604
INFO:tensorflow:Saving 'checkpoint_path' summary for global step 20000: ./models/train/model.ckpt-20000
{'accuracy': 0.9684, 'loss': 0.10172604, 'global_step': 20000}

You can get an accuracy of 96.84% on our test data set.