Build time measurements

lstm.py

#!/usr/bin/python3
'''
Running: "./lstm.py <num_timesteps>"
'''
 
import tvm
import topi
import numpy as np
import time
import sys
 
# Training Parameters
batch_size = 64
 
# Network Parameters
num_timesteps = int(sys.argv[1])
num_input = 32
num_hidden = 128
num_classes = 10
 
# Weights and biases sizes
sizes = [
    (num_hidden, num_input + num_hidden),
    (num_hidden,),
    (num_hidden, num_input + num_hidden),
    (num_hidden,),
    (num_hidden, num_input + num_hidden),
    (num_hidden,),
    (num_hidden, num_input + num_hidden),
    (num_hidden,),
    (num_classes, num_hidden),
    (num_classes,)
]
 
# Structure for weights and biases initialization
inits = [
    (np.zeros, 'shape'),
    (np.zeros, 'shape'),
    (np.zeros, 'shape'),
    (np.zeros, 'shape'),
    (np.zeros, 'shape'),
    (np.ones, 'shape'),
    (np.zeros, 'shape'),
    (np.zeros, 'shape'),
    (np.random.normal, 'size'),
    (np.random.normal, 'size')
]
 
# Graph input
x = tvm.placeholder((batch_size, num_timesteps * num_input), 'float32')
y = tvm.placeholder((batch_size, num_classes), 'float32')
s = tvm.placeholder((batch_size, num_hidden), 'float32')
h = tvm.placeholder((batch_size, num_hidden), 'float32')
 
# Tensors and vars for training graph
weights = [tvm.placeholder(x, 'float32') for x in sizes]
 
#Construct model
xs = topi.split(topi.reshape(x, (batch_size, num_timesteps, num_input)), num_timesteps, axis=1)
xs = [topi.reshape(x, (batch_size, num_input)) for x in xs]
new_s = s
new_h = h
 
for i in range(num_timesteps):
    inp = topi.concatenate([xs[i], new_h], 1)
    g = topi.tanh(topi.nn.dense(inp, weights[0], weights[1]))
    j = topi.sigmoid(topi.nn.dense(inp, weights[2], weights[3]))
    f = topi.sigmoid(topi.nn.dense(inp, weights[4], weights[5]))
    o = topi.sigmoid(topi.nn.dense(inp, weights[6], weights[7]))
 
    new_s = new_s * f + g * j
    new_h = topi.tanh(new_s) * o
 
logits = topi.nn.dense(new_h, weights[8], weights[9])
 
# Define loss
loss = topi.sum(-topi.sum(y * topi.nn.log_softmax(logits), axis=1)) / batch_size
 
# Define model
sched = tvm.create_schedule([loss.op])
 
f = open('lstm.txt', 'a')
start_time = time.time()
model = tvm.build(sched, [x, y, s, h, loss, *weights])
print(num_timesteps, time.time() - start_time, file=f)
f.close()