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)