train_mnist_conv

#!/usr/bin/env python3

import os
import sys
import time
import random
import scipy.ndimage
import numpy as np
import tensorflow as tf
import tensorflow.examples.tutorials.mnist as tf_mnist

import loadsave
from utils import *

def dense(x, out_dim):
    with tf.name_scope("dense"):
        in_dim = np.prod([
            1 if v is None else v
            for v in x.shape.as_list()
        ])
        x = tf.reshape(x, [-1, in_dim])
        b = tf.Variable(tf.constant(0.1, shape = [out_dim]))
        w = tf.Variable(tf.truncated_normal(
            [in_dim, out_dim],
            stddev = 0.1
        ))
        return tf.matmul(x, w) + b

def conv(x, ksize, filters, stride = 1, padding = "SAME"):
    with tf.name_scope("conv"):
        in_channels = x.shape[-1].value
        b = tf.Variable(tf.constant(0.1, shape = [filters]))
        w = tf.Variable(tf.truncated_normal(
            [ksize, ksize, in_channels, filters],
            stddev = 0.1
        ))
        return tf.nn.conv2d(
            x, w,
            strides = [1, stride, stride, 1],
            padding = padding
        ) + b

def maxpool(x, size, padding = "SAME"):
    with tf.name_scope("maxpool"):
        return tf.nn.max_pool(
            x,
            ksize = [1, size, size, 1],
            strides = [1, size, size, 1],
            padding = padding
        )

def trainable_parameters():
    return sum([
        np.prod(v.shape.as_list())
        for v in tf.trainable_variables()
    ])

def create_layers(inp, layers):
    params = trainable_parameters()
    print_info(
        "%30s Input:  %d params"
        % (inp.shape, params)
    )
    y = inp
    for i in range(len(layers)):
        with tf.name_scope("layer_%d" % (i+1)):
            y = layers[i](y)
        new_params = trainable_parameters() - params
        print_info(
            "%30s Layer %d: %d params"
            % (y.shape, i+1, new_params)
        )
        params += new_params
    return y

def create_model():
    x = tf.placeholder(tf.float32, [None, 28, 28], name = "x")
    t = tf.placeholder(tf.float32, [None, 10], name = "t")
    dropout = tf.placeholder_with_default(0.0, [], name = "dropout")

    layers = [
        lambda l: tf.nn.relu(conv(l, 3, 32)),
        lambda l: tf.nn.relu(conv(l, 3, 32)),
        lambda l: tf.nn.dropout(l, 1.0 - dropout),
        lambda l: maxpool(l, 2),
        lambda l: tf.nn.relu(conv(l, 3, 64)),
        lambda l: tf.nn.relu(conv(l, 3, 64)),
        lambda l: tf.nn.dropout(l, 1.0 - dropout),
        lambda l: maxpool(l, 2),
        lambda l: tf.nn.relu(conv(l, 3, 64)),
        lambda l: tf.nn.relu(conv(l, 3, 64)),
        lambda l: tf.nn.dropout(l, 1.0 - dropout),
        lambda l: maxpool(l, 2),
        lambda l: tf.nn.relu(dense(l, 1024)),
        lambda l: tf.nn.dropout(l, 1.0 - dropout),
        lambda l: dense(l, 10)
    ]

    x = tf.clip_by_value(x * 1.1 - 0.1, 0.0, 1.0)
    x = tf.reshape(x, [-1, 28, 28, 1])
    y = create_layers(x, layers)

    correct = tf.equal(tf.argmax(y, 1), tf.argmax(t, 1))
    correct = tf.cast(correct, tf.float32)

    tf.cast(tf.reduce_sum(1.0 - correct), tf.int64, name = "mistakes")
    tf.reduce_mean(correct, name = "accuracy")

    loss = tf.reduce_mean(
        tf.nn.softmax_cross_entropy_with_logits(labels=t, logits=y),
        name = "loss"
    )

    y = tf.nn.softmax(y, 1, name = "y")

    global_step = tf.contrib.framework.get_or_create_global_step()

    tf.train.AdamOptimizer(0.0001).minimize(
        loss,
        global_step = global_step,
        name = "train"
    )

def compute_stats(s, data):
    accuracy, loss, mistakes = 0, 0, 0
    for x, t in random_batches(data, 1000):
        v = s.run(
            [s.accuracy, s.loss, s.mistakes],
            feed_dict = {s.x: x, s.t: t}
        )
        accuracy += v[0] * len(x)
        loss += v[1] * len(x)
        mistakes += v[2]
    return {
        "accuracy": accuracy / len(data[0]),
        "loss": loss / len(data[0]),
        "mistakes": mistakes,
        "samples": len(data[0])
    }

def format_stats(v):
    return "%5.2f%% (%d/%d), loss=%f" % (
        v["accuracy"] * 100.0,
        v["samples"] - v["mistakes"],
        v["samples"],
        v["loss"]
    )

def add_scar(i, size):
    y = random.randrange(i.shape[0] - size + 1)
    x = random.randrange(i.shape[1] - size + 1)
    i[y:y+size,x:x+size] = 0.0

def augment_image(i):
    i = i.reshape((28, 28)).copy()
    add_scar(i, 3)
    if np.random.rand() < 0.5:
        add_scar(i, 4)
        i = scipy.ndimage.rotate(
            i, 20.0 * (np.random.rand() - 0.5),
            reshape = False
        )
    if np.random.rand() < 0.2:
        i = i * 3 - np.random.rand() * 2
    if np.random.rand() < 0.1:
        add_scar(i, 5)
    return i

def augmented_batches(data, batch_size):
    for b in random_batches(data, batch_size):
        yield (np.array([augment_image(i) for i in b[0]]), b[1])

def train_batch(s, batch):
    t = time.time()
    s.run(
        s.train,
        feed_dict = {s.x: batch[0], s.t: batch[1], s.dropout: 0.5}
    )
    return time.time() - t

def train(s, train, valid, program):
    epoch = 0
    next_target = 0.99

    while True:
        epoch += 1
        start_time = time.time()
        calc_time = 0.0

        for b in iterate_in_thread(augmented_batches(train, 100)):
            calc_time += train_batch(s, b)

        epoch_time = time.time() - start_time
        s_train = compute_stats(s, next(augmented_batches(train, 5000)))
        s_valid = compute_stats(s, valid)

        stat = "epoch %-5d train: %s   valid: %s" % (epoch,
            format_stats(s_train), format_stats(s_valid))

        time_stat = "[C+%.1fs/X+%.1fs/S+%.1fs]" % (
            calc_time,
            epoch_time - calc_time,
            time.time() - start_time - epoch_time
        )

        print_info("%-30s %s" % (time_stat, stat))

        accuracy = s_valid["accuracy"]
        if accuracy >= next_target:
            next_target = accuracy * 0.9 + 0.1
            loadsave.save(s, sys.argv[1])
            save_result(program, stat)

        if accuracy >= 0.99985:
            return

def run():
    if len(sys.argv) < 2:
        sys.stderr.write("\nUsage:\n\n");
        sys.stderr.write("\ttrain_mnist_conv <output.model>\n\n")
        sys.exit(1)

    with open(sys.argv[0], "r") as f:
        program = f.read()

    mnist = tf_mnist.input_data.read_data_sets(
        "__mnist__",
        one_hot = True
    )

    with loadsave.load(sys.argv[1], create_model = create_model) as s:
        print_info("Training model...")
        train(
            s,
            (mnist.train.images.reshape((-1, 28, 28)),
                mnist.train.labels),
            (mnist.validation.images.reshape((-1, 28, 28)),
                mnist.validation.labels),
            program
        )

if __name__ == "__main__":
    run()