def _create_layers(self): self.latent_dim = self.model_params["categ_encoding"]["num_dimensions"] model_func = lambda c_out : AutoregressiveLSTMModel(c_in=self.latent_dim, c_out=c_out, max_seq_len=self.max_seq_len, num_layers=self.model_params["coupling_hidden_layers"], hidden_size=self.model_params["coupling_hidden_size"], dp_rate=self.model_params["coupling_dropout"], input_dp_rate=self.model_params["coupling_input_dropout"]) self.model_params["categ_encoding"]["flow_config"]["model_func"] = model_func self.encoding_layer = create_encoding(self.model_params["categ_encoding"], dataset_class=self.dataset_class, vocab_size=self.vocab_size, vocab=self.vocab) num_flows = self.model_params["coupling_num_flows"] layers = [] for flow_index in range(num_flows): layers += [ActNormFlow(self.latent_dim)] if flow_index > 0: layers += [InvertibleConv(self.latent_dim)] layers += [ AutoregressiveMixtureCDFCoupling( c_in=self.latent_dim, model_func=model_func, block_type="LSTM model", num_mixtures=self.model_params["coupling_num_mixtures"]) ] self.flow_layers = nn.ModuleList([self.encoding_layer] + layers)
def _create_layers(self):
# Load global model params
self.num_node_types = self.dataset_class.num_node_types()
# Create essential modules of the flow
self.node_embed_flow = create_encoding(
self.model_params["categ_encoding"],
dataset_class=self.dataset_class,
vocab_size=self.num_node_types)
self.embed_dim = self.node_embed_flow.D
main_flow_layers = self._create_node_flow_layers()
self.flow_layers = nn.ModuleList([self.node_embed_flow] +
main_flow_layers)
def _create_encoding_layers(self):
self.node_encoding = create_encoding(
self.model_params["categ_encoding_nodes"],
dataset_class=self.dataset_class,
vocab_size=self.
num_node_types, # 9 for zinc250k dataset; not include the one for virtual nodes;
category_prior=self.dataset_class.get_node_prior(
data_root="data/"))
self.edge_attr_encoding = create_encoding(
self.model_params["categ_encoding_edges"],
dataset_class=self.dataset_class,
vocab_size=self.
num_edge_types, #3 for zinc250k dataset; # Removing the virtual edges here
category_prior=self.dataset_class.get_edge_prior(
data_root="data/"))
self.encoding_dim_nodes = self.node_encoding.D # seems like this is the dimension of node features;though in current case, encoding_dim_nodes is one, which is the node type
self.encoding_dim_edges = self.edge_attr_encoding.D # not very confident about this. will double check it
# Virtual edges are encoded by a single mixture
self.edge_virtual_encoding = LinearCategoricalEncoding(
num_dimensions=self.encoding_dim_edges,
flow_config={
"num_flows": self.model_params["encoding_virtual_num_flows"],
"hidden_layers": 2,
"hidden_size": 128
},
dataset_class=self.dataset_class,
vocab_size=1)
# Posterior needs to be a separate network as the true cannot be easily found.
self.edge_virtual_decoder = DecoderLinear(
num_categories=2,
embed_dim=self.encoding_dim_edges,
hidden_size=128,
num_layers=2,
class_prior_log=np.log(np.array(
[0.9,
0.1]))) # Molecules are sparse and usually have ~10% density
def _create_encoding_layers(self):
self.node_encoding = create_encoding(
self.model_params["categ_encoding_nodes"],
dataset_class=self.dataset_class,
vocab_size=self.num_node_types,
category_prior=self.dataset_class.get_node_prior(
data_root="data/"))
self.edge_attr_encoding = create_encoding(
self.model_params["categ_encoding_edges"],
dataset_class=self.dataset_class,
vocab_size=self.num_edge_types, # Removing the virtual edges here
category_prior=self.dataset_class.get_edge_prior(
data_root="data/"))
self.encoding_dim_nodes = self.node_encoding.D
self.encoding_dim_edges = self.edge_attr_encoding.D
# Virtual edges are encoded by a single mixture
self.edge_virtual_encoding = LinearCategoricalEncoding(
num_dimensions=self.encoding_dim_edges,
flow_config={
"num_flows": self.model_params["encoding_virtual_num_flows"],
"hidden_layers": 2,
"hidden_size": 128
},
dataset_class=self.dataset_class,
vocab_size=1)
# Posterior needs to be a separate network as the true cannot be easily found.
self.edge_virtual_decoder = DecoderLinear(
num_categories=2,
embed_dim=self.encoding_dim_edges,
hidden_size=128,
num_layers=2,
class_prior_log=np.log(np.array(
[0.9,
0.1]))) # Molecules are sparse and usually have ~10% density
def _create_layers(self):
self.latent_dim = self.model_params["categ_encoding"]["num_dimensions"]
model_func = lambda c_out: CouplingTransformerNet(
c_in=self.latent_dim,
c_out=c_out,
num_layers=self.model_params["coupling_hidden_layers"],
hidden_size=self.model_params["coupling_hidden_size"])
self.model_params["categ_encoding"]["flow_config"][
"model_func"] = model_func
self.model_params["categ_encoding"]["flow_config"][
"block_type"] = "Transformer"
self.encoding_layer = create_encoding(
self.model_params["categ_encoding"],
dataset_class=self.dataset_class,
vocab_size=self.vocab_size)
num_flows = self.model_params["coupling_num_flows"]
if self.latent_dim > 1:
coupling_mask = CouplingLayer.create_channel_mask(
self.latent_dim,
ratio=self.model_params["coupling_mask_ratio"])
coupling_mask_func = lambda flow_index: coupling_mask
else:
coupling_mask = CouplingLayer.create_chess_mask()
coupling_mask_func = lambda flow_index: coupling_mask if flow_index % 2 == 0 else 1 - coupling_mask
layers = []
for flow_index in range(num_flows):
layers += [
ActNormFlow(self.latent_dim),
InvertibleConv(self.latent_dim),
MixtureCDFCoupling(
c_in=self.latent_dim,
mask=coupling_mask_func(flow_index),
model_func=model_func,
block_type="Transformer",
num_mixtures=self.model_params["coupling_num_mixtures"])
]
self.flow_layers = nn.ModuleList([self.encoding_layer] + layers)