def __init__(self,model_type,model_size,learning_rate = 0.001):
self.model_type = model_type
self.constrain_op = None
self.x = tf.placeholder(dtype=tf.float32,shape=[None,None,6])
self.target_y = tf.placeholder(dtype=tf.float32,shape=[None,None,1])
self.model_size = model_size
head = self.x
if(model_type == "lstm"):
self.fused_cell = tf.nn.rnn_cell.LSTMCell(model_size)
head,_ = tf.nn.dynamic_rnn(self.fused_cell,head,dtype=tf.float32,time_major=True)
elif(model_type.startswith("ltc")):
learning_rate = 0.01 # LTC needs a higher learning rate
self.wm = ltc.LTCCell(model_size)
if(model_type.endswith("_rk")):
self.wm._solver = ltc.ODESolver.RungeKutta
elif(model_type.endswith("_ex")):
self.wm._solver = ltc.ODESolver.Explicit
else:
self.wm._solver = ltc.ODESolver.SemiImplicit
head,_ = tf.nn.dynamic_rnn(self.wm,head,dtype=tf.float32,time_major=True)
self.constrain_op = self.wm.get_param_constrain_op()
elif(model_type == "node"):
self.fused_cell = NODE(model_size,cell_clip=10)
head,_ = tf.nn.dynamic_rnn(self.fused_cell,head,dtype=tf.float32,time_major=True)
elif(model_type == "ctgru"):
self.fused_cell = CTGRU(model_size,cell_clip=-1)
head,_ = tf.nn.dynamic_rnn(self.fused_cell,head,dtype=tf.float32,time_major=True)
elif(model_type == "ctrnn"):
self.fused_cell = CTRNN(model_size,cell_clip=-1,global_feedback=True)
head,_ = tf.nn.dynamic_rnn(self.fused_cell,head,dtype=tf.float32,time_major=True)
else:
raise ValueError("Unknown model type '{}'".format(model_type))
# target_y = tf.expand_dims(self.target_y,axis=-1)
self.y = tf.layers.Dense(1,activation=None,kernel_initializer=tf.keras.initializers.TruncatedNormal())(head)
print("logit shape: ",str(self.y.shape))
self.loss = tf.reduce_mean(tf.square(self.target_y-self.y))
optimizer = tf.train.AdamOptimizer(learning_rate)
self.train_step = optimizer.minimize(self.loss)
self.accuracy = tf.reduce_mean(tf.abs(self.target_y-self.y))
self.sess = tf.InteractiveSession()
self.sess.run(tf.global_variables_initializer())
self.result_file = os.path.join("results","power","{}_{}.csv".format(model_type,model_size))
if(not os.path.exists("results/power")):
os.makedirs("results/power")
if(not os.path.isfile(self.result_file)):
with open(self.result_file,"w") as f:
f.write("best epoch, train loss, train mae, valid loss, valid mae, test loss, test mae\n")
self.checkpoint_path = os.path.join("tf_sessions","power","{}".format(model_type))
if(not os.path.exists("tf_sessions/power")):
os.makedirs("tf_sessions/power")
self.saver = tf.train.Saver()