def build_model(args: TrainArgs) -> MoleculeModel:
"""
Builds a MoleculeModel, which is a message passing neural network + feed-forward layers.
:param args: Arguments.
:return: A MoleculeModel containing the MPN encoder along with final linear layers with parameters initialized.
"""
output_size = args.num_tasks
args.output_size = output_size
if args.dataset_type == 'multiclass':
args.output_size *= args.multiclass_num_classes
model = MoleculeModel(
classification=args.dataset_type == 'classification',
multiclass=args.dataset_type == 'multiclass'
)
model.create_encoder(args)
model.create_ffn(args)
# Check hasattr for compatibility with loaded args from previous versions of chemprop
if hasattr(args, 'pytorch_seed') and args.pytorch_seed is not None:
torch.manual_seed(args.pytorch_seed)
initialize_weights(model)
return model
def __init__(self, args: TrainArgs):
"""
:param args: A :class:`~chemprop.args.TrainArgs` object containing model arguments.
"""
super(MoleculeModel, self).__init__()
self.classification = args.dataset_type == 'classification'
self.multiclass = args.dataset_type == 'multiclass'
# when using cross entropy losses, no sigmoid or softmax during training. But they are needed for mcc loss.
if self.classification or self.multiclass:
self.no_training_normalization = args.loss_function in [
'cross_entropy', 'binary_cross_entropy'
]
self.output_size = args.num_tasks
if self.multiclass:
self.output_size *= args.multiclass_num_classes
if self.classification:
self.sigmoid = nn.Sigmoid()
if self.multiclass:
self.multiclass_softmax = nn.Softmax(dim=2)
self.create_encoder(args)
self.create_ffn(args)
initialize_weights(self)
def build_model_autograd_a(args: Namespace) -> nn.Module:
"""
Builds a MoleculeModel, which is a message passing neural network + feed-forward layers.
:param args: Arguments.
:return: A MoleculeModel containing the MPN encoder along with final linear layers with parameters initialized.
"""
output_size = args.num_tasks
args.output_size = output_size
if args.dataset_type == 'multiclass':
args.output_size *= args.multiclass_num_classes
if args.epistemic == 'mc_dropout':
args.reg_acc = RegularizationAccumulator()
model = MoleculeModel_autograd_a(
classification=args.dataset_type == 'classification',
multiclass=args.dataset_type == 'multiclass',
aleatoric=args.aleatoric,
epistemic=args.epistemic)
model.create_encoder(args)
model.create_ffn(args)
initialize_weights(model)
if args.epistemic == 'mc_dropout':
args.reg_acc.initialize(cuda=args.cuda)
return model
def __init__(self, args: TrainArgs, featurizer: bool = False):
"""
Initializes the MoleculeModel.
:param args: Arguments.
:param featurizer: Whether the model should act as a featurizer, i.e. outputting
learned features in the final layer before prediction.
"""
super(MoleculeModel, self).__init__()
self.classification = args.dataset_type == 'classification'
self.multiclass = args.dataset_type == 'multiclass'
self.featurizer = featurizer
self.output_size = args.num_tasks
if self.multiclass:
self.output_size *= args.multiclass_num_classes
if self.classification:
self.sigmoid = nn.Sigmoid()
if self.multiclass:
self.multiclass_softmax = nn.Softmax(dim=2)
self.create_encoder(args)
self.create_ffn(args)
initialize_weights(self)
# create log noise parameter
# this must take place after initialize_weights
self.create_log_noise(args)
def __init__(self, args: TrainArgs, featurizer: bool = False):
"""
:param args: A :class:`~chemprop.args.TrainArgs` object containing model arguments.
:param featurizer: Whether the model should act as a featurizer, i.e., outputting the
learned features from the last layer prior to prediction rather than
outputting the actual property predictions.
"""
super(MoleculeModel, self).__init__()
self.classification = args.dataset_type == 'classification'
self.multiclass = args.dataset_type == 'multiclass'
self.featurizer = featurizer
try:
self.output_size = args.output_size
except:
self.output_size = args.num_tasks
if self.multiclass:
self.output_size *= args.multiclass_num_classes
if self.classification:
self.sigmoid = nn.Sigmoid()
if self.multiclass:
self.multiclass_softmax = nn.Softmax(dim=2)
self.create_encoder(args)
if type(args.ffn_hidden_size) is tuple:
self.create_ffn_from_tuple(args)
else:
self.create_ffn(args)
initialize_weights(self)
def build_model(args: Namespace) -> nn.Module:
"""
Builds a message passing neural network including final linear layers and initializes parameters.
:param args: Arguments.
:return: An nn.Module containing the MPN encoder along with final linear layers with parameters initialized.
"""
# Regression with binning
output_size = args.num_tasks # ? Is this width of output, i.e. number of classes? TODO:
encoder = MPNEncoder(
args, args.atom_fdim,
args.bond_fdim) # Obtain the message passing network component
first_linear_dim = args.hidden_size * (1 + args.jtnn
) # What exactly is jtnn? TODO
drop_layer = lambda p: nn.Dropout(p)
linear_layer = lambda input_dim, output_dim, p: nn.Linear(
input_dim, output_dim)
# Create FFN layers
if args.ffn_num_layers == 1:
ffn = [
drop_layer(args.ffn_input_dropout),
linear_layer(first_linear_dim, output_size, args.ffn_input_dropout)
]
else:
ffn = [
drop_layer(args.ffn_input_dropout),
linear_layer(first_linear_dim, args.ffn_hidden_size,
args.ffn_input_dropout)
]
for _ in range(args.ffn_num_layers - 2):
ffn.extend([
get_activation_function(args.activation),
drop_layer(args.ffn_dropout),
linear_layer(args.ffn_hidden_size, args.ffn_hidden_size,
args.ffn_dropout),
])
ffn.extend([
get_activation_function(args.activation),
drop_layer(args.ffn_dropout),
linear_layer(args.ffn_hidden_size, output_size, args.ffn_dropout),
])
# Classification
if args.dataset_type == 'classification':
ffn.append(nn.Sigmoid())
# Combined model
ffn = nn.Sequential(*ffn)
model = MoleculeModel(encoder, ffn)
initialize_weights(model, args)
return model
def prepare_model(args):
args.output_size = args.num_tasks
inv_model = MoleculeModel(
classification=args.dataset_type == 'classification',
multiclass=args.dataset_type == 'multiclass')
inv_model.create_encoder(
args) # phi(x), shared across source and target domain
inv_model.create_ffn(args) # source function
inv_model.src_ffn = inv_model.ffn
inv_model.create_ffn(args) # target function
initialize_weights(inv_model)
return inv_model.cuda()
def build_model(args: Namespace,
params: Dict[str, nn.Parameter] = None) -> nn.Module:
"""
Builds a message passing neural network including final linear layers and initializes parameters.
:param args: Arguments.
:param params: Parameters to use instead of creating parameters.
:return: An nn.Module containing the MPN encoder along with final linear layers with parameters initialized.
"""
# Regression with binning
if args.dataset_type == 'regression_with_binning':
output_size = args.num_bins * args.num_tasks
elif args.dataset_type == 'unsupervised':
output_size = args.unsupervised_n_clusters
elif args.dataset_type == 'kernel':
output_size = args.ffn_hidden_size # there will be another output layer later, for the pair of encodings
else:
output_size = args.num_tasks
if args.maml:
output_size = 1
args.output_size = output_size
if args.moe:
if params is not None:
raise ValueError(
'Setting parameters not yet supported for MOE or GAN')
model = MOE(args)
if args.adversarial:
model = GAN(args, prediction_model=model, encoder=model.encoder)
initialize_weights(model, args)
return model
model = MoleculeModel()
model.create_encoder(args, params=params)
model.create_ffn(args, params=params)
if args.adversarial:
if params is not None:
raise ValueError('Setting parameters not yet supported for GAN')
model = GAN(args, prediction_model=model, encoder=model.encoder)
initialize_weights(model, args)
return model
def build_model(args: Namespace) -> nn.Module:
"""
Builds a MoleculeModel, which is a message passing neural network + feed-forward layers.
:param args: Arguments.
:return: A MoleculeModel containing the MPN encoder along with final linear layers with parameters initialized.
"""
output_size = args.num_tasks
args.output_size = output_size
model = MoleculeModel(classification=args.dataset_type == 'classification')
model.create_encoder(args)
model.create_ffn(args)
initialize_weights(model)
return model
def __init__(self, args: TrainArgs, featurizer: bool = False):
"""
:param args: A :class:`~chemprop.args.TrainArgs` object containing model arguments.
:param featurizer: Whether the model should act as a featurizer, i.e., outputting the
learned features from the last layer prior to prediction rather than
outputting the actual property predictions.
"""
super(MoleculeModel, self).__init__()
self.classification = args.dataset_type == "classification"
self.multiclass = args.dataset_type == "multiclass"
self.featurizer = featurizer
self.device = args.device
self.output_size = args.num_tasks
if self.multiclass:
self.output_size *= args.multiclass_num_classes
if self.classification:
self.sigmoid = nn.Sigmoid()
self.ablation = args.ablation
self.ffn_input_size = (
(("lstm" not in self.ablation) *
(args.lstm_layers * args.lstm_hidden_size)
) # Output size of SMILES Encoder
+ (("fp_ffn" not in self.ablation) *
(args.fp_ffn_output_size)) # Output size of FP Encoder
+ (("mpnn" not in self.ablation) * (args.hidden_size))
# Output size of Graph Encoder
)
if self.multiclass:
self.multiclass_softmax = nn.Softmax(dim=2)
if "mpnn" not in self.ablation:
self.GraphEncoder = MPN(args=args)
if "lstm" not in self.ablation:
self.SMILESEncoder = SMILESEncoder(args=args)
if "fp_ffn" not in self.ablation:
self.FingerprintEncoder = create_FPEncoder(args=args)
self.FFNetwork = create_ffn(output_size=self.output_size,
input_size=self.ffn_input_size,
args=args)
initialize_weights(self)
def __init__(self, args: TrainArgs):
"""
:param args: A :class:`~chemprop.args.TrainArgs` object containing model arguments.
"""
super(MoleculeModel, self).__init__()
self.classification = args.dataset_type == 'classification'
self.multiclass = args.dataset_type == 'multiclass'
self.loss_function = args.loss_function
if hasattr(args, 'train_class_sizes'):
self.train_class_sizes = args.train_class_sizes
else:
self.train_class_sizes = None
# when using cross entropy losses, no sigmoid or softmax during training. But they are needed for mcc loss.
if self.classification or self.multiclass:
self.no_training_normalization = args.loss_function in [
'cross_entropy', 'binary_cross_entropy'
]
self.output_size = args.num_tasks
if self.multiclass:
self.output_size *= args.multiclass_num_classes
if self.loss_function == 'mve':
self.output_size *= 2 # return means and variances
if self.loss_function == 'dirichlet' and self.classification:
self.output_size *= 2 # return dirichlet parameters for positive and negative class
if self.loss_function == 'evidential':
self.output_size *= 4 # return four evidential parameters: gamma, lambda, alpha, beta
if self.classification:
self.sigmoid = nn.Sigmoid()
if self.multiclass:
self.multiclass_softmax = nn.Softmax(dim=2)
if self.loss_function in ['mve', 'evidential', 'dirichlet']:
self.softplus = nn.Softplus()
self.create_encoder(args)
self.create_ffn(args)
initialize_weights(self)
def build_model(args: Namespace) -> nn.Module:
"""
Builds a MoleculeModel, which is a message passing neural network + feed-forward layers.
:param args: Arguments.
:return: A MoleculeModel containing the MPN encoder along with final linear layers with parameters initialized.
"""
output_size = args.num_tasks
args.output_size = output_size
if args.dataset_type == 'multiclass':
raise NotImplementedError
model = MoleculeModel(output_raw=args.output_raw, use_cuda=args.cuda)
model.create_encoder(args)
model.create_ffn(args)
initialize_weights(model)
return model
def __init__(self, args: TrainArgs):
"""
:param args: A :class:`~chemprop.args.TrainArgs` object containing model arguments.
"""
super(MoleculeModel, self).__init__()
self.classification = args.dataset_type == 'classification'
self.multiclass = args.dataset_type == 'multiclass'
self.output_size = args.num_tasks
if self.multiclass:
self.output_size *= args.multiclass_num_classes
if self.classification:
self.sigmoid = nn.Sigmoid()
if self.multiclass:
self.multiclass_softmax = nn.Softmax(dim=2)
self.create_encoder(args)
self.create_ffn(args)
initialize_weights(self)
