TensorFlow 2.0 CIFAR-10 Sample

train_eval.py

import tensorflow as tf
import numpy as np

import os

from tensorflow.keras.layers import Dense, Flatten, Conv2D
from tensorflow.keras import Model
print(tf.__version__)

DATA_DIR = './data'

train_list = filter(lambda f: f.startswith('data_') and f.endswith('.bin'), os.listdir(DATA_DIR))
train_path_list = list(map(lambda f: os.path.join(DATA_DIR, f), train_list))

test_list = filter(lambda f: f.startswith('test_') and f.endswith('.bin'), os.listdir(DATA_DIR))
test_path_list = list(map(lambda f: os.path.join(DATA_DIR, f), test_list))

print(train_path_list)
print(test_path_list)


LABEL_BYTE = 1
IMAGE_BYTES = 32 * 32 * 3
RECORD_BYTES = LABEL_BYTE + IMAGE_BYTES


def _load_dataset(data_path_list):
    def _process_record(record):
        value = tf.io.decode_raw(record, tf.uint8)
        label = value[0]
        image = value[1:]
        image = tf.reshape(image, (3, 32, 32))
        image = tf.transpose(image, (1, 2, 0))
        image = tf.cast(image, tf.float32)
        image = image / 255
        return image, label

    dataset = tf.data.FixedLengthRecordDataset(
        data_path_list,
        RECORD_BYTES)
    return dataset.map(_process_record)

class MyModel(Model):
    def __init__(self):
        super(MyModel, self).__init__()
        self.conv1 = Conv2D(32, 3, activation='relu')
        self.flatten = Flatten()
        self.d1 = Dense(128, activation='relu')
        self.d2 = Dense(10, activation='softmax')

    def call(self, x):
        x = self.conv1(x)
        x = self.flatten(x)
        x = self.d1(x)
        return self.d2(x)

# Create an instance of the model
model = MyModel()

loss_object = tf.keras.losses.SparseCategoricalCrossentropy()
optimizer = tf.keras.optimizers.Adam()

@tf.function
def train_step(images, labels):
    with tf.GradientTape() as tape:
        predictions = model(images)
        loss = loss_object(labels, predictions)
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))

test_loss = tf.keras.metrics.Mean(name='test_loss')
test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='test_accuracy')

@tf.function
def test_step(images, labels):
    predictions = model(images)
    t_loss = loss_object(labels, predictions)

    test_loss(t_loss)
    test_accuracy(labels, predictions)

EPOCHS = 5
SUMMARY_DIR = './summary'

TRAIN_BATCH_SIZE = 32
TEST_BATCH_SIZE = 32

import time

train_dataset = _load_dataset(train_path_list).batch(TRAIN_BATCH_SIZE)
test_dataset = _load_dataset(test_path_list).batch(TEST_BATCH_SIZE)

summary_writer = tf.summary.create_file_writer(SUMMARY_DIR)

for epoch in range(EPOCHS):
    start = time.time()
    for images, labels in train_dataset:
        train_step(images, labels)

    for test_images, test_labels in test_dataset:
        test_step(test_images, test_labels)

    elapsed = time.time() - start
    print('elapsed: %f' % elapsed)

    template = 'Epoch {}, Test Loss: {}, Test Accuracy: {}'
    print(template.format(epoch+1,
                        test_loss.result(),
                        test_accuracy.result()*100))

    # Reset the metrics for the next epoch
    test_loss.reset_states()
    test_accuracy.reset_states()

print('Training Finished.')