def test(net, test_input_image_dir, direction):
# prepare dataset
test_dataset = helper.Dataset(test_input_image_dir,
convert_to_lab_color=False,
direction=direction,
is_test=True)
saver = tf.train.Saver()
with tf.Session() as sess:
saver.restore(sess, tf.train.latest_checkpoint('./checkpoints'))
for ii in range(test_dataset.n_images):
test_image = test_dataset.get_image_by_index(ii)
test_a = [x for x, y in test_image]
test_b = [y for x, y in test_image]
test_a = np.array(test_a)
test_b = np.array(test_b)
gen_image = sess.run(generator(net.gen_inputs,
net.input_channel,
reuse=True,
is_training=True),
feed_dict={net.gen_inputs: test_a})
image_fn = './assets/test_result{:d}_tf.png'.format(ii)
helper.save_result(image_fn,
gen_image,
input_image=test_a,
target_image=test_b)
def train(net, epochs, batch_size, train_input_image_dir, test_image, direction, dataset_name, print_every=30):
losses = []
steps = 0
# prepare saver for saving trained model
saver = tf.train.Saver()
# prepare dataset
train_dataset = helper.Dataset(train_input_image_dir, convert_to_lab_color=False, direction=direction, is_test=False)
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
for e in range(epochs):
for ii in range(train_dataset.n_images//batch_size):
steps += 1
# will return list of tuples [ (inputs, targets), (inputs, targets), ... , (inputs, targets)]
batch_images_tuple = train_dataset.get_next_batch(batch_size)
a = [x for x, y in batch_images_tuple]
b = [y for x, y in batch_images_tuple]
a = np.array(a)
b = np.array(b)
fd = {
net.dis_inputs: a,
net.dis_targets: b,
net.gen_inputs: a
}
d_opt_out = sess.run(net.d_train_opt, feed_dict=fd)
g_opt_out = sess.run(net.g_train_opt, feed_dict=fd)
if steps % print_every == 0:
# At the end of each epoch, get the losses and print them out
train_loss_d = net.d_loss.eval(fd)
train_loss_g = net.g_loss.eval(fd)
train_loss_gan = net.gen_loss_GAN.eval(fd)
train_loss_l1 = net.gen_loss_L1.eval(fd)
print("Epoch {}/{}...".format(e + 1, epochs),
"Discriminator Loss: {:.4f}...".format(train_loss_d),
"Generator Loss GAN: {:.4f}".format(train_loss_gan),
"Generator Loss L1: {:.4f}".format(train_loss_l1),
"Generator Loss: {:.4f}".format(train_loss_g))
# Save losses to view after training
losses.append((train_loss_d, train_loss_gan))
# save generated images on every epochs
test_a = [x for x, y in test_image]
test_a = np.array(test_a)
gen_image = sess.run(generator(net.gen_inputs, net.input_channel, reuse=True, is_training=False),
feed_dict={net.gen_inputs: test_a})
image_fn = './assets/epoch_{:d}_tf.png'.format(e)
helper.save_result(image_fn, gen_image)
ckpt_fn = './checkpoints/pix2pix-{}.ckpt'.format(dataset_name)
saver.save(sess, ckpt_fn)
return losses
def train(net,
dataset_name,
train_data_loader,
val_data_loader,
epochs,
batch_size,
print_every=30,
save_every=100):
losses = []
steps = 0
# prepare saver for saving trained model
saver = tf.train.Saver()
start_time = time.time()
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
for e in range(epochs):
# shuffle data randomly at every epoch
train_data_loader.reset()
# val_data_loader.reset()
for ii in range(train_data_loader.n_images // batch_size):
steps += 1
batch_image_u, batch_image_v = train_data_loader.get_next_batch(
batch_size)
net.run_optim(sess, batch_image_u, batch_image_v, steps,
start_time)
if steps % save_every == 0:
# save generated images on every epochs
random_index = np.random.randint(0,
val_data_loader.n_images)
test_image_u, test_image_v = val_data_loader.get_image_by_index(
random_index)
fd_val = {
net.real_A: test_image_u,
net.real_B: test_image_v
}
g_image_u_to_v_to_u, g_image_u_to_v = sess.run(
[net.A2B2A, net.A2B], feed_dict=fd_val)
g_image_v_to_u_to_v, g_image_v_to_u = sess.run(
[net.B2A2B, net.B2A], feed_dict=fd_val)
image_fn = './assets/{:s}/epoch_{:d}-{:d}_tf.png'.format(
dataset_name, e, steps)
helper.save_result(image_fn, test_image_u, g_image_u_to_v,
g_image_u_to_v_to_u, test_image_v,
g_image_v_to_u, g_image_v_to_u_to_v)
ckpt_fn = './checkpoints/DualGAN-{:s}.ckpt'.format(dataset_name)
saver.save(sess, ckpt_fn)
return losses
def test(net, dataset_name, val_data_loader):
ckpt_fn = './checkpoints/DualGAN-{:s}.ckpt'.format(dataset_name)
saver = tf.train.Saver()
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
saver.restore(sess, ckpt_fn)
for ii in range(val_data_loader.n_images):
test_image_u, test_image_v = val_data_loader.get_image_by_index(ii)
gen_A_out, gen_AB_out = sess.run(
[net.gen_A_out, net.gen_AB_out],
feed_dict={net.input_u: test_image_u})
gen_B_out, gen_BA_out = sess.run(
[net.gen_B_out, net.gen_BA_out],
feed_dict={net.input_v: test_image_v})
image_fn = './assets/{:s}/{:s}_result_{:04d}_tf.png'.format(
dataset_name, dataset_name, ii)
helper.save_result(image_fn, test_image_u, gen_A_out, gen_AB_out,
test_image_v, gen_B_out, gen_BA_out)
def train(net,
dataset_name,
train_data_loader,
val_data_loader,
epochs,
batch_size,
print_every=30,
save_every=100):
steps = 0
discriminator_losses, generator_losses = [], []
# prepare saver for saving trained model
saver = tf.train.Saver()
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
for e in range(epochs):
# shuffle data randomly at every epoch
train_data_loader.reset()
# val_data_loader.reset()
epoch_disc_loss, epoch_gen_loss = [], []
for ii in range(train_data_loader.n_images // batch_size):
steps += 1
batch_image_u, batch_image_v = train_data_loader.get_next_batch(
batch_size)
fd = {net.input_u: batch_image_u, net.input_v: batch_image_v}
_ = sess.run(net.d_train_opt, feed_dict=fd)
_ = sess.run(net.g_train_opt, feed_dict=fd)
_ = sess.run(net.g_train_opt, feed_dict=fd)
if steps % print_every == 0:
# At the end of each epoch, get the losses and print them out
train_loss_d = net.d_loss.eval(fd)
train_loss_g = net.g_loss.eval(fd)
train_loss_A_l1 = net.gen_A_l1_loss.eval(fd)
train_loss_B_l1 = net.gen_B_l1_loss.eval(fd)
epoch_disc_loss.append(train_loss_d)
epoch_gen_loss.append(train_loss_g)
# print(f"This epoch, discriminator loss {epoch_disc_loss}, generator loss {epoch_gen_loss}")
# to provide visual feedback during the process.
print("Epoch {}/{}...".format(e + 1, epochs),
"Discriminator Loss: {:.4f}...".format(train_loss_d),
"Generator Loss: {:.4f}".format(train_loss_g),
"A-L1 Loss: {:.4f}".format(train_loss_A_l1),
"B-L1 Loss: {:.4f}".format(train_loss_B_l1))
if steps % save_every == 0:
# save generated images on every epochs
random_index = np.random.randint(0,
val_data_loader.n_images)
test_image_u, test_image_v = val_data_loader.get_image_by_index(
random_index)
gen_A_out, gen_AB_out = sess.run(
[net.gen_A_out, net.gen_AB_out],
feed_dict={net.input_u: test_image_u})
gen_B_out, gen_BA_out = sess.run(
[net.gen_B_out, net.gen_BA_out],
feed_dict={net.input_v: test_image_v})
image_fn = './assets/{:s}/epoch_{:d}-{:d}_tf.png'.format(
dataset_name, e, steps)
helper.save_result(image_fn, test_image_u, gen_A_out,
gen_AB_out, test_image_v, gen_B_out,
gen_BA_out)
discriminator_losses.append(
sum(epoch_disc_loss) / len(epoch_disc_loss))
generator_losses.append(sum(epoch_gen_loss) / len(epoch_gen_loss))
# print(f"After epoch {e+1}, discriminator loss {discriminator_losses}, generator loss {generator_losses}")
print(30 * "=")
# save checkpoint
ckpt_fn = './checkpoints/DualGAN-{:s}.ckpt'.format(dataset_name)
saver.save(sess, ckpt_fn)
epoch_list = [i for i in range(1, epochs + 1)]
# Single plot for discriminator and generator losses.
plt.figure()
plt.plot(epoch_list,
discriminator_losses,
color='red',
label='Discriminator')
plt.plot(epoch_list,
generator_losses,
color='blue',
label='Generator',
linestyle='dashed')
plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.legend(loc="upper left")
plt.savefig('./assets/loss_single_plot.png')
plt.close()
# 2 Subplots in 1 figure for generator and discriminator losses.
plt.figure()
plt.subplots_adjust(hspace=0.5, wspace=0.5)
#Subplot 1
plt.subplot(211)
plt.title("Generator")
plt.plot(epoch_list, generator_losses, color='blue', label='Generator')
plt.xlabel('Epochs')
plt.ylabel('Loss')
# Subplot 2
plt.subplot(212)
plt.title("Discriminator")
plt.plot(epoch_list, discriminator_losses, color='red')
plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.savefig('./assets/loss_subplots.png')
plt.close()
return
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--model',
choices=["DQN", "DDPG"],
default="DDPG",
help='Reinforcement Learning Algorithms')
parser.add_argument('--n_agent',
type=int,
default=5,
help='The number of agents')
parser.add_argument('--sustainable_weight',
type=float,
default=0.5,
help='Weight of sustainability goal')
parser.add_argument('--run_mode',
choices=["train", "test"],
default="train",
help='Train or test mode')
args = parser.parse_args()
with open('config/config.json', 'r') as f:
conf = json.load(f)
conf["hyper_parameter"]["sustainable_weight"] = args.sustainable_weight
timestamp = datetime.now().strftime("%d-%m-%Y_%H-%M-%S")
env_conf = conf["env"]
save_result_path = Path(env_conf["log_path"]). \
joinpath(f"{args.model}_model_{timestamp}_{args.sustainable_weight}_{args.n_agent}")
save_model_path = os.path.join(save_result_path, "checkpoints")
# Init game environment
cpr = CPREnvironment(conf["hyper_parameter"])
states = np.zeros((args.n_agent, env_conf["state_space"]))
next_states = np.zeros((args.n_agent, env_conf["state_space"]))
rewards = np.zeros(args.n_agent)
# -------------- train mode --------------
if args.run_mode == "train":
# Init agents
agents = list(
map(lambda _: Agent(conf["model"][args.model]),
range(args.n_agent)))
avg_scores = []
epsilon = env_conf["init_epsilon"]
epoch = 0
while epsilon >= env_conf["min_epsilon"]:
# Reset Game Environment
cpr.reset()
efforts = np.array([env_conf["total_init_effort"] / args.n_agent] *
args.n_agent)
score = 0.0
for _ in range(env_conf["max_train_round"]):
for i, agent in enumerate(agents):
if args.model == "DDPG":
action = agent.act(states[i],
epsilon=epsilon,
upper_bound=cpr.pool / args.n_agent)
else:
action = agent.act(states[i],
epsilon=epsilon,
pre_action=effort_list[i])
efforts[i] = action
agent.remember(states[i], action, rewards[i],
next_states[i])
next_states, rewards, done = cpr.step(efforts)
score += sum(rewards) / args.n_agent
if done:
break
for agent in agents:
agent.learn()
print(f"epoch: {epoch}, score: {score}, e: {epsilon:.2}")
epsilon -= env_conf["epsilon_decay"]
epoch += 1
avg_scores.append(score)
for agent in agents:
agent.close(save_model_path)
helper.save_result({"train_avg_score.txt": avg_scores},
save_result_path)
# -------------- test mode --------------
elif args.run_mode == "test":
# Init agents
agent_list = [
Agent("agent_{}".format(i), conf["model"][args.model],
save_model_path) for i in range(args.n_agent)
]
avg_asset_seq = [0]
pool_level_seq = []
avg_score_seq = []
for t in range(env_conf["max_test_round"]):
pool_level_seq.append(cpr.pool)
effort_list = [env_conf["total_init_effort"] / args.n_agent
] * args.n_agent
for i, agent in enumerate(agent_list):
if args.model == "DDPG":
action = agent.act(states[i],
upper_bound=cpr.pool / args.n_agent)
elif args.model == "DQN":
action = agent.act(states[i], pre_action=effort_list[i])
effort_list[i] = action
agent.remember(states[i], action, rewards[i], next_states[i])
effort_list[i] = action
next_states, rewards, done = cpr.step(effort_list)
avg_score_seq.append(sum(rewards) / args.n_agent)
avg_asset_seq.append(avg_asset_seq[-1] +
next_states[0][3] / args.n_agent)
if done:
break
for agent in agent_list:
agent.close()
helper.save_result(
{
"test_avg_score.txt": avg_asset_seq,
"test_assets.txt": avg_asset_seq,
"test_resource_level.txt": pool_level_seq
}, save_result_path)
def train(net,
dataset_name,
train_data_loader,
val_data_loader,
epochs,
batch_size,
print_every=30,
save_every=100):
steps = 0
# prepare saver for saving trained model
saver = tf.train.Saver()
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
for e in range(epochs):
# shuffle data randomly at every epoch
train_data_loader.reset()
# val_data_loader.reset()
for ii in range(train_data_loader.n_images // batch_size):
steps += 1
batch_image_u, batch_image_v = train_data_loader.get_next_batch(
batch_size)
fd = {net.input_u: batch_image_u, net.input_v: batch_image_v}
_ = sess.run(net.d_train_opt, feed_dict=fd)
_ = sess.run(net.g_train_opt, feed_dict=fd)
_ = sess.run(net.g_train_opt, feed_dict=fd)
if steps % print_every == 0:
# At the end of each epoch, get the losses and print them out
train_loss_d = net.d_loss.eval(fd)
train_loss_g = net.g_loss.eval(fd)
train_loss_A_l1 = net.gen_A_l1_loss.eval(fd)
train_loss_B_l1 = net.gen_B_l1_loss.eval(fd)
print("Epoch {}/{}...".format(e + 1, epochs),
"Discriminator Loss: {:.4f}...".format(train_loss_d),
"Generator Loss: {:.4f}".format(train_loss_g),
"A-L1 Loss: {:.4f}".format(train_loss_A_l1),
"B-L1 Loss: {:.4f}".format(train_loss_B_l1))
if steps % save_every == 0:
# save generated images on every epochs
random_index = np.random.randint(0,
val_data_loader.n_images)
test_image_u, test_image_v = val_data_loader.get_image_by_index(
random_index)
gen_A_out, gen_AB_out = sess.run(
[net.gen_A_out, net.gen_AB_out],
feed_dict={net.input_u: test_image_u})
gen_B_out, gen_BA_out = sess.run(
[net.gen_B_out, net.gen_BA_out],
feed_dict={net.input_v: test_image_v})
image_fn = './assets/{:s}/epoch_{:d}-{:d}_tf.png'.format(
dataset_name, e, steps)
helper.save_result(image_fn, test_image_u, gen_A_out,
gen_AB_out, test_image_v, gen_B_out,
gen_BA_out)
ckpt_fn = './checkpoints/DualGAN-{:s}.ckpt'.format(dataset_name)
saver.save(sess, ckpt_fn)
return