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)