import numpy as np from matplotlib import pyplot as plt from keras.datasets import mnist from gan import GAN (x_train, y_train), (x_test, y_test) = mnist.load_data() x_train = x_train.reshape([-1, 28, 28, 1]) / 255 x_test = x_test.reshape([-1, 28, 28, 1]) / 255 x = np.concatenate([x_train, x_test]) gan = GAN() gan.gan.summary() gan.train(x, 100, batch_size=100) gan.save('number_gan/gan.h5') noise = np.random.normal(size=[10, 8 * 8]) images = gan.generator.predict(noise) images = images.reshape([-1, 28, 28]) plt.gray() for i in range(10): plt.subplot(2, 5, i + 1) plt.imshow(images[i]) plt.show()
opts.generator_frequency = 1 opts.generator_dropout = 0.3 opts.label_softness = 0.2 opts.batch_size = 128 opts.epoch_number = 300 #Note: if you are using jupyter notebook you might need to disable workers: opts.workers_nbr = 0 #List of tranformations applied to each input image opts.transforms = [ transforms.Resize(int(opts.image_size), Image.BICUBIC), transforms.ToTensor(), #do not forget to transform image into a tensor transforms.Normalize((.5, .5, .5), (.5, .5, .5)) #RandomNoise(0.01) #adding noise might prevent the discriminator from over-fitting ] #Build GAN model = GAN(opts) #Display GAN architecture (note: only work if cuda is enabled) summary(model.generator.cuda(), input_size=(opts.latent_dim, 1, 1)) summary(model.discriminator.cuda(), input_size=(opts.channels_nbr, opts.image_size, opts.image_size)) #Start training model.train() #Save model model.save(opts.output_path)
model = GAN(gex_size=train.shape[1], num_cells_generate=test.shape[0]) model.compile() model.build(input_shape=(model.hyperparams.batch_size, model.hyperparams.latent_dim)) # req. for subclassed models # process data for training train_tf = tf.data.Dataset.from_tensor_slices(train.X). \ cache(). \ shuffle(buffer_size=train.shape[0], seed=utils.RANDOM). \ batch(batch_size=model.hyperparams.batch_size * strategy.num_replicas_in_sync, num_parallel_calls=tf.data.AUTOTUNE). \ prefetch(buffer_size=tf.data.AUTOTUNE) train_tf_distributed = strategy.experimental_distribute_dataset(train_tf) test_tf = tf.data.Dataset.from_tensor_slices(test.X). \ cache(). \ shuffle(buffer_size=test.shape[0], seed=utils.RANDOM). \ prefetch(buffer_size=tf.data.AUTOTUNE) tb_callback = tf.keras.callbacks.TensorBoard(log_dir=utils.LOG_DIR + datetime.datetime.now().strftime("%Y%m%d-%H%M%S"), update_freq='epoch', write_graph=False, profile_batch=0) model.fit(x=train_tf_distributed, epochs=model.hyperparams.num_epochs, steps_per_epoch=int (train.shape[0] / model.hyperparams.batch_size), # steps = # batches per epoch callbacks=[tb_callback, TrainingMetrics(tb_callback), tf.keras.callbacks.ModelCheckpoint(period=100, filepath='training/model_{epoch}')]) # Save model model.save("training/model_final")
def main(num_faces, num_epochs, cuda, verbose=False): # set computation device (None/CPU if in development mode, CUDA otherwise) device = torch.device("cuda:0") if cuda else None # load faces masked_dir = "../data/masked" masked_suffix = "_Mask.jpg" unmasked_dir = "../data/unmasked" unmasked_suffix = ".png" masked_faces, unmasked_faces, idx_to_face_id = utils.load_faces( num_faces, masked_dir, masked_suffix, unmasked_dir, unmasked_suffix) if verbose: print("loaded {} faces...".format(num_faces)) # split data into training and testing sets split = int(0.8 * num_faces) train_input, train_output = ( masked_faces[:split], torch.Tensor((range(0, split))).long(), ) test_input, test_output = ( masked_faces[split:], torch.Tensor(range(split, num_faces)).long(), ) static_faces = unmasked_faces[:split] # instantiate GAN generator = Generator(learning_rate=2e-3) projector = Projector(load_path="../models/projector.pt") discriminator = Discriminator() gan = GAN(generator, projector, discriminator, device=device) if verbose: print("instantiated GAN...") # compute and store unmasked discriminator embeddings gan.compute_unmasked_embeddings(unmasked_faces) # train if verbose: print("training initiated...") gan.fit(train_input, static_faces, train_output, num_epochs=num_epochs, verbose=verbose) if verbose: print("\ntraining complete...") # save models save_dir = "../models" suffix = time.strftime("%Y%m%d_%H%M%S") gan.save(save_dir, suffix) if verbose: print("models saved under '{}/<model>_{}'...".format(save_dir, suffix)) # display sample masks and faces plt.figure() fig, axes = plt.subplots(2, 5) fig.set_figwidth(20) fig.set_figheight(7) for idx in range(5): # original image face_id = idx_to_face_id[idx] original_img = Image.open("../data/masked/{}_Mask.jpg".format(face_id)) axes[0, idx].imshow(original_img) axes[0, idx].get_xaxis().set_ticks([]) axes[0, idx].get_yaxis().set_ticks([]) # generated mask image mask = (gan.generator(torch.rand(1, 100).to(device)) if device else gan.generator()) masked_tensor = masked_faces[idx].unsqueeze(0) if device: masked_tensor = masked_tensor.to(device) masked_image = gan.project_mask(mask, masked_tensor, process=True)[0] masked_image = torch.transpose(masked_image, 0, 1) masked_image = torch.transpose(masked_image, 1, 2) masked_image = masked_image.cpu().detach().numpy() axes[1, idx].imshow(masked_image) axes[1, idx].get_xaxis().set_ticks([]) axes[1, idx].get_yaxis().set_ticks([]) plt.savefig("../figures/sample_masks.png") # evaluate accuracy train_accuracy = gan.evaluate(train_input, train_output) test_accuracy = gan.evaluate(test_input, test_output) masked_accuracy = gan.discriminator_evaluate(masked_faces, unmasked_faces) unmasked_accuracy = gan.discriminator_evaluate(unmasked_faces, unmasked_faces) print("\nfacial recognition accuracy for...") print(" random choice:\t\t{:.1f}%".format(100 / num_faces)) print(" training images:\t\t{:.1f}%".format(100 * train_accuracy)) print(" testing images:\t\t{:.1f}%".format(100 * test_accuracy)) print(" original masked images:\t{:.1f}%".format(100 * masked_accuracy)) print(" original unmasked images:\t{:.1f}%".format(100 * unmasked_accuracy)) # write results to file file_path = "../data/accuracy.txt" with open(file_path, "w") as file: file.write("facial recognition accuracy for...") file.write("\n random choice:\t\t{:.1f}%".format(100 / num_faces)) file.write("\n training images:\t\t{:.1f}%".format(100 * train_accuracy)) file.write("\n testing images:\t\t{:.1f}%".format(100 * test_accuracy)) file.write("\n original masked images:\t{:.1f}%".format( 100 * masked_accuracy)) file.write("\n original unmasked images:\t{:.1f}%".format( 100 * unmasked_accuracy)) if verbose: print("\nsaved results...") print("done:)")