Implementing Multiple Profiles for Dynamic Scaling of Neural Network

I was trying to replicate the results in the paper titled “Incomplete Dot Products for Dynamic Computation Scaling in Neural Network Inference” using Tensorflow, but couldn’t. After several emails to the authors, one of them had managed to reply and sent me the code, which was written in Chainer. All efforts to understand the code and convert it to Tensorflow or Keras proved abortive as I couldn’t make sense out of many lines in the code. However, I managed to reproduce the first part of the results (Single Profile). My biggest challenge is how to implement the Multiple Profile. My research direction is heavily dependent on my ability to reproduce the results in the paper mentioned above. I will appreciate if you could help me because I have been in it for over five months and still couldn’t figure out how to do it. Thank you. Below is my code, and the one written in Chainer by the authors can be found in https://github.com/kunglab/idp

import numpy as np
import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data
import math
import matplotlib.pyplot as plt
from tensorflow.python.framework import ops
mnist = input_data.read_data_sets("/tmp/data/", one_hot=True)

Network Parameters

n_1 = 100               # 1st layer number of neurons
n_2 = 100               # 2nd layer number of neurons
n_input = 784           #MNIST data input (img shape: 28*28)
n_classes = 10          # MNIST total classes (0-9 digits)
learning_rate = 0.0008
training_epochs = 20
batch_size = 30
display_step = 1

np.random.seed(1)

tf Graph input

tf.reset_default_graph()
X = tf.placeholder(tf.float32, [None, n_input])
Y = tf.placeholder(tf.float32, [None, n_classes])

Store layers weight & bias

def initialize_param(n_input, n_1, n_2, n_class):
    tf.set_random_seed(1)
    W1 = tf.get_variable(“W1”, shape = [n_input, n_1], initializer = tf.contrib.layers.xavier_initializer(seed = 1))
    b1 = tf.get_variable(“b1”, shape = [n_1], initializer = tf.zeros_initializer())
    W2 = tf.get_variable(“W2”, shape = [n_1, n_2], initializer = tf.contrib.layers.xavier_initializer(seed = 1))
    b2 = tf.get_variable(“b2”, shape = [n_2], initializer = tf.zeros_initializer())
    W3 = tf.get_variable(“W3”, shape = [n_2, n_class], initializer = tf.contrib.layers.xavier_initializer(seed = 1))
    b3 = tf.get_variable(“b3”, shape = [n_class], initializer = tf.zeros_initializer())

    parameters = {“W1”: W1,“b1”: b1,“W2”: W2,“b2”: b2,“W3”: W3,“b3”: b3}
    return parameters
parameters = initialize_param(784, 100, 100, 10)

def linear_func(n):
    return[np.float32(1.0 - 1.0 * i/n) for i in range(1, n + 1)]
L = linear_func(100)

def linear_profile(lp, n_1):
    p_L = tf.constant(L, shape = [1, 100])
    L_11 = tf.constant(1.0, shape = [1, int(np.round((lp) * n_1))])
    L_12 = tf.zeros(shape = [1, int(np.round((1 - lp) * n_1))])
    L1 = tf.concat((L_11, L_12), axis = 1)
    p_L1 = tf.multiply(L1, p_L)
    return p_L1

Creating Multiple Profile

pc1 = np.linspace(0, 0.2, 10)
pc2 = np.linspace(0.2, 0.4, 10)
pc3 = np.linspace(0.4, 1.0, 10)

profile_1 = []
profile_2 = []
profile_3 = []

for j in pc1:
    p_L1 = linear_profile(j, 100)
    profile_11 = tf.stack(p_L1, axis = 0)
    profile_1.append(profile_11)
    profile1 = tf.convert_to_tensor(profile_1, dtype=tf.float32)

for j in pc2:
    p_L1 = linear_profile(j, 100)
    profile_22 = tf.stack(p_L1, axis = 0)
    profile_2.append(profile_22)
    profile2 = tf.convert_to_tensor(profile_2, dtype=tf.float32)

for j in pc3:
    p_L1 = linear_profile(j, 100)
    profile_33 = tf.stack(p_L1, axis = 0)
    profile_3.append(profile_33)
    profile3 = tf.convert_to_tensor(profile_3, dtype=tf.float32)
    
def mlp_1(x, profile_type):
    logist_t = []
    for j in range(len(pc1)):
        Z_ML11 = tf.add(tf.matmul(x, parameters[‘W1’]), parameters[‘b1’])  
        A_ML11 = tf.nn.relu(Z_ML11)
        P_ML11 = tf.multiply(profile_type[j], A_ML11)
        Z_ML12 = tf.add(tf.matmul(P_ML11, parameters[‘W2’]), parameters[‘b2’])  
        A_ML12 = tf.nn.relu(Z_ML12)
        P_ML12 = tf.multiply(profile_type[j], A_ML12)
        out_layer = tf.add(tf.matmul(P_ML12, parameters[‘W3’]), parameters[‘b3’])
        logist_t.append(out_layer)
    return logist_t
logist_1 = mlp_1(X, profile_1)
logist_2 = mlp_1(X, profile_2)
logist_3 = mlp_1(X, profile_3)

def optimize_param(logits):
    for j in range(len(pc1)):
        loss_op[j] = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits = logits[j], labels = Y))
        optimizer = tf.train.MomentumOptimizer(learning_rate = learning_rate, momentum = 0.98).minimize(loss_op)
    return loss_op, optimizer

loss_op_1,optimizer_1 = optimize_param(logist_1)
loss_op_2,optimizer_2 = optimize_param(logist_2)
loss_op_3,optimizer_3 = optimize_param(logist_3)

Initializing the variables

init = tf.global_variables_initializer()
with tf.Session() as sess:
    sess.run(init)
    
    # Training Loop
    cost_1 = []
    for epoch in range(training_epochs):
        avg_cost1 = 0.
        total_batch = int(mnist.train.num_examples/batch_size)
                
        # Loop over all batches
        for i in range(total_batch):
            batch_x, batch_y = mnist.train.next_batch(batch_size)
            # Run optimization op (backprop) and cost op (to get loss value)
            _, c_1 = sess.run([loss_op_1, optimizer_1], feed_dict = {X: batch_x, Y: batch_y})
                    
            # Compute average losses
            avg_cost1 += c_1 / total_batch
            cost_1.append(avg_cost1)
                
            train_accuracy_11 = []
            test_accuracy_11 = []
            for k in range(len(logist_1)):
                if i % 5000 == 0:
                    pred_1 = tf.nn.softmax(logits_1[k])  # Apply softmax to logits
                    correct_prediction_1 = tf.equal(tf.argmax(pred_1, 1), tf.argmax(Y, 1))
                    accuracy_1 = tf.reduce_mean(tf.cast(correct_prediction_1, “float”))
                    trian_accuracy_1 = accuracy_1.eval({X: mnist.train.images, Y: mnist.train.labels})
                    test_accuracy_1 = accuracy_1.eval({X: mnist.test.images, Y: mnist.test.labels})*100
                    train_accuracy_11.append(train_accuracy_1)
                    test_accuracy_11.append(test_accuracy_1)
                        
            # Display logs per epoch step
            for k in range(len(logist_1)):
                if epoch % display_step == 0:
                    print(“Epoch_1:”, ‘%03d’ % (epoch + 1), “cost = {:.9f}”.format(avg_cost1[k]))
            
        for l in pc1:
            if l % display_step == 0:
                print(“Percentage_Profile_1:”, “Test Accuracy_1 = {}”.format(test_accuracy_11))  
                
sess.close()