from utils import (choose, split_trn_val, update_task, get_max_of_db_column,

get_a_task, get_task, ExploitationNeeded,

LossIsNaN, get_task_ids_and_scores, PopulationFinished,

get_col_from_populations, RemainingTasksTaken,

print_with_time, ExploitationOcurring,

create_new_population)

import argparse

import os

import time

import pathlib

from psycopg2.extensions import TransactionRollbackError

import numpy as np

import pickle

import multiprocessing

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' # or any {'0', '1', '2'}

HYPERPARAM_NAMES = ["lr", "momentum"] # This is unfortunate.

EPOCHS = 10

BATCH_SIZE = 64

POPULATION_SIZE = 5 # Number of models in a population

EXPLOIT_INTERVAL = 0.5 # When to exploit, in number of epochs

USE_SQLITE = True # If False, you'll need to set up a local Postgres server

# tf.enable_eager_execution()

def data():

return (x_train, y_train), (x_test, y_test)

def dnn_model():

return model

def get_optimizer():

return optimizer_class(lr=lr, momentum=momentum)

def save_obj(obj, path):

out_put = open(path, "wb")

pickle.dump(obj, out_put)

def load_obj(path):

read_file = open(path, "rb")

return pickle.load(read_file)

class Trainer:

self.optimizer = self.optimizer.from_config(param_group)

def init():

KTF.set_session(sess)

def exploit_and_explore(connect_str_or_path, population_id):

tf.keras.backend.clear_session()

def tran(x_train, y_train, x_test, y_test, epochs, batch_size, task_id, population_id,

# break

if __name__ == "__main__":

# tf.logging.set_verbosity(tf.logging.WARN)

# TODO: Does this help?

print("mian:", os.getpid())

parser = argparse.ArgumentParser(description="Population Based Training")

parser.add_argument("-p", "--population_id", type=int, default=None,

help="Resumes work on the population with the given ID. Use -1 to select the most recently created population. Without this flag, a new population will be created.")

parser.add_argument("-e", "--exploiter", action="store_true",

help="Set this process as the exploiter. It will be responsible for running the exploit step over the entire population at the end of each interval.")

args = parser.parse_args()

population_id = args.population_id

exploiter = args.exploiter

(x_train, y_train), (x_test, y_test) = data()

pathlib.Path('checkpoints').mkdir(exist_ok=True)

checkpoint_str = "checkpoints/pop-%03d_task-%03d.h5"

interval_limit = int(np.ceil(EPOCHS / EXPLOIT_INTERVAL))

table_name = "populations"

if USE_SQLITE:

sqlite_path = "database.sqlite3"

connect_str_or_path = sqlite_path

ready_for_exploitation_False = 0

ready_for_exploitation_True = 1

active_False = 0

active_True = 1

else: # Postgres

db_env_var_names = ['PGDATABASE', 'PGUSER', 'PGPORT', 'PGHOST']

db_params = [os.environ[var_name] for var_name in db_env_var_names]

db_connect_str = "dbname={} user={} port={} host={}".format(*db_params)

connect_str_or_path = db_connect_str

ready_for_exploitation_False = False

ready_for_exploitation_True = True

active_False = False

active_True = True

if population_id is None:

population_id = create_new_population(connect_str_or_path, USE_SQLITE,

POPULATION_SIZE)

msg = "Population added to populations table. Population ID: %s"

print_with_time(msg % population_id)

elif population_id == -1:

population_id = get_max_of_db_column(connect_str_or_path, USE_SQLITE,

table_name, "population_id")

# Train each available task for an interval

task_wait_count = 0

exploitation_wait_count = 0

start_time = int(time.time())

# global exploitation_wait_count

# global task_wait_count

while True:

try:

# Find a task that's incomplete and inactive, and set it to active

pool = multiprocessing.Pool(processes=POPULATION_SIZE, initializer=init) #

try:

tasks = get_task(connect_str_or_path, USE_SQLITE, population_id,

interval_limit, POPULATION_SIZE)

# task_id, intervals_trained, seed_for_shuffling = task

except RemainingTasksTaken:

if task_wait_count == 0:

print_with_time("Waiting for a task to be available.")

time.sleep(10)

task_wait_count += 1

continue

except PopulationFinished:

task_ids, scores = get_task_ids_and_scores(connect_str_or_path,

USE_SQLITE,

population_id)

print("Population finished. Best score: %.2f" % scores[0])

checkpoint_path = (checkpoint_str % (population_id, task_ids[0]))

pre, suf = checkpoint_path.split('.')

weights_path = pre + "_weights." + suf

print("Best weights saved to: %s" % weights_path)

break

except (ExploitationNeeded, ExploitationOcurring):

if exploiter:

pool.apply_async(exploit_and_explore, (connect_str_or_path, population_id))

pool.close()

pool.join()

time.sleep(1)

continue

else:

print_with_time("Waiting for exploiter to finish.")

time.sleep(10)

exploitation_wait_count += 1

if exploitation_wait_count > 11:

print_with_time(

"Exploiter is taking too long. Ending process.")

quit()

continue

except TransactionRollbackError:

print_with_time("Deadlock?")

time.sleep(1)

continue

# multiprocessing train

for task_id, intervals_trained, seed_for_shuffling in tasks:

pool.apply_async(tran, (

x_train, y_train, x_test, y_test, 1, BATCH_SIZE, task_id, population_id,

ready_for_exploitation_False,

ready_for_exploitation_True, active_False, active_True, connect_str_or_path, intervals_trained,

seed_for_shuffling))

pool.close()

pool.join()

# time.sleep(10)

except KeyboardInterrupt:

pool.terminate()

break