def train(batch_size, epoch_count, lamda, datasetA_folder_path,
datasetB_folder_path, output_path):
dataset_A = data_io.dataset_load(datasetA_folder_path)
train_iter_A = SerialIterator(dataset_A,
batch_size,
repeat=True,
shuffle=True)
dataset_B = data_io.dataset_load(datasetB_folder_path)
train_iter_B = SerialIterator(dataset_B,
batch_size,
repeat=True,
shuffle=True)
g_ab = Generator()
g_ba = Generator()
d_a = Discriminator()
d_b = Discriminator()
g_ab.to_gpu(0)
g_ba.to_gpu(0)
d_a.to_gpu(0)
d_b.to_gpu(0)
opt_g_ab = chainer.optimizers.Adam(alpha=0.0002, beta1=0.5)
opt_g_ab.setup(g_ab)
opt_g_ba = chainer.optimizers.Adam(alpha=0.0002, beta1=0.5)
opt_g_ba.setup(g_ba)
opt_d_a = chainer.optimizers.Adam(alpha=0.0002, beta1=0.5)
opt_d_a.setup(d_a)
opt_d_b = chainer.optimizers.Adam(alpha=0.0002, beta1=0.5)
opt_d_b.setup(d_b)
iteration = 0
train_iter_A.reset()
train_iter_B.reset()
log_list = []
image_path = output_path + "image/"
disA_model_path = output_path + "dis_A/"
disB_model_path = output_path + "dis_B/"
genAB_model_path = output_path + "gen_AB/"
genBA_model_path = output_path + "gen_BA/"
os.mkdir(output_path)
os.mkdir(image_path)
os.mkdir(disA_model_path)
os.mkdir(disB_model_path)
os.mkdir(genAB_model_path)
os.mkdir(genBA_model_path)
for epoch in range(epoch_count):
d_a_loss_list = []
d_b_loss_list = []
g_AB_loss_list = []
g_BA_loss_list = []
while True:
mini_batch_images_A = train_iter_A.next()
mini_batch_images_A = np.array(mini_batch_images_A)
mini_batch_images_A = (mini_batch_images_A - 128.0) / 128.0
real_a = Variable(np.array(mini_batch_images_A))
real_a.to_gpu(0)
mini_batch_images_B = train_iter_B.next()
mini_batch_images_B = np.array(mini_batch_images_B)
mini_batch_images_B = (mini_batch_images_B - 128.0) / 128.0
real_b = Variable(np.array(mini_batch_images_B))
real_b.to_gpu(0)
fake_b = g_ab(real_a)
fake_a = g_ba(real_b)
reconstr_a = g_ba(fake_b)
reconstr_b = g_ab(fake_a)
d_a_real_result = d_a(real_a)
d_a_fake_result = d_a(fake_a)
loss_d_a = loss_dis(batch_size, d_a_real_result, d_a_fake_result)
d_b_real_result = d_b(real_b)
d_b_fake_result = d_b(fake_b)
loss_d_b = loss_dis(batch_size, d_b_real_result, d_b_fake_result)
d_a.cleargrads()
loss_d_a.backward()
opt_d_a.update()
d_b.cleargrads()
loss_d_b.backward()
opt_d_b.update()
"""generatorのloss計算"""
loss_g_ab = loss_gen(batch_size, d_b_fake_result, real_a,
reconstr_a, lamda)
loss_g_ba = loss_gen(batch_size, d_a_fake_result, real_b,
reconstr_b, lamda)
g_ab.cleargrads()
loss_g_ab.backward()
opt_g_ab.update()
g_ba.cleargrads()
loss_g_ba.backward()
opt_g_ba.update()
loss_d_a.to_cpu()
loss_d_b.to_cpu()
loss_g_ab.to_cpu()
loss_g_ba.to_cpu()
iteration += batch_size
d_a_loss_list.append(loss_d_a.array)
d_b_loss_list.append(loss_d_b.array)
g_AB_loss_list.append(loss_g_ab.array)
g_BA_loss_list.append(loss_g_ba.array)
if train_iter_A.is_new_epoch or train_iter_B.is_new_epoch:
break
real_a.to_cpu()
fake_b.to_cpu()
reconstr_a.to_cpu()
real_b.to_cpu()
fake_a.to_cpu()
reconstr_b.to_cpu()
real_a_images = real_a.array.transpose(0, 2, 3, 1)
fake_b_images = fake_b.array.transpose(0, 2, 3, 1)
reconstr_a_images = reconstr_a.array.transpose(0, 2, 3, 1)
real_b_images = real_b.array.transpose(0, 2, 3, 1)
fake_a_images = fake_a.array.transpose(0, 2, 3, 1)
reconstr_b_images = reconstr_b.array.transpose(0, 2, 3, 1)
data_io.output_images(image_path + str(epoch), real_a_images,
fake_b_images, reconstr_a_images, real_b_images,
fake_a_images, reconstr_b_images)
print("epoch: " + str(epoch) + ", interation: " + str(iteration) + \
", d_A_loss: " + str(np.mean(d_a_loss_list)) + ", d_B_loss: " + str(np.mean(d_b_loss_list)) + \
", g_AB_loss: " + str(np.mean(g_AB_loss_list)) + ", g_BA_loss: " + str(np.mean(g_BA_loss_list)))
log_json = {"epoch": str(epoch), "interation": str(iteration), \
"d_A_loss": str(np.mean(d_a_loss_list)), "d_B_loss": str(np.mean(d_b_loss_list)), \
"g_AB_loss": str(np.mean(g_AB_loss_list)), "g_BA_loss": str(np.mean(g_BA_loss_list))}
log_list.append(log_json)
with open(output_path + 'log.json', 'w') as log_file:
json.dump(log_list, log_file, indent=4)
if (epoch % 100 == 0):
g_ab.to_cpu()
g_ba.to_cpu()
d_a.to_cpu()
d_b.to_cpu()
save_npz(genAB_model_path + str(epoch) + '.npz', g_ab)
save_npz(genBA_model_path + str(epoch) + '.npz', g_ba)
save_npz(disA_model_path + str(epoch) + '.npz', d_a)
save_npz(disB_model_path + str(epoch) + '.npz', d_b)
g_ab.to_gpu(0)
g_ba.to_gpu(0)
d_a.to_gpu(0)
d_b.to_gpu(0)
g_ab.to_cpu()
g_ba.to_cpu()
d_a.to_cpu()
d_b.to_cpu()
save_npz(genAB_model_path + 'last.npz', g_ab)
save_npz(genBA_model_path + 'last.npz', g_ba)
save_npz(disA_model_path + 'last.npz', d_a)
save_npz(disB_model_path + 'last.npz', d_b)
def train(batch_size, epoch_count, lamda, datasetA_folder_path,
datasetB_folder_path, output_path):
print("Start load images data from " + datasetA_folder_path)
dataset_A = data_io.dataset_load(datasetA_folder_path)
print("Finish load images data from " + datasetA_folder_path)
print("Start load images data from " + datasetB_folder_path)
dataset_B = data_io.dataset_load(datasetB_folder_path)
print("Finish load images data from " + datasetB_folder_path)
if len(dataset_A) != len(dataset_B):
print("Error! Datasets are not paired data.")
exit()
gen = Generator()
dis = Discriminator()
gen.to_gpu(0)
dis.to_gpu(0)
opt_g = chainer.optimizers.Adam(alpha=0.0002, beta1=0.5)
opt_g.setup(gen)
opt_d = chainer.optimizers.Adam(alpha=0.0002, beta1=0.5)
opt_d.setup(dis)
iteration = 0
log_list = []
image_path = output_path + "image/"
dis_model_path = output_path + "dis/"
gen_model_path = output_path + "gen/"
os.mkdir(output_path)
os.mkdir(image_path)
os.mkdir(dis_model_path)
os.mkdir(gen_model_path)
dataset_num = min(dataset_A.shape[0], dataset_B.shape[0])
for epoch in range(epoch_count):
d_loss_list = []
g_loss_list = []
for i in range(dataset_num // batch_size):
input_image = dataset_A[i * batch_size:(i + 1) * batch_size]
input_image = Variable(input_image)
input_image.to_gpu(0)
correct_image = dataset_B[i * batch_size:(i + 1) * batch_size]
correct_image = Variable(correct_image)
correct_image.to_gpu(0)
fake_image = gen(input_image)
d_real_result = dis(correct_image)
d_fake_result = dis(fake_image)
loss_d = loss_dis(batch_size, d_real_result, d_fake_result)
dis.cleargrads()
loss_d.backward()
opt_d.update()
"""generatorのloss計算"""
loss_g = loss_gen(d_fake_result, fake_image, correct_image, lamda)
gen.cleargrads()
loss_g.backward()
opt_g.update()
loss_d.to_cpu()
loss_g.to_cpu()
iteration += batch_size
d_loss_list.append(loss_d.array)
g_loss_list.append(loss_g.array)
input_image.to_cpu()
correct_image.to_cpu()
fake_image.to_cpu()
input_images = input_image.array.transpose(0, 2, 3, 1)
correct_images = correct_image.array.transpose(0, 2, 3, 1)
fake_images = fake_image.array.transpose(0, 2, 3, 1)
data_io.output_images(image_path + str(epoch), input_images,
correct_images, fake_images)
print("epoch: " + str(epoch) + ", interation: " + str(iteration) + \
", d_loss: " + str(np.mean(d_loss_list)) + ", g_loss: " + str(np.mean(g_loss_list)))
log_json = {"epoch": str(epoch), "interation": str(iteration), \
"d_loss": str(np.mean(d_loss_list)),"g_loss": str(np.mean(g_loss_list))}
log_list.append(log_json)
with open(output_path + 'log.json', 'w') as log_file:
json.dump(log_list, log_file, indent=4)
if (epoch % 100 == 0):
gen.to_cpu()
dis.to_cpu()
save_npz(gen_model_path + str(epoch) + '.npz', gen)
save_npz(dis_model_path + str(epoch) + '.npz', dis)
gen.to_gpu(0)
dis.to_gpu(0)
gen.to_cpu()
dis.to_cpu()
save_npz(gen_model_path + 'last.npz', gen)
save_npz(dis_model_path + 'last.npz', dis)
def train(batch_size, epoch_count, dataset_folder_path, n_hidden, output_path):
dataset = data_io.dataset_load(dataset_folder_path)
train_iter = SerialIterator(dataset, batch_size, repeat=True, shuffle=True)
gen = Generator(n_hidden=n_hidden)
dis = Discriminator()
gen.to_gpu(0)
dis.to_gpu(0)
opt_gen = chainer.optimizers.Adam(alpha=0.0002, beta1=0.5)
opt_gen.setup(gen)
opt_dis = chainer.optimizers.Adam(alpha=0.0002, beta1=0.5)
opt_dis.setup(dis)
iteration = 0
train_iter.reset()
log_list = []
image_path = output_path + "image/"
dis_model_path = output_path + "dis/"
gen_model_path = output_path + "gen/"
os.mkdir(output_path)
os.mkdir(image_path)
os.mkdir(dis_model_path)
os.mkdir(gen_model_path)
for epoch in range(epoch_count):
d_loss_list = []
g_loss_list = []
while True:
mini_batch_images = train_iter.next()
mini_batch_images = np.array(mini_batch_images)
mini_batch_images = (mini_batch_images - 128.0) / 128.0
x_real = Variable(np.array(mini_batch_images))
x_real.to_gpu(0)
y_real = dis(x_real)
noise = xp.random.uniform(-1,
1, (batch_size, n_hidden),
dtype=np.float32)
z = Variable(noise)
x_fake = gen(z, batch_size)
y_fake = dis(x_fake)
d_loss = loss_dis(batch_size, y_real, y_fake)
g_loss = loss_gen(batch_size, y_fake)
dis.cleargrads()
d_loss.backward()
opt_dis.update()
gen.cleargrads()
g_loss.backward()
opt_gen.update()
d_loss.to_cpu()
g_loss.to_cpu()
iteration += batch_size
d_loss_list.append(d_loss.array)
g_loss_list.append(g_loss.array)
if train_iter.is_new_epoch:
break
x_fake.to_cpu()
generated_images = x_fake.array
generated_images = generated_images.transpose(0, 2, 3, 1)
Image.fromarray(
np.clip(generated_images[0] * 255, 0.0, 255.0).astype(
np.uint8)).save(image_path + str(epoch) + ".png")
print("epoch: " + str(epoch) + ", interation: " + str(iteration) +
", d_loss: " + str(np.mean(d_loss_list)) + ", g_loss: " +
str(np.mean(g_loss_list)))
log_json = {
"epoch": str(epoch),
"iteration": str(iteration),
"d_loss": str(np.mean(d_loss_list)),
"g_loss": str(np.mean(g_loss_list))
}
log_list.append(log_json)
with open(output_path + 'log.json', 'w') as log_file:
json.dump(log_list, log_file, indent=4)
if (epoch % 100 == 0):
dis.to_cpu()
save_npz(dis_model_path + str(epoch) + '.npz', dis)
gen.to_cpu()
save_npz(gen_model_path + str(epoch) + '.npz', gen)
gen.to_gpu(0)
dis.to_gpu(0)
logGraph.save_log_graph(output_path + 'log.json',
output_path + "lossGraph.png")
dis.to_cpu()
save_npz(dis_model_path + 'last.npz', dis)
gen.to_cpu()
save_npz(gen_model_path + 'last.npz', gen)
