zhry10 / phycnn Goto Github PK

def save_NN(self, model_name):
    saver = tf.train.Saver()
    save_path = saver.save(self.sess, model_name)
    print("Model saved in path: %s" % save_path)

def load_NN(self, model_name):
    saver = tf.train.Saver()
    saver.restore(self.sess, model_name)
    print("Model restored.")

# with tf.device('/device:GPU:1'):
with tf.device('/cpu:0'):

    config = tf.ConfigProto()
    config.gpu_options.allow_growth = True
    # config.gpu_options.per_process_gpu_memory_fraction = 0.4
    session = tf.Session(config=config)
    # tf.Session(config=tf.ConfigProto(log_device_placement=True))

    # Training
    model = DeepPhyLSTM(eta_tt_train, ag_train, Phi_t)

    Loss = model.train(num_epochs=200, batch_size=N_train, learning_rate=1e-3, bfgs=1)
    model.save_NN('./model_dir/model_name')

with tf.device('/cpu:0'):

    config = tf.ConfigProto()
    config.gpu_options.allow_growth = True
    # config.gpu_options.per_process_gpu_memory_fraction = 0.4
    session = tf.Session(config=config)
    # tf.Session(config=tf.ConfigProto(log_device_placement=True))

    # Training
    model2 = DeepPhyLSTM(eta_tt_train, ag_train, Phi_t)
    model2.load_NN('./model_dir_5')