tf_cifar10_dls.py

import tensorflow as tf

__weights_dict = dict()

is_train = False

def load_weights(weight_file):
    import numpy as np

    if weight_file == None:
        return

    try:
        weights_dict = np.load(weight_file).item()
    except:
        weights_dict = np.load(weight_file, encoding='bytes').item()

    return weights_dict


def KitModel(weight_file = None):
    global __weights_dict
    __weights_dict = load_weights(weight_file)

    Input_0         = tf.placeholder(tf.float32,  shape = (None, 3, 32, 32), name = 'Input_0')
    convolution2d_1 = convolution(Input_0, group=1, strides=[1, 1], padding='SAME', name='convolution2d_1')
    activation_1    = tf.nn.relu(convolution2d_1, name = 'activation_1')
    convolution2d_2 = convolution(activation_1, group=1, strides=[1, 1], padding='VALID', name='convolution2d_2')
    activation_2    = tf.nn.relu(convolution2d_2, name = 'activation_2')
    maxpooling2d_1  = tf.nn.max_pool(activation_2, [1, 2, 2, 1], [1, 2, 2, 1], padding='VALID', name='maxpooling2d_1')
    convolution2d_3 = convolution(maxpooling2d_1, group=1, strides=[1, 1], padding='SAME', name='convolution2d_3')
    activation_3    = tf.nn.relu(convolution2d_3, name = 'activation_3')
    convolution2d_4 = convolution(activation_3, group=1, strides=[1, 1], padding='VALID', name='convolution2d_4')
    activation_4    = tf.nn.relu(convolution2d_4, name = 'activation_4')
    maxpooling2d_2  = tf.nn.max_pool(activation_4, [1, 2, 2, 1], [1, 2, 2, 1], padding='VALID', name='maxpooling2d_2')
    flatten_1       = tf.contrib.layers.flatten(maxpooling2d_2)
    dense_1         = tf.layers.dense(flatten_1, 512, kernel_initializer = tf.constant_initializer(__weights_dict['dense_1']['weights']), bias_initializer = tf.constant_initializer(__weights_dict['dense_1']['bias']), use_bias = True)
    activation_5    = tf.nn.relu(dense_1, name = 'activation_5')
    dense_2         = tf.layers.dense(activation_5, 10, kernel_initializer = tf.constant_initializer(__weights_dict['dense_2']['weights']), bias_initializer = tf.constant_initializer(__weights_dict['dense_2']['bias']), use_bias = True)
    activation_6    = tf.nn.softmax(dense_2, name = 'activation_6')
    return Input_0, activation_6


def convolution(input, name, group, **kwargs):
    w = tf.Variable(__weights_dict[name]['weights'], trainable=is_train, name=name + "_weight")
    if group == 1:
        layer = tf.nn.convolution(input, w, name=name, **kwargs)
    else:
        weight_groups = tf.split(w, num_or_size_splits=group, axis=-1)
        xs = tf.split(input, num_or_size_splits=group, axis=-1)
        convolved = [tf.nn.convolution(x, weight, name=name, **kwargs) for
                    (x, weight) in zip(xs, weight_groups)]
        layer = tf.concat(convolved, axis=-1)

    if 'bias' in __weights_dict[name]:
        b = tf.Variable(__weights_dict[name]['bias'], trainable=is_train, name=name + "_bias")
        layer = layer + b
    return layer