keras.py

import keras
import keras.layers as L
import keras.regularizers as R
import keras.callbacks as C
from keras.engine import Model, Input
from alea import dataset

import tensorflow as tf


def gen_session(gpu_id=None):
  config = tf.ConfigProto()
  config.gpu_options.allow_growth = True
  if gpu_id is not None:
    config.gpu_options.visible_device_list = f"{gpu_id}"
  sess = tf.Session(config=config)
  return sess

def gen_model():
  x = Input([28, 28, 1])
  reg = R.l2(0.005)

  h = x
  h = L.Conv2D(
          filters=16, kernel_size=(3, 3), strides=(2, 2),
          padding='same', activation='relu',
          kernel_regularizer=reg,
          )(h)
  h = L.Conv2D(
          filters=32, kernel_size=(5, 5), strides=(3, 3),
          padding='same', activation='relu',
          kernel_regularizer=reg,
          )(h)

  h = L.Flatten()(h)
  h = L.Dense(1024, activation='relu')(h)
  h = L.Dense(10, activation='softmax')(h)

  model = Model(inputs=[x], outputs=[h])
  return model

def main():
    sess = gen_session(gpu_id=3)

    dataset = dataset.MnistDataset()

    npz = dataset.load()
    images_train = npz['images_train']
    images_test = npz['images_test']
    labels_train = npz['labels_train']
    labels_test = npz['labels_test']

    model = gen_model()

    m.compile(
        optimizer='adam',
        loss='sparse_categorical_crossentropy',
        metrics=['accuracy'],
    )

    m.summary()

    m.fit(
      images_train, labels_train,
      epochs=100,
      validation_data=[images_test, labels_test],
      batch_size=128,
      callbacks=[
        C.CSVLogger('log.csv'),
        C.EarlyStopping(patience=10),
      ],
    )



if __name__ == '__main__':
    main()