def __init__(self, config):
"""Initialize configurations."""
self.config = config
# Data loader.
self.data = SparseMolecularDataset()
self.data.load(config.mol_data_dir)
# Model configurations.
self.z_dim = config.z_dim
self.m_dim = self.data.atom_num_types
self.b_dim = self.data.bond_num_types
self.g_conv_dim = config.g_conv_dim
self.d_conv_dim = config.d_conv_dim
self.g_repeat_num = config.g_repeat_num
self.d_repeat_num = config.d_repeat_num
self.lambda_cls = config.lambda_cls
self.lambda_rec = config.lambda_rec
self.lambda_gp = config.lambda_gp
self.post_method = config.post_method
self.metric = 'validity,sas'
# Training configurations.
self.batch_size = config.batch_size
self.num_iters = config.num_iters
self.num_iters_decay = config.num_iters_decay
self.g_lr = config.g_lr
self.d_lr = config.d_lr
self.dropout = config.dropout
self.n_critic = config.n_critic
self.beta1 = config.beta1
self.beta2 = config.beta2
self.resume_iters = config.resume_iters
# Test configurations.
self.test_iters = config.test_iters
# Miscellaneous.
self.use_tensorboard = config.use_tensorboard
self.device = torch.device(
'cuda' if torch.cuda.is_available() else 'cpu')
# Directories.
self.log_dir = config.log_dir
self.sample_dir = config.sample_dir
self.model_save_dir = config.model_save_dir
self.result_dir = config.result_dir
# Step size.
self.log_step = config.log_step
self.sample_step = config.sample_step
self.model_save_step = config.model_save_step
self.lr_update_step = config.lr_update_step
# Build the model and tensorboard.
self.build_model()
if self.use_tensorboard:
self.build_tensorboard()
def __init__(self, config, log=None):
"""Initialize configurations."""
# Log
self.log = log
# Data loader.
self.data = SparseMolecularDataset()
self.data.load(config.mol_data_dir)
# Model configurations.
self.z_dim = config.z_dim
self.m_dim = self.data.atom_num_types
self.b_dim = self.data.bond_num_types
self.f_dim = self.data.features
self.g_conv_dim = config.g_conv_dim
self.d_conv_dim = config.d_conv_dim
self.lambda_wgan = config.lambda_wgan
self.lambda_rec = config.lambda_rec
self.post_method = config.post_method
self.metric = 'validity,qed'
# Training configurations.
self.batch_size = config.batch_size
self.num_epochs = config.num_epochs
self.num_steps = (len(self.data) // self.batch_size)
self.g_lr = config.g_lr
self.d_lr = config.d_lr
self.dropout_rate = config.dropout
self.n_critic = config.n_critic
self.resume_epoch = config.resume_epoch
# Training or testing.
self.mode = config.mode
# Miscellaneous.
self.device = torch.device(
'cuda' if torch.cuda.is_available() else 'cpu')
print('Device: ', self.device)
# Directories.
self.log_dir_path = config.log_dir_path
self.model_dir_path = config.model_dir_path
self.img_dir_path = config.img_dir_path
# Step size.
self.model_save_step = config.model_save_step
# VAE KL weight.
self.kl_la = 1.
# Build the model.
self.build_model()
def __init__(self, data_dir):
"""Initialize and preprocess the CelebA dataset."""
self.data = SparseMolecularDataset()
self.data.load(data_dir)