main.py

# Copyright 2017 Abien Fred Agarap

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#    http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Implementation of the CNN classes"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

__version__ = '0.1.0'
__author__ = 'Abien Fred Agarap'

import argparse
from model.cnn_softmax import CNN
from model.cnn_svm import CNNSVM
from tensorflow.examples.tutorials.mnist import input_data


def parse_args():
    parser = argparse.ArgumentParser(description='CNN & CNN-SVM for Image Classification')
    group = parser.add_argument_group('Arguments')
    group.add_argument('-m', '--model', required=True, type=str,
                       help='[1] CNN-Softmax, [2] CNN-SVM')
    group.add_argument('-d', '--dataset', required=True, type=str,
                       help='path of the MNIST dataset')
    group.add_argument('-p', '--penalty_parameter', required=False, type=int,
                       help='the SVM C penalty parameter')
    group.add_argument('-c', '--checkpoint_path', required=True, type=str,
                       help='path where to save the trained model')
    group.add_argument('-l', '--log_path', required=True, type=str,
                       help='path where to save the TensorBoard logs')
    arguments = parser.parse_args()
    return arguments


if __name__ == '__main__':
    args = parse_args()

    mnist = input_data.read_data_sets(args.dataset, one_hot=True)
    num_classes = mnist.train.labels.shape[1]
    sequence_length = mnist.train.images.shape[1]
    model_choice = args.model

    assert model_choice == '1' or model_choice == '2', "Invalid choice: Choose between 1 and 2 only."

    if model_choice == '1':
        model = CNN(alpha=1e-3, batch_size=128, num_classes=num_classes, num_features=sequence_length)
        model.train(checkpoint_path=args.checkpoint_path, epochs=10000, log_path=args.log_path,
                    train_data=mnist.train, test_data=mnist.test)
    elif model_choice == '2':
        model = CNNSVM(alpha=1e-3, batch_size=128, num_classes=num_classes, num_features=sequence_length,
                       penalty_parameter=args.penalty_parameter)
        model.train(checkpoint_path=args.checkpoint_path, epochs=10000, log_path=args.log_path,
                    train_data=mnist.train, test_data=mnist.test)