def __init__(self, args: Namespace, atom_fdim: int, bond_fdim: int):
"""Initializes the MPNEncoder.
:param args: Arguments.
:param atom_fdim: Atom features dimension.
:param bond_fdim: Bond features dimension.
"""
super(MPNEncoder, self).__init__()
self.atom_fdim = atom_fdim # 133
self.bond_fdim = bond_fdim # 147
self.hidden_size = args.hidden_size
self.bias = args.bias
self.depth = args.depth
self.dropout = args.dropout
self.layers_per_message = 1
self.undirected = args.undirected
self.atom_messages = args.atom_messages # Use messages on atoms instead of messages on bonds, default=False
self.features_only = args.features_only
self.use_input_features = args.use_input_features
self.epistemic = args.epistemic
self.mc_dropout = self.epistemic == 'mc_dropout'
self.args = args
self.features_generator = args.features_generator
if self.features_only or self.features_generator:
return
# Dropout
self.dropout_layer = nn.Dropout(p=self.dropout)
# Activation
self.act_func = get_activation_function(args.activation)
# Cached zeros
self.cached_zero_vector = nn.Parameter(torch.zeros(self.hidden_size), requires_grad=False)
# Concrete Dropout for Bayesian NN
wd, dd = get_cc_dropout_hyper(args.train_data_size, args.regularization_scale)
# Input
input_dim = self.atom_fdim if self.atom_messages else self.bond_fdim # 默認input dim -> bond_fdim 147
if self.mc_dropout:
self.W_i = ConcreteDropout(layer=nn.Linear(input_dim, self.hidden_size, bias=self.bias), reg_acc=args.reg_acc, weight_regularizer=wd, dropout_regularizer=dd)
else:
self.W_i = nn.Linear(input_dim, self.hidden_size, bias=self.bias) # in 147 out 1000 self.bias-> False (no bias)
if self.atom_messages:
w_h_input_size = self.hidden_size + self.bond_fdim
else:
w_h_input_size = self.hidden_size # 1000
# Shared weight matrix across depths (default)
if self.mc_dropout:
self.W_h = ConcreteDropout(layer=nn.Linear(w_h_input_size, self.hidden_size, bias=self.bias), reg_acc=args.reg_acc, weight_regularizer=wd, dropout_regularizer=dd, depth=self.depth - 1)
self.W_o = ConcreteDropout(layer=nn.Linear(self.atom_fdim + self.hidden_size, self.hidden_size), reg_acc=args.reg_acc, weight_regularizer=wd, dropout_regularizer=dd)
else:
self.W_h = nn.Linear(w_h_input_size, self.hidden_size, bias=self.bias) # in 1000 out 1000 (no bias)
self.W_o = nn.Linear(self.atom_fdim + self.hidden_size, self.hidden_size) # in 1000+133 out 1000 (with bias 1000)
def create_ffn(self, args: Namespace):
"""
Creates the feed-forward network for the model.
:param args: Arguments.
"""
self.multiclass = args.dataset_type == 'multiclass'
if self.multiclass:
self.num_classes = args.multiclass_num_classes
if args.features_only:
first_linear_dim = args.features_size
else:
first_linear_dim = args.hidden_size
if args.use_input_features:
first_linear_dim += args.features_dim
dropout = nn.Dropout(args.dropout)
activation = get_activation_function(args.activation)
wd, dd = get_cc_dropout_hyper(args.train_data_size, args.regularization_scale)
# Create FFN layers
if args.ffn_num_layers == 1:
ffn = [
dropout,
]
last_linear_dim = first_linear_dim
else:
ffn = [
dropout,
ConcreteDropout(layer=nn.Linear(first_linear_dim, args.ffn_hidden_size),
reg_acc=args.reg_acc, weight_regularizer=wd,
dropout_regularizer=dd) if self.mc_dropout else
nn.Linear(first_linear_dim, args.ffn_hidden_size)
]
for _ in range(args.ffn_num_layers - 2):
ffn.extend([
activation,
dropout,
ConcreteDropout(layer=nn.Linear(args.ffn_hidden_size, args.ffn_hidden_size),
reg_acc=args.reg_acc, weight_regularizer=wd,
dropout_regularizer=dd) if self.mc_dropout else
nn.Linear(args.ffn_hidden_size, args.ffn_hidden_size)
])
ffn.extend([
activation,
dropout,
])
last_linear_dim = args.ffn_hidden_size
# Create FFN model
self._ffn = nn.Sequential(*ffn)
if self.aleatoric:
self.output_layer = nn.Linear(last_linear_dim, args.output_size)
self.logvar_layer = nn.Linear(last_linear_dim, args.output_size)
else:
self.output_layer = nn.Linear(last_linear_dim, args.output_size)