def make_batch_generators(mode, X_train = None, X_test = None, batch_size=8, siamese = False ): batch_gen_test = None if not siamese: batch_gen_train = batch_utils.gen_batch( X_train, batch_size, augment=mode=='train', randomize=mode=='train' ) if X_test is not None: batch_gen_test = batch_utils.gen_batch( X_test, batch_size, augment=mode=='train', randomize = False ) else: batch_gen_train = batch_utils.gen_siamese_batch( X_train, batch_size, augment=mode=='train', write_examples = False, randomize = mode=='train' ) if X_test is not None: batch_gen_test = batch_utils.gen_siamese_batch( X_test, batch_size, augment=mode=='train', write_examples = False, randomize = False ) return batch_gen_train, batch_gen_test
def train_model(training_data=None, n_levels=3, n_nodes=[10, 20, 40], input_dim=100, n_epochs=25, batch_size=64, n_batch_per_epoch=100, d_iters=20, lr_discriminator=0.005, lr_generator=0.00005, weight_constraint=[-0.01, 0.01], rule='mgd', verbose=True): """ Train the hierarchical model. Progressively generate trees with more and more nodes. Parameters ---------- training_data: dict of dicts each inner dict is an array 'geometry': 3-d arrays (locations) n_samples x n_nodes - 1 x 3 'morphology': 2-d arrays n_samples x n_nodes - 2 (prufer sequences) example: training_data['geometry']['n20'][0:10, :, :] gives the geometry for the first 10 neurons training_data['geometry']['n20'][0:10, :] gives the prufer sequences for the first 10 neurons here, 'n20' indexes a key corresponding to 20-node downsampled neurons. n_levels: int number of levels in the hierarchy n_nodes: list of length n_levels specifies the number of nodes for each level. should be consistent with training data. input_dim: int dimensionality of noise input n_epochs: number of epochs over training data batch_size: batch size n_batch_per_epoch: int number of batches per epoch d_iters: int number of iterations to train discriminator lr_discriminator: float learning rate for optimization of discriminator lr_generator: float learning rate for optimization of generator weight_constraint: array upper and lower bounds of weights (to clip) verbose: bool print relevant progress throughout training Returns ------- geom_model: list of keras model objects geometry generators for each level cond_geom_model: list of keras model objects conditional geometry generators for each level morph_model: list of keras model objects morphology generators for each level cond_morph_model: list of keras model objects conditional morphology generators for each level disc_model: list of keras model objects discriminators for each level gan_model: list of keras model objects discriminators stacked on generators for each level """ # ################################### # Initialize models at all levels # ################################### geom_model = list() cond_geom_model = list() morph_model = list() cond_morph_model = list() disc_model = list() gan_model = list() for level in range(n_levels): # Discriminator d_model = models.discriminator(n_nodes_in=n_nodes[level]) # Generators and GANs # If we are in the first level, no context if level == 0: g_model, cg_model, m_model, cm_model = \ models.generator(use_context=False, n_nodes_in=n_nodes[level-1], n_nodes_out=n_nodes[level]) mgd_model = \ models.discriminator_on_generators(g_model, cg_model, m_model, cm_model, d_model, conditioning_rule=rule, input_dim=input_dim, n_nodes_in=n_nodes[level-1], n_nodes_out=n_nodes[level], use_context=False) # In subsequent levels, we need context else: g_model, cg_model, m_model, cm_model = \ models.generator(use_context=True, n_nodes_in=n_nodes[level-1], n_nodes_out=n_nodes[level]) mgd_model = \ models.discriminator_on_generators(g_model, cg_model, m_model, cm_model, d_model, conditioning_rule=rule, input_dim=input_dim, n_nodes_in=n_nodes[level-1], n_nodes_out=n_nodes[level], use_context=True) # Collect all models into a list disc_model.append(d_model) geom_model.append(g_model) cond_geom_model.append(cg_model) morph_model.append(m_model) cond_morph_model.append(cm_model) gan_model.append(mgd_model) # ############### # Optimizers # ############### optim_d = Adagrad() # RMSprop(lr=lr_discriminator) optim_g = Adagrad() # RMSprop(lr=lr_generator) # ############## # Train # ############## for level in range(n_levels): # --------------- # Compile models # --------------- g_model = geom_model[level] m_model = morph_model[level] d_model = disc_model[level] mgd_model = gan_model[level] g_model.compile(loss='mse', optimizer=optim_g) m_model.compile(loss='mse', optimizer=optim_g) d_model.trainable = False mgd_model.compile(loss=models.wasserstein_loss, optimizer=optim_g) d_model.trainable = True d_model.compile(loss=models.wasserstein_loss, optimizer=optim_d) if verbose: print("") print(20 * "=") print("Level #{0}".format(level)) print(20 * "=") # ----------------- # Loop over epochs # ----------------- for e in range(n_epochs): batch_counter = 1 g_iters = 0 if verbose: print("") print(" Epoch #{0}".format(e)) print("") while batch_counter < n_batch_per_epoch: list_d_loss = list() list_g_loss = list() # ---------------------------- # Step 1: Train discriminator # ---------------------------- for d_iter in range(d_iters): # Clip discriminator weights d_model = clip_weights(d_model, weight_constraint) # Create a batch to feed the discriminator model X_locations_real, X_prufer_real = \ batch_utils.get_batch(training_data=training_data, batch_size=batch_size, batch_counter=batch_counter, n_nodes=n_nodes[level]) y_real = -np.ones((X_locations_real.shape[0], 1, 1)) #print X_locations_real.shape, X_prufer_real.shape, y_real.shape X_locations_gen, X_prufer_gen = \ batch_utils.gen_batch(batch_size=batch_size, n_nodes=n_nodes, level=level, input_dim=input_dim, geom_model=geom_model, cond_geom_model=cond_geom_model, morph_model=morph_model, cond_morph_model=cond_morph_model, conditioning_rule=rule) y_gen = np.ones((X_locations_gen.shape[0], 1, 1)) #print X_locations_gen.shape, X_prufer_gen.shape, y_gen.shape X_locations = np.concatenate( (X_locations_real, X_locations_gen), axis=0) X_prufer = np.concatenate((X_prufer_real, X_prufer_gen), axis=0) y = np.concatenate((y_real, y_gen), axis=0) # Update the discriminator #d_model.summary() disc_loss = \ d_model.train_on_batch([X_locations, X_prufer], y) list_d_loss.append(disc_loss) if verbose: print(" After {0} iterations".format(d_iters)) print(" Discriminator Loss \ = {0}".format(disc_loss)) # ------------------------- # Step 2: Train generators # ------------------------- # Freeze the discriminator d_model.trainable = False if level > 0: X_locations_prior_gen, X_prufer_prior_gen = \ batch_utils.gen_batch(batch_size=batch_size, n_nodes=n_nodes, level=level-1, input_dim=input_dim, geom_model=geom_model, cond_geom_model=cond_geom_model, morph_model=morph_model, cond_morph_model=cond_morph_model, conditioning_rule=rule) noise_input = np.random.randn(batch_size, 1, input_dim) if level == 0: gen_loss = \ mgd_model.train_on_batch([noise_input], y_real) else: gen_loss = \ mgd_model.train_on_batch([X_locations_prior_gen, X_prufer_prior_gen, noise_input], y_real) list_g_loss.append(gen_loss) if verbose: print("") print(" Generator_Loss: {0}".format(gen_loss)) # Unfreeze the discriminator d_model.trainable = True # --------------------- # Step 3: Housekeeping # --------------------- g_iters += 1 batch_counter += 1 # Save model weights (few times per epoch) print(batch_counter) if batch_counter % 25 == 0: #save_model_weights(g_model, # m_model, # level, # epoch, # batch_counter) if verbose: print(" Level #{0} Epoch #{1} Batch #{2}".format( level, e, batch_counter)) neuron_object = \ plot_example_neuron(X_locations_gen[0, :, :], X_prufer_gen[0, :, :]) plt.show() # Display loss trace if 0: plt.figure(figsize=(3, 2)) plt.plot(list_d_loss) plt.show() # Save models geom_model[level] = g_model cond_geom_model[level] = cg_model morph_model[level] = m_model cond_morph_model[level] = cm_model disc_model[level] = d_model gan_model[level] = mgd_model return geom_model, \ cond_geom_model, \ morph_model, \ cond_morph_model, \ disc_model, \ gan_model
def train_model(training_data=None, n_nodes=20, input_dim=100, n_epochs=25, batch_size=32, n_batch_per_epoch=100, d_iters=20, lr_discriminator=0.005, lr_generator=0.00005, d_weight_constraint=[-.03, .03], g_weight_constraint=[-.03, .03], m_weight_constraint=[-.03, .03], rule='none', train_loss='wasserstein_loss', verbose=True): """ Train the hierarchical model. Progressively generate trees with more and more nodes. Parameters ---------- training_data: dict of dicts each inner dict is an array 'geometry': 3-d arrays (locations) n_samples x n_nodes - 1 x 3 'morphology': 2-d arrays n_samples x n_nodes - 1 (parent sequences) example: training_data['geometry']['n20'][0:10, :, :] gives the geometry for the first 10 neurons training_data['geometry']['n20'][0:10, :] gives the parent sequences for the first 10 neurons here, 'n20' indexes a key corresponding to 20-node downsampled neurons. n_nodes: array specifies the number of nodes. input_dim: int dimensionality of noise input n_epochs: number of epochs over training data batch_size: batch size n_batch_per_epoch: int number of batches per epoch d_iters: int number of iterations to train discriminator lr_discriminator: float learning rate for optimization of discriminator lr_generator: float learning rate for optimization of generator weight_constraint: array upper and lower bounds of weights (to clip) verbose: bool print relevant progress throughout training Returns ------- geom_model: list of keras model objects geometry generators morph_model: list of keras model objects morphology generators disc_model: list of keras model objects discriminators gan_model: list of keras model objects discriminators stacked on generators """ # ################################### # Initialize models # ################################### geom_model = list() morph_model = list() disc_model = list() gan_model = list() # Discriminator d_model = models.discriminator(n_nodes=n_nodes, batch_size=batch_size, train_loss=train_loss) # Generators and GANs g_model, m_model = \ models.generator(n_nodes=n_nodes, batch_size=batch_size) stacked_model = \ models.discriminator_on_generators(g_model, m_model, d_model, conditioning_rule=rule, input_dim=input_dim, n_nodes=n_nodes) # Collect all models into a list disc_model.append(d_model) geom_model.append(g_model) morph_model.append(m_model) gan_model.append(stacked_model) # ############### # Optimizers # ############### optim_d = Adagrad() # RMSprop(lr=lr_discriminator) optim_g = Adagrad() # RMSprop(lr=lr_generator) # ############## # Train # ############## # --------------- # Compile models # --------------- g_model.compile(loss='mse', optimizer=optim_g) m_model.compile(loss='mse', optimizer=optim_g) d_model.trainable = False if train_loss == 'wasserstein_loss': stacked_model.compile(loss=models.wasserstein_loss, optimizer=optim_g) else: stacked_model.compile(loss='binary_crossentropy', optimizer=optim_g) d_model.trainable = True if train_loss == 'wasserstein_loss': d_model.compile(loss=models.wasserstein_loss, optimizer=optim_d) else: d_model.compile(loss='binary_crossentropy', optimizer=optim_d) if verbose: print("") print(20 * "=") # ----------------- # Loop over epochs # ----------------- for e in range(n_epochs): batch_counter = 1 g_iters = 0 if verbose: print("") print("Epoch #{0}".format(e)) print("") while batch_counter < n_batch_per_epoch: list_d_loss = list() list_g_loss = list() # ---------------------------- # Step 1: Train discriminator # ---------------------------- for d_iter in range(d_iters): # Clip discriminator weights d_model = clip_weights(d_model, d_weight_constraint) # Create a batch to feed the discriminator model select = range((batch_counter - 1) * batch_size, batch_counter * batch_size) X_locations_real = \ training_data['geometry']['n'+str(n_nodes)][select, :, :] X_locations_real = np.reshape(X_locations_real, [batch_size, (n_nodes - 1), 3]) X_parent_cut = \ np.reshape(training_data['morphology']['n'+str(n_nodes)][select, :], [1, (n_nodes - 1) * batch_size]) X_parent_real = \ batch_utils.get_batch(X_parent_cut=X_parent_cut, batch_size=batch_size, n_nodes=n_nodes) if train_loss == 'wasserstein_loss': y_real = -np.ones((X_locations_real.shape[0], 1, 1)) else: y_real = np.ones((X_locations_real.shape[0], 1, 1)) X_locations_gen, X_parent_gen = \ batch_utils.gen_batch(batch_size=batch_size, n_nodes=n_nodes, input_dim=input_dim, geom_model=g_model, morph_model=m_model, conditioning_rule=rule) if train_loss == 'wasserstein_loss': y_gen = np.ones((X_locations_gen.shape[0], 1, 1)) else: y_gen = np.zeros((X_locations_gen.shape[0], 1, 1)) # make data in half of real and generated cutting = int(batch_size / 2) X_locations_real_first_half = np.append( X_locations_real[:cutting, :, :], X_locations_gen[:cutting, :, :], axis=0) X_parent_real_first_half = np.append( X_parent_real[:cutting, :, :], X_parent_gen[:cutting, :, :], axis=0) y_real_first_half = np.append(y_real[:cutting, :, :], y_gen[:cutting, :, :], axis=0) X_locations_real_second_half = np.append( X_locations_real[cutting:, :, :], X_locations_gen[cutting:, :, :], axis=0) X_parent_real_second_half = np.append( X_parent_real[cutting:, :, :], X_parent_real[cutting:, :, :], axis=0) y_real_second_half = np.append(y_real[cutting:, :, :], y_gen[cutting:, :, :], axis=0) # Update the discriminator disc_loss = \ d_model.train_on_batch([X_locations_real_first_half, X_parent_real_first_half], y_real_first_half) list_d_loss.append(disc_loss) disc_loss = \ d_model.train_on_batch([X_locations_real_second_half, X_parent_real_second_half], y_real_second_half) list_d_loss.append(disc_loss) if verbose: print(" After {0} iterations".format(d_iters)) print(" Discriminator Loss \ = {0}".format(disc_loss)) # ------------------------------------- # Step 2: Train generators alternately # ------------------------------------- # Freeze the discriminator d_model.trainable = False noise_input = np.random.rand(batch_size, 1, input_dim) gen_loss = \ stacked_model.train_on_batch([noise_input], y_real) # Clip generator weights g_model = clip_weights(g_model, g_weight_constraint) m_model = clip_weights(m_model, m_weight_constraint) list_g_loss.append(gen_loss) if verbose: print("") print(" Generator_Loss: {0}".format(gen_loss)) # Unfreeze the discriminator d_model.trainable = True # --------------------- # Step 3: Housekeeping # --------------------- g_iters += 1 batch_counter += 1 # Save model weights (few times per epoch) print(batch_counter) if batch_counter % 2 == 0: #m_model = clip_weights(m_model, m_weight_constraint) #g_model = clip_weights(g_model, g_weight_constraint) if verbose: print(" Level #{0} Epoch #{1} Batch #{2}".format( 1, e, batch_counter)) neuron_object = \ plot_utils.plot_example_neuron_from_parent( X_locations_gen[0, :, :], X_parent_gen[0, :, :]) plt.plot(np.squeeze(X_locations_gen[0, :, :])) plot_utils.plot_adjacency(X_parent_real[0:2, :, :], X_parent_gen[0:2, :, :]) # Display loss trace #if verbose: plot_utils.plot_loss_trace(list_d_loss) # save the models save_model_weights(g_model, m_model, d_model, 0, e, batch_counter, list_d_loss, model_path_root='../model_weights') # Save models geom_model = g_model morph_model = m_model disc_model = d_model gan_model = stacked_model return geom_model, \ morph_model, \ disc_model, \ gan_model