def get_fitter(model, train_gen, valid_gen, loss_obj, fit_plot_prefix='', model_process_fun=None, lr=None, loss_display_cap=float('inf')): nice_params = filter(lambda p: p.requires_grad, model.parameters()) optimizer = optim.Adam(nice_params, lr=lr) scheduler = lr_scheduler.ReduceLROnPlateau( optimizer, patience=100) #.StepLR(optimizer, step_size=100, gamma=0.99) if dashboard is not None: index_to_lang_ordered = OrderedDict() for lang in languages: index_to_lang_ordered[meta['predefined'][lang]] = lang metric_monitor = MetricPlotter( plot_prefix=fit_plot_prefix, loss_display_cap=loss_display_cap, dashboard_name=dashboard, plot_ignore_initial=plot_ignore_initial, process_model_fun=model_process_fun, extra_metric_fun=partial(language_metrics_for_monitor, index_to_lang=index_to_lang_ordered), smooth_weight=0.9) else: metric_monitor = None checkpointer = Checkpointer(valid_batches_to_checkpoint=1, save_path=save_path, save_always=False) fitter = fit(train_gen=train_gen, valid_gen=valid_gen, model=model, optimizer=optimizer, scheduler=scheduler, epochs=EPOCHS, loss_fn=loss_obj, batches_to_valid=4, metric_monitor=metric_monitor, checkpointer=checkpointer) return fitter
def train_validity(grammar=True, model=None, EPOCHS=None, BATCH_SIZE=None, lr=2e-4, main_dataset=None, new_datasets=None, plot_ignore_initial=0, save_file=None, plot_prefix='', dashboard='main', preload_weights=False): root_location = os.path.dirname( os.path.abspath(inspect.getfile(inspect.currentframe()))) root_location = root_location + '/../' if save_file is not None: save_path = root_location + 'pretrained/' + save_file else: save_path = None molecules = True # checking for validity only makes sense for molecules settings = get_settings(molecules=molecules, grammar=grammar) # TODO: separate settings for this? if EPOCHS is not None: settings['EPOCHS'] = EPOCHS if BATCH_SIZE is not None: settings['BATCH_SIZE'] = BATCH_SIZE if preload_weights: try: model.load(save_path) except: pass nice_params = filter(lambda p: p.requires_grad, model.parameters()) optimizer = optim.Adam(nice_params, lr=lr) # create the composite loaders train_loader, valid_loader = train_valid_loaders(main_dataset, valid_fraction=0.1, batch_size=BATCH_SIZE, pin_memory=use_gpu) valid_smile_ds, invalid_smile_ds = new_datasets valid_train, valid_val = valid_smile_ds.get_train_valid_loaders(BATCH_SIZE) invalid_train, invalid_val = valid_smile_ds.get_train_valid_loaders( BATCH_SIZE) train_gen = MixedLoader(train_loader, valid_train, invalid_train) valid_gen = MixedLoader(valid_loader, valid_val, invalid_val) scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, factor=0.2, patience=3, min_lr=0.0001, eps=1e-08) #scheduler = lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.9) loss_obj = nn.BCELoss(size_average=True) fitter = fit(train_gen=train_gen, valid_gen=valid_gen, model=model, optimizer=optimizer, scheduler=scheduler, epochs=settings['EPOCHS'], loss_fn=loss_obj, save_path=save_path, dashboard_name=dashboard, plot_ignore_initial=plot_ignore_initial, plot_prefix=plot_prefix) return fitter
def train_mol_descriptor(grammar=True, EPOCHS=None, BATCH_SIZE=None, lr=2e-4, gradient_clip=5, drop_rate=0.0, plot_ignore_initial=0, save_file=None, preload_file=None, encoder_type='rnn', plot_prefix='', dashboard='properties', aux_dataset=None, preload_weights=False): root_location = os.path.dirname( os.path.abspath(inspect.getfile(inspect.currentframe()))) root_location = root_location + '/../' save_path = root_location + 'pretrained/' + save_file if preload_file is None: preload_path = save_path else: preload_path = root_location + 'pretrained/' + preload_file batch_mult = 1 if aux_dataset is None else 2 settings = get_settings(molecules=True, grammar=grammar) max_steps = settings['max_seq_length'] if EPOCHS is not None: settings['EPOCHS'] = EPOCHS if BATCH_SIZE is not None: settings['BATCH_SIZE'] = BATCH_SIZE if False: pre_model, _ = get_decoder(True, grammar, z_size=settings['z_size'], decoder_hidden_n=200, feature_len=settings['feature_len'], max_seq_length=max_steps, drop_rate=drop_rate, decoder_type=encoder_type, batch_size=BATCH_SIZE * batch_mult) class AttentionSimulator(nn.Module): def __init__(self, pre_model, drop_rate): super().__init__() self.pre_model = pre_model pre_model_2 = AttentionAggregatingHead(pre_model, drop_rate=drop_rate) pre_model_2.model_out_transform = lambda x: x[1] self.model = MeanVarianceSkewHead(pre_model_2, 4, drop_rate=drop_rate) def forward(self, x): self.pre_model.policy = PolicyFromTarget(x) return self.model(None) model = to_gpu(AttentionSimulator(pre_model, drop_rate=drop_rate)) else: pre_model = get_encoder(feature_len=settings['feature_len'], max_seq_length=settings['max_seq_length'], cnn_encoder_params={ 'kernel_sizes': (2, 3, 4), 'filters': (2, 3, 4), 'dense_size': 100 }, drop_rate=drop_rate, encoder_type=encoder_type) model = MeanVarianceSkewHead(pre_model, 4, drop_rate=drop_rate) nice_params = filter(lambda p: p.requires_grad, model.parameters()) optimizer = optim.Adam(nice_params, lr=lr) main_dataset = MultiDatasetFromHDF5(settings['data_path'], ['actions', 'smiles']) train_loader, valid_loader = train_valid_loaders(main_dataset, valid_fraction=0.1, batch_size=BATCH_SIZE, pin_memory=use_gpu) def scoring_fun(x): if isinstance(x, tuple) or isinstance(x, list): x = {'actions': x[0], 'smiles': x[1]} out_x = to_gpu(x['actions']) end_of_slice = randint(3, out_x.size()[1]) #TODO inject random slicing back out_x = out_x[:, 0:end_of_slice] smiles = x['smiles'] scores = to_gpu( torch.from_numpy(property_scorer(smiles).astype(np.float32))) return out_x, scores train_gen_main = IterableTransform(train_loader, scoring_fun) valid_gen_main = IterableTransform(valid_loader, scoring_fun) if aux_dataset is not None: train_aux, valid_aux = SamplingWrapper(aux_dataset) \ .get_train_valid_loaders(BATCH_SIZE, dataset_name=['actions', 'smiles']) train_gen_aux = IterableTransform(train_aux, scoring_fun) valid_gen_aux = IterableTransform(valid_aux, scoring_fun) train_gen = CombinedLoader([train_gen_main, train_gen_aux], num_batches=90) valid_gen = CombinedLoader([valid_gen_main, valid_gen_aux], num_batches=10) else: train_gen = train_gen_main #CombinedLoader([train_gen_main, train_gen_aux], num_batches=90) valid_gen = valid_gen_main #CombinedLoader([valid_gen_main, valid_gen_aux], num_batches=10) scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, factor=0.2, patience=3, min_lr=min(0.0001, 0.1 * lr), eps=1e-08) #scheduler = lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.9) loss_obj = VariationalLoss(['valid', 'logP', 'SA', 'cyc_sc']) metric_monitor = MetricPlotter(plot_prefix=plot_prefix, loss_display_cap=4.0, dashboard_name=dashboard, plot_ignore_initial=plot_ignore_initial) checkpointer = Checkpointer(valid_batches_to_checkpoint=10, save_path=save_path) fitter = fit(train_gen=train_gen, valid_gen=valid_gen, model=model, optimizer=optimizer, scheduler=scheduler, grad_clip=gradient_clip, epochs=settings['EPOCHS'], loss_fn=loss_obj, metric_monitor=metric_monitor, checkpointer=checkpointer) return model, fitter, main_dataset
def train_vae(molecules = True, grammar = True, EPOCHS = None, BATCH_SIZE = None, lr = 2e-4, drop_rate = 0.0, plot_ignore_initial = 0, reg_weight = 1, epsilon_std = 0.01, sample_z = True, save_file = None, preload_file = None, encoder_type='cnn', decoder_type='step', plot_prefix = '', dashboard = 'main', preload_weights=False): root_location = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))) root_location = root_location + '/../' save_path = root_location + 'pretrained/' + save_file if preload_file is None: preload_path = save_path else: preload_path = root_location + 'pretrained/' + preload_file settings = get_settings(molecules=molecules,grammar=grammar) if EPOCHS is not None: settings['EPOCHS'] = EPOCHS if BATCH_SIZE is not None: settings['BATCH_SIZE'] = BATCH_SIZE model,_ = get_vae(molecules=molecules, grammar=grammar, drop_rate=drop_rate, sample_z = sample_z, rnn_encoder=encoder_type, decoder_type = decoder_type, weights_file=preload_path if preload_weights else None, epsilon_std=epsilon_std ) nice_params = filter(lambda p: p.requires_grad, model.parameters()) optimizer = optim.Adam(nice_params, lr=lr) main_dataset = DatasetFromHDF5(settings['data_path'],'data') train_loader, valid_loader = train_valid_loaders(main_dataset, valid_fraction=0.1, batch_size=BATCH_SIZE, pin_memory=use_gpu) train_gen = TwinGenerator(train_loader) valid_gen = TwinGenerator(valid_loader) scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, factor=0.2, patience=3, min_lr=min(0.0001,0.1*lr), eps=1e-08) #scheduler = lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.9) loss_obj = VAELoss(settings['grammar'], sample_z, reg_weight) metric_monitor = MetricPlotter(plot_prefix=plot_prefix, loss_display_cap=4.0, dashboard_name=dashboard, plot_ignore_initial=plot_ignore_initial) checkpointer = Checkpointer(valid_batches_to_checkpoint=1, save_path=save_path) fitter = fit(train_gen=train_gen, valid_gen=valid_gen, model=model, optimizer=optimizer, scheduler=scheduler, epochs=settings['EPOCHS'], loss_fn=loss_obj, metric_monitor = metric_monitor, checkpointer = checkpointer) return model, fitter, main_dataset
def train_reinforcement(grammar=True, model=None, EPOCHS=None, BATCH_SIZE=None, lr=2e-4, main_dataset=None, new_datasets=None, plot_ignore_initial=0, save_file=None, plot_prefix='', dashboard='main', preload_weights=False): root_location = os.path.dirname( os.path.abspath(inspect.getfile(inspect.currentframe()))) root_location = root_location + '/../' if save_file is not None: save_path = root_location + 'pretrained/' + save_file else: save_path = None molecules = True # checking for validity only makes sense for molecules settings = get_settings(molecules=molecules, grammar=grammar) # TODO: separate settings for this? if EPOCHS is not None: settings['EPOCHS'] = EPOCHS if BATCH_SIZE is not None: settings['BATCH_SIZE'] = BATCH_SIZE if preload_weights: try: model.load(save_path) except: pass nice_params = filter(lambda p: p.requires_grad, model.parameters()) optimizer = optim.Adam(nice_params, lr=lr) # # create the composite loaders # train_loader, valid_loader = train_valid_loaders(main_dataset, # valid_fraction=0.1, # batch_size=BATCH_SIZE, # pin_memory=use_gpu) train_l = [] valid_l = [] for ds in new_datasets: train_loader, valid_loader = SamplingWrapper(ds)\ .get_train_valid_loaders(BATCH_SIZE, valid_batch_size = 1+int(BATCH_SIZE/5), dataset_name=['actions','seq_len','valid','sample_seq_ind'], window=1000) train_l.append(train_loader) valid_l.append(valid_loader) train_gen = CombinedLoader(train_l, num_batches=90) valid_gen = CombinedLoader(valid_l, num_batches=10) scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, factor=0.2, patience=3, min_lr=0.0001, eps=1e-08) #scheduler = lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.9) loss_obj = ReinforcementLoss() fitter = fit(train_gen=train_gen, valid_gen=valid_gen, model=model, optimizer=optimizer, scheduler=scheduler, epochs=settings['EPOCHS'], loss_fn=loss_obj, save_path=save_path, save_always=True, dashboard_name=dashboard, plot_ignore_initial=plot_ignore_initial, plot_prefix=plot_prefix, loss_display_cap=200) return fitter