def train_discriminator_wrapper(x, x_gen, batch_size=1, vocab_size=10):
y = gen_label(len(x),fixed_value=1)
y_gen = gen_label(len(x_gen),fixed_value=0)
x_train = torch.cat([x.int(),x_gen.int()], dim=0)
y_train = torch.cat([y,y_gen], dim=0)
discriminator = train_discriminator(x_train, y_train, batch_size, vocab_size)
return discriminator
def sanityCheck_rewards(batch_size=5):
''' test rewards generation '''
log = openLog('test.txt')
log.write('\n\nTest rollout.sanityCheck_rewards: {}'.format(
datetime.now()))
try:
generator, _, y_output_all = sanityCheck_generator(
batch_size=batch_size, sample_size=batch_size * 2)
gen_output = y_output_all[-batch_size:, :]
rollout = Rollout(generator=generator)
rollout = nn.DataParallel(rollout)
rollout.to(DEVICE)
discriminator = train_discriminator(
batch_size=batch_size,
vocab_size=generator.pretrain_model.module.vocab_size)
rewards = getReward(gen_output, rollout, discriminator)
log.write('\n rollout.sanityCheck_rewards SUCCESSFUL. {}\n'.format(
datetime.now()))
log.close()
return rewards
except:
log.write(
'\n rollout.sanityCheck_rewards !!!!!! UNSUCCESSFUL !!!!!! {}\n'.
format(datetime.now()))
log.close()
return None
def train( num_generations, generation_size, starting_selfies, max_molecules_len,
disc_epochs_per_generation, disc_enc_type, disc_layers, training_start_gen,
device, properties_calc_ls, num_processors, beta,
max_fitness_collector, impose_time_adapted_pen):
# Collecting all generated molecules and counting occurances
smiles_all = []
selfies_all = []
smiles_all_counter = {}
#Preparing generator and discriminator
initial_population = du.sanitize_multiple_smiles([decoder(selfie) for selfie in starting_selfies])
reference_mols = dict.fromkeys(du.read_dataset_encoding(disc_enc_type), '')
discriminator, d_opt, d_loss = D.build_discriminator(disc_enc_type, disc_layers, max_molecules_len, device)
total_time = time.time()
for generation_index in range(1, num_generations+1):
print(f"###Generation {generation_index} of {num_generations}")
start_time = time.time()
# Molecules of previous population
smiles_here, selfies_here = G.get_prev_gen_mols(initial_population, starting_selfies, generation_size,
generation_index, selfies_all, smiles_all)
# Fitness of previous population
fitness_here, order, fitness_ordered, smiles_ordered, selfies_ordered = G.get_fitness( disc_enc_type, smiles_here, selfies_here,
properties_calc_ls, discriminator, generation_index,
max_molecules_len, device, generation_size,
num_processors, writer, beta,
image_dir, data_dir, max_fitness_collector, impose_time_adapted_pen
)
# Molecules to replace and to keep
to_replace, to_keep = G.cutoff_population(order, generation_size)
# Molecules of current population
smiles_mutated, selfies_mutated = G.get_next_gen_mols(order, to_replace, to_keep,
selfies_ordered, smiles_ordered, max_molecules_len)
# Results of generator
smiles_all, selfies_all, smiles_all_counter = G.update_gen_res(smiles_all, smiles_mutated, selfies_all, selfies_mutated, smiles_all_counter)
# Data for discriminator training
dataset_x, dataset_y = G.get_dis_data(disc_enc_type, reference_mols, smiles_mutated, selfies_mutated, max_molecules_len, num_processors, generation_index)
if generation_index >= training_start_gen:
discriminator = D.train_discriminator(dataset_x, dataset_y, discriminator, d_opt, d_loss , disc_epochs_per_generation, generation_index-1, device, writer, data_dir)
D.save_model(discriminator, generation_index-1, saved_models_dir)
print('Generation time: ', round((time.time()-start_time), 2), ' seconds')
print('Experiment time: ', round((time.time()-total_time)/60, 2), ' mins')
print('Number of unique molecules: ', len(smiles_all_counter))
return smiles_all_counter
def train_adversarial(sess, saver, MODEL_STRING, generator, discriminator,
rollout, dis_data_loader, likelihood_data_loader,
task, log, n):
print('#################################################################')
print('Start Adversarial Training...')
log.write('adversarial training...\n')
saver.restore(sess, tf.train.latest_checkpoint(MODEL_STRING))
small_loss = float('inf')
for total_batch in range(n):
# Train the generator for one step
samples = generator.generate(sess)
rewards = rollout.get_reward(sess, samples, 16, discriminator) #I might actually need to change the value 16 here.
feed = {generator.x: samples, generator.rewards: rewards}
_ = sess.run(generator.g_updates, feed_dict=feed)
# Test
sample = generate_samples(sess, generator, BATCH_SIZE,
task.generated_num, task.eval_file)
print("Examples from generator:")
for sample in task.vocab.decode(samples)[:5]:
print(sample)
likelihood_data_loader.create_batches(task.valid_file)
test_loss = target_loss(sess, generator, likelihood_data_loader)
if test_loss < small_loss:
small_loss = test_loss
saver.save(sess, MODEL_STRING +"/model")
print("Saving checkpoint ...")
print("total_batch: ", total_batch, "test_loss: ", test_loss)
buffer = "total_batch: " + str(total_batch) + "test_loss: " + str(test_loss)
log.write(buffer)
# Update roll-out parameters
rollout.update_params()
# Train the discriminator for 5 steps
train_discriminator(sess, generator, discriminator, dis_data_loader,
task, log, 5, BATCH_SIZE, task.generated_num,
dis_dropout_keep_prob)
def main():
# Create a parser to parse user input
def parse_arguments():
parser = argparse.ArgumentParser(description='Program for running several SeqGan applications.')
parser.add_argument('app', metavar='application', type=str, choices=['obama', 'haiku', 'synth'],
help='Enter either \'obama\' or \'haiku\'')
parser.add_argument('gen_n', type = int,
help='Number of generator pre-training steps')
parser.add_argument('disc_n', type = int,
help='Number of discriminator pre-training steps')
parser.add_argument('adv_n', type = int,
help='Number of adversarial pre-training steps')
parser.add_argument('-mn', metavar="model_name", type = str, default = "",
help = "Name for the checkpoint files. Will be stored at ./<app>/models/<model_name>")
parser.add_argument('-numeat', metavar="num_eat", type = int, default = 500,
help = "For synthetic data generation. Determines number of eaters in vocab.")
parser.add_argument('-numfeed', metavar="num_feed", type = int, default = 500,
help = "For synthetic data generation. Determines number of feeders in vocab.")
parser.add_argument('-numsent', metavar="num_sent", type = int, default = 10000,
help = "For synthetic data generation. Determines number of sentences generated.")
args = parser.parse_args()
synth_gen_params = ("NA", "NA", "NA")
if args.app == "synth":
synth_gen_params = (args.numsent, args.numfeed, args.numeat)
generate_random_sents("../data/synth/input.txt", args.numsent, args.numfeed, args.numeat)
task = load_task(args.app)
#Make the /models directory if its not there.
model_string = task.path +"/models/"
if not os.path.exists("./"+model_string):
os.mkdir("./"+model_string)
#make the checkpoint directory if its not there.
if args.mn == "":
model_string += str(args.gen_n)+ "_" + str(args.disc_n) + "_" + str(args.adv_n)
model_string += time.strftime("_on_%m_%d_%y", time.gmtime())
else:
model_string += args.mn
if not os.path.exists("./"+model_string):
os.mkdir("./"+model_string)
return args.gen_n, args.disc_n, args.adv_n, model_string, task, synth_gen_params
gen_n, disc_n, adv_n, MODEL_STRING, task, SYNTH_GEN_PARAMS = parse_arguments()
assert START_TOKEN == 0
tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)
# Initialize the data loaders
gen_data_loader = Gen_Data_loader(BATCH_SIZE, task.max_seq_length)
likelihood_data_loader = Gen_Data_loader(BATCH_SIZE, task.max_seq_length) # For validation
dis_data_loader = Dis_dataloader(BATCH_SIZE, task.max_seq_length)
# Initialize the Generator
generator = Generator(len(task.vocab), BATCH_SIZE, EMB_DIM, HIDDEN_DIM,
task.max_seq_length, START_TOKEN)
# Initialize the Discriminator
discriminator = Discriminator(sequence_length=task.max_seq_length,
num_classes=2,
vocab_size=len(task.vocab),
embedding_size=dis_embedding_dim,
filter_sizes=dis_filter_sizes,
num_filters=dis_num_filters,
l2_reg_lambda=dis_l2_reg_lambda)
# Set session configurations.
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
saver = tf.train.Saver()
sess = tf.Session(config=config)
sess.run(tf.global_variables_initializer())
# If restoring from a previous run ....
if len(os.listdir("./"+MODEL_STRING)) > 0:
saver.restore(sess, tf.train.latest_checkpoint(MODEL_STRING))
# Create batches from the positive file.
gen_data_loader.create_batches(task.train_file)
# Open log file for writing
log = open(task.log_file, 'w')
# Pre_train the generator with MLE.
pre_train_generator(sess, saver, MODEL_STRING, generator, gen_data_loader,
likelihood_data_loader, task, log, gen_n, BATCH_SIZE,
task.generated_num)
print('Start pre-training discriminator...')
# Do the discriminator pre-training steps
saver.restore(sess, tf.train.latest_checkpoint(MODEL_STRING))
train_discriminator(sess, generator, discriminator, dis_data_loader,
task, log, disc_n, BATCH_SIZE, task.generated_num,
dis_dropout_keep_prob)
print("Saving checkpoint ...")
saver.save(sess, MODEL_STRING+ "/model")
# Do the adversarial training steps
rollout = ROLLOUT(generator, 0.8)
train_adversarial(sess, saver, MODEL_STRING, generator, discriminator,
rollout, dis_data_loader, likelihood_data_loader,
task, log, adv_n)
#Use the best model to generate final sample
saver.restore(sess, tf.train.latest_checkpoint(MODEL_STRING))
generate_samples(sess, generator, BATCH_SIZE, task.generated_num,
task.eval_file)
#Writing results to CSV
with open(task.eval_file) as f:
generated = []
for line in f:
line = line.strip().split()
generated.append(line)
generated = task.vocab.decode(generated)
f.close()
with open(task.test_file) as f:
references = []
for line in f:
line = line.strip().split()
references.append(line)
references = task.vocab.decode(references)
f.close()
blue = corpus_bleu([references]*len(generated), generated)
print("Run with args {} {} {}: BLEUscore = {}\n".format(gen_n, disc_n, adv_n, blue))
prop = "NA"
if task.name == "synth":
total_correct = 0
for sentence in generated:
if is_valid_phrase(sentence):
total_correct +=1
prop = total_correct/len(generated)
if not os.path.exists("./results.csv"):
os.mknod("./results.csv")
with open("./results.csv", 'a') as csvfile:
fieldnames = ["name", "task_name", "num_gen", "num_disc", "num_adv",
"num_sents", "num_feeders", "num_eaters", "BLEU", "prop_valid"]
writer = csv.DictWriter(csvfile, fieldnames = fieldnames)
writer.writeheader()
writer.writerow({"name": MODEL_STRING, "task_name": task.name, "num_gen": gen_n,
"num_disc":disc_n, "num_adv": adv_n, "num_sents":SYNTH_GEN_PARAMS[0],
"num_feeders":SYNTH_GEN_PARAMS[1], "num_eaters":SYNTH_GEN_PARAMS[2],
"BLEU": blue, "prop_valid": prop})
f.close()
log.close()
from graph_freezer import freezing_graph
from discriminator import train_discriminator
from estimator import train_estimator
discriminator = train_discriminator()
estimator = train_estimator()
freezing_graph("discriminator")
freezing_graph("estimator")