def fit(self, model, train_data, val_data=None):
logger = Logger() if self.config.log_file is not None else None
train_loader = self.get_dataloader(model, train_data, shuffle=True)
val_loader = None if val_data is None else self.get_dataloader(
model, val_data, shuffle=False)
self._train(model, train_loader, val_loader, logger)
return model
def fit(self, model, train_data, val_data=None):
logger = Logger() if self.config.log_file is not None else None
# Generator
gen_collate_fn = self.generator_collate_fn(model)
gen_train_loader = self.get_dataloader(model,
train_data,
gen_collate_fn,
shuffle=True)
gen_val_loader = None if val_data is None else self.get_dataloader(
model, val_data, gen_collate_fn, shuffle=False)
self._pretrain_generator(model, gen_train_loader, gen_val_loader,
logger)
# Discriminator
dsc_collate_fn = self.discriminator_collate_fn(model)
dsc_train_loader = self.get_dataloader(model,
train_data,
dsc_collate_fn,
shuffle=True)
dsc_val_loader = None if val_data is None else self.get_dataloader(
model, val_data, dsc_collate_fn, shuffle=False)
self._pretrain_discriminator(model, dsc_train_loader, dsc_val_loader,
logger)
# Policy gradient
self.ref_smiles, self.ref_mols = None, None
if model.metrics_reward is not None:
(self.ref_smiles, self.ref_mols
) = model.metrics_reward.get_reference_data(train_data)
pg_train_loader = dsc_train_loader
self._train_policy_gradient(model, pg_train_loader, logger)
del self.ref_smiles
del self.ref_mols
return model
def _train(self, model, train_loader, val_loader=None):
criterions = {
'autoencoder': nn.CrossEntropyLoss(),
'generator': lambda t: -torch.mean(F.logsigmoid(t)),
'discriminator': nn.BCEWithLogitsLoss()
}
optimizers = {
'autoencoder':
torch.optim.Adam(list(model.encoder.parameters()) +
list(model.decoder.parameters()),
lr=self.config.lr),
'generator':
torch.optim.Adam(model.encoder.parameters(), lr=self.config.lr),
'discriminator':
torch.optim.Adam(model.discriminator.parameters(),
lr=self.config.lr)
}
schedulers = {
k: torch.optim.lr_scheduler.StepLR(v, self.config.step_size,
self.config.gamma)
for k, v in optimizers.items()
}
device = torch.device(self.config.device)
log = Logger()
for epoch in range(self.config.train_epochs):
tqdm_data = tqdm(train_loader,
desc='Training (epoch #{})'.format(epoch))
for scheduler in schedulers.values():
scheduler.step()
log.append(
self._train_epoch(model, tqdm_data, criterions, optimizers))
log.write(self.config.log_file)
if val_loader is not None:
tqdm_data = tqdm(val_loader,
desc='Validation (epoch #{})'.format(epoch))
self._train_epoch(model, tqdm_data, criterions)
if epoch % self.config.save_frequency == 0:
model.to('cpu')
torch.save(
model.state_dict(),
self.config.model_save[:-3] + '_{0:03d}.pt'.format(epoch))
model.to(device)
def fit(self, model, data):
def get_params():
return (p for p in model.parameters() if p.requires_grad)
model.train()
log = Logger()
n_epoch = self.config.num_epochs
optimizer = optim.Adam(get_params(), lr=self.config.lr)
for epoch in range(n_epoch):
if epoch < self.config.kl_start:
kl_w = 0
else:
kl_w = self.config.kl_w
word_acc, topo_acc, assm_acc, steo_acc, all_kl = 0, 0, 0, 0, 0
with tqdm.tqdm(data) as train_dataloader:
train_dataloader.set_description('Train (epoch #{})'.format(epoch))
for it, batch in enumerate(train_dataloader):
model.zero_grad()
loss, kl_div, wacc, tacc, sacc, dacc = model(batch, kl_w)
loss.backward()
optimizer.step()
word_acc += wacc
topo_acc += tacc
assm_acc += sacc
steo_acc += dacc
all_kl += kl_div
postfix = {'kl': all_kl / (it + 1),
'word': word_acc / (it + 1) * 100,
'topo': topo_acc / (it + 1) * 100,
'assm': assm_acc / (it + 1) * 100,
'steo': steo_acc / (it + 1) * 100}
train_dataloader.set_postfix(postfix)
log.append(postfix)
log.save(self.config.log_file)
if epoch % self.config.save_frequency == 0:
model.to('cpu')
torch.save(model.state_dict(), self.config.model_save[:-3]+'_{0:03d}.pt'.format(epoch))
model.to(device)
def fit(self, model, data):
def get_params():
return (p for p in model.parameters() if p.requires_grad)
if isinstance(data, tuple):
train_dataloader = data[0]
val_dataloader = data[1]
else:
train_dataloader = data
val_dataloader = None
num_epochs = self.config.num_epochs
device = torch.device(self.config.device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(get_params(), lr=self.config.lr)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
self.config.step_size,
self.config.gamma)
elog = Logger()
for epoch in range(num_epochs):
scheduler.step()
model.train()
train_dataloader = tqdm.tqdm(train_dataloader)
train_dataloader.set_description('Train (epoch #{})'.format(epoch))
loss = self._pass_data(model, train_dataloader, criterion,
optimizer)
elog.append({'loss': loss})
if val_dataloader is not None:
val_dataloader = tqdm.tqdm(val_dataloader)
val_dataloader.set_description(
'Validation (epoch #{})'.format(epoch))
self._pass_data(model, val_dataloader, criterion)
if epoch % self.config.save_frequency == 0:
model.to('cpu')
torch.save(
model.state_dict(),
self.config.model_save[:-3] + '_{0:03d}.pt'.format(epoch))
model.to(device)
elog.save(self.config.log_file)
torch.save(model.state_dict(), self.config.model_save)
############ load training data
print("Loading training set...")
train_split = DatasetSplit("train", config.train_load)
train_dataloader = train_split.get_dataloader(batch_size=config.train_bsz)
if config.valid_load is not None:
print("Loading validation set...")
valid_split = DatasetSplit("valid", config.valid_load)
valid_dataloader = valid_split.get_dataloader(batch_size=config.train_bsz,
shuffle=False)
vocab = train_split._vocab
############ get model and train
print("Initializing model...")
model = LVAE(vocab, config)
## (optional) load trained model
if config.model_load is not None:
model.load_state_dict(torch.load(config.train_from))
model.to(device)
## log training process to csv file
logger = Logger() if config.log_path is not None else None
print("Start training...")
train(model, config, train_dataloader, valid_dataloader, logger)
def fit(self, model, content_train, style_instance, content_test=None):
self.encoder_c = model.encoder_c
self.encoder_s1 = model.encoder_s1
self.encoder_s2 = model.encoder_s2
self.decoder = model.decoder
self.discriminator = model.Discriminator
self.heteroencoder = model.heteroencoder
cuda = True if torch.cuda.is_available() else False
if cuda:
print('Using CUDA')
self.discriminator.cuda()
self.encoder_c.cuda()
self.encoder_s1.cuda()
self.encoder_s2.cuda()
self.decoder.cuda()
self.heteroencoder.cuda()
else:
print('Using CPU')
logger = Logger() if self.config.log_file is not None else None
_, smi2vec = load_model()
print("Training GAN.")
def get_latent(smiles):
vec = smi2vec(smiles)
latent = self.heteroencoder.encode(vec)
latent = latent.reshape(latent.shape[0], self.latent_size)
return latent
os.makedirs('./latent_save', exist_ok=True)
if not os.path.exists(
f'./latent_save/latent_content_train_{self.config.target}.pkl'
):
latent_content_train = get_latent(content_train)
latent_style_instance = get_latent(style_instance)
pickle.dump(
latent_content_train,
open(
f'./latent_save/latent_content_train_{self.config.target}.pkl',
'wb'))
pickle.dump(
latent_style_instance,
open(
f'./latent_save/latent_style_instance_{self.config.target}.pkl',
'wb'))
if content_test is not None:
latent_content_test = get_latent(content_test)
pickle.dump(
latent_content_test,
open(
f'./latent_save/latent_content_test_{self.config.target}.pkl',
'wb'))
else:
latent_content_train = pickle.load(
open(
f'./latent_save/latent_content_train_{self.config.target}.pkl',
'rb'))
latent_style_instance = pickle.load(
open(
f'./latent_save/latent_style_instance_{self.config.target}.pkl',
'rb'))
if content_test is not None:
latent_content_test = pickle.load(
open(
f'./latent_save/latent_content_test_{self.config.target}.pkl',
'rb'))
train_loader = self.get_dataloader(model,
LatentMolsDataset(
latent_content_train,
latent_style_instance,
is_train=True),
shuffle=True)
val_loader = None if content_test is None else self.get_dataloader(
model,
LatentMolsDataset(latent_content_test, latent_style_instance),
shuffle=False)
self._train(model, train_loader, val_loader, logger)
return model
def fit(self,
model,
train_data,
val_data=None):
from ddc_pub import ddc_v3 as ddc
self.generator = model.Generator
self.discriminator = model.Discriminator
cuda = True if torch.cuda.is_available() else False
if cuda:
self.discriminator.cuda()
self.generator.cuda()
logger = Logger() if self.config.log_file is not None else None
if self.config.heteroencoder_version == 'new':
# Train the heteroencoder first
print("Training heteroencoder.")
currentDirectory = os.getcwd()
path = '{}/moses/latentgan/heteroencoder_models/new_model' \
.format(currentDirectory)
encoder_checkpoint_path = \
'{}/moses/latentgan/heteroencoder_models/checkpoints/' \
.format(currentDirectory)
# Convert all SMILES to binary RDKit mols to be
# compatible with the heteroencoder
heteroencoder_mols = [Chem.rdchem.Mol
.ToBinary(Chem.MolFromSmiles(smiles))
for smiles in train_data]
# Dataset information
dataset_info = self._get_dataset_info(
train_data, name="heteroencoder_train_data")
# Initialize heteroencoder with default parameters
heteroencoder_model = ddc.DDC(x=np.array(heteroencoder_mols),
y=np.array(heteroencoder_mols),
dataset_info=dataset_info,
scaling=False,
noise_std=self.config.
heteroencoder_noise_std,
lstm_dim=self.config.
heteroencoder_layer_dim,
dec_layers=self.config.
heteroencoder_dec_layers,
td_dense_dim=0,
batch_size=self.config.
heteroencoder_batch_size,
codelayer_dim=self.latent_size)
# Train heteroencoder
heteroencoder_model.fit(epochs=self.config.heteroencoder_epochs,
lr=self.config.heteroencoder_lr,
model_name="new_model",
mini_epochs=self.config.
heteroencoder_mini_epochs,
patience=self.config.
heteroencoder_patience,
save_period=self.config.
heteroencoder_save_period,
checkpoint_dir=encoder_checkpoint_path,
gpus=1,
use_multiprocessing=False,
workers=1,
lr_decay=self.config.
heteroencoder_lr_decay,
sch_epoch_to_start=self.config.
heteroencoder_lr_decay_start)
heteroencoder_model.save(path)
heteroencoder = load_model(
model_version=self.config.heteroencoder_version)
print("Training GAN.")
mols_in = [Chem.rdchem.Mol.ToBinary(
Chem.MolFromSmiles(smiles)) for smiles in train_data]
latent_train = heteroencoder.transform(
heteroencoder.vectorize(mols_in))
# Now encode the GAN training set to latent vectors
latent_train = latent_train.reshape(latent_train.shape[0],
self.latent_size)
if val_data is not None:
mols_val = [Chem.rdchem.Mol.ToBinary(Chem.MolFromSmiles(smiles))
for smiles in val_data]
latent_val = heteroencoder.transform(
heteroencoder.vectorize(mols_val))
latent_val = latent_val.reshape(latent_val.shape[0],
self.latent_size)
train_loader = self.get_dataloader(model,
LatentMolsDataset(latent_train),
shuffle=True)
val_loader = None if val_data is None else self.get_dataloader(
model, LatentMolsDataset(latent_val), shuffle=False
)
self._train(model, train_loader, val_loader, logger)
return model
def fit(self, model, data):
def get_params():
return (p for p in model.vae.parameters() if p.requires_grad)
model.train()
n_epoch = self._n_epoch()
kl_annealer = KLAnnealer(n_epoch, self.config)
optimizer = optim.Adam(get_params(), lr=self.config.lr_start)
lr_annealer = CosineAnnealingLRWithRestart(optimizer, self.config)
device = torch.device(self.config.device)
n_last = self.config.n_last
elog, ilog = Logger(), Logger()
for epoch in range(n_epoch):
# Epoch start
kl_weight = kl_annealer(epoch)
# Iters
T = tqdm.tqdm(data)
for i, x in enumerate(T):
# Forward
kl_loss, recon_loss = model(x)
loss = kl_weight * kl_loss + recon_loss
# Backward
optimizer.zero_grad()
loss.backward()
clip_grad_norm_(get_params(), self.config.grad_clipping)
optimizer.step()
# Log
lr = optimizer.param_groups[0]['lr']
ilog.append({
'epoch': epoch,
'kl_loss': kl_loss.item(),
'recon_loss': recon_loss.item(),
'loss': loss.item(),
'kl_weight': kl_weight,
'lr': lr
})
# Update T
kl_loss_value = np.mean(ilog['kl_loss'][-n_last:])
recon_loss_value = np.mean(ilog['recon_loss'][-n_last:])
loss_value = np.mean(ilog['loss'][-n_last:])
postfix = [
f'loss={loss_value:.5f}', f'(kl={kl_loss_value:.5f}',
f'recon={recon_loss_value:.5f})',
f'klw={kl_weight:.5f} lr={lr:.5f}'
]
T.set_postfix_str(' '.join(postfix))
T.set_description(f'Train (epoch #{epoch})')
T.refresh()
# Log
elog.append({
**{k: v
for k, v in ilog[-1].items() if 'loss' not in k}, 'kl_loss':
kl_loss_value,
'recon_loss': recon_loss_value,
'loss': loss_value
})
# Save model at each epoch
if epoch % self.config.save_frequency == 0:
model.to('cpu')
torch.save(
model.state_dict(),
self.config.model_save[:-3] + '_{0:03d}.pt'.format(epoch))
model.to(device)
elog.save(self.config.log_file)
# Epoch end
lr_annealer.step()
torch.save(model.state_dict(), self.config.model_save)
return elog, ilog
def fit(self, model, train_data, fps_data, val_data=None):
log = Logger()
self._pretrain_generator(model, train_data, fps_data, log, val_data)
self._pretrain_discriminator(model, train_data, fps_data, log,
val_data)
self._train_policy_gradient(model, train_data, fps_data, log)