merge & squash to refresh branch

docs/implementation_notes.md

@@ -33,4 +33,4 @@ The code contains a specific categorical distribution type for graph actions, `G
 
 Consider for example the `AddNode` and `SetEdgeAttr` actions, one applies to nodes and one to edges. An efficient way to produce logits for these actions would be to take the node/edge embeddings and project them (e.g. via an MLP) to a `(n_nodes, n_node_actions)` and `(n_edges, n_edge_actions)` tensor respectively. We thus obtain a list of tensors representing the logits of different actions, but logits are mixed between graphs in the minibatch, so one cannot simply apply a `softmax` operator on the tensor. 
 
-The `GraphActionCategorical` class handles this and can be used to compute various other things, such as entropy, log probabilities, and so on; it can also be used to sample from the distribution.
+The `GraphActionCategorical` class handles this and can be used to compute various other things, such as entropy, log probabilities, and so on; it can also be used to sample from the distribution.

src/gflownet/algo/graph_sampling.py

@@ -113,7 +113,9 @@ def not_done(lst):
  ]
  if self.sample_temp != 1:
  sample_cat = copy.copy(fwd_cat)
- sample_cat.logits = [i / self.sample_temp for i in fwd_cat.logits]
+ sample_cat.logits = [
+ i * m / self.sample_temp - 1000 * (1 - m) for i, m in zip(fwd_cat.logits, fwd_cat.masks)
+ ]
  actions = sample_cat.sample()
  else:
  actions = fwd_cat.sample()

src/gflownet/algo/trajectory_balance.py

@@ -6,7 +6,7 @@
 import torch.nn as nn
 import torch_geometric.data as gd
 from torch import Tensor
-from torch_scatter import scatter, scatter_sum
+from torch_scatter import scatter, scatter_sum, scatter_logsumexp
 
 from gflownet.algo.graph_sampling import GraphSampler
 from gflownet.config import Config
@@ -309,22 +309,15 @@ def compute_batch_losses(
  # Indicate that the `batch` corresponding to each action is the above
  ip_log_prob = fwd_cat.log_prob(batch.ip_actions, batch=ip_batch_idces)
  # take the logsumexp (because we want to sum probabilities, not log probabilities)
- # TODO: numerically stable version:
- p = scatter(ip_log_prob.exp(), ip_batch_idces, dim=0, dim_size=batch_idx.shape[0], reduce="sum")
- # As a (reasonable) band-aid, ignore p < 1e-30, this will prevent underflows due to
- # scatter(small number) = 0 on CUDA
- log_p_F = p.clamp(1e-30).log()
+ log_p_F = scatter_logsumexp(ip_log_prob, ip_batch_idces, dim=0, dim_size=batch_idx.shape[0])
 
  if self.cfg.do_parameterize_p_b:
  # Now we repeat this but for the backward policy
  bck_ip_batch_idces = torch.arange(batch.bck_ip_lens.shape[0], device=dev).repeat_interleave(
  batch.bck_ip_lens
  )
  bck_ip_log_prob = bck_cat.log_prob(batch.bck_ip_actions, batch=bck_ip_batch_idces)
- bck_p = scatter(
- bck_ip_log_prob.exp(), bck_ip_batch_idces, dim=0, dim_size=batch_idx.shape[0], reduce="sum"
- )
- log_p_B = bck_p.clamp(1e-30).log()
+ log_p_B = scatter_logsumexp(bck_ip_log_prob, bck_ip_batch_idces, dim=0, dim_size=batch_idx.shape[0])
  else:
  # Else just naively take the logprob of the actions we took
  log_p_F = fwd_cat.log_prob(batch.actions)
@@ -496,7 +489,7 @@ def subtb_loss_fast(self, P_F, P_B, F, R, traj_lengths):
  cumul_lens = torch.cumsum(torch.cat([torch.zeros(1, device=dev), traj_lengths]), 0).long()
  total_loss = torch.zeros(num_trajs, device=dev)
  ar = torch.arange(max_len, device=dev)
- car = torch.cumsum(ar, 0)
+ car = torch.cumsum(ar, 0) if self.length_normalize_losses else torch.ones_like(ar)
  F_and_R = torch.cat([F, R])
  R_start = F.shape[0]
  for ep in range(traj_lengths.shape[0]):

src/gflownet/config.py

@@ -81,6 +81,7 @@ class Config:
  """
 
  log_dir: str = MISSING
+ log_sampled_data: bool = True
  device: str = "cuda"
  seed: int = 0
  validate_every: int = 1000

src/gflownet/data/sampling_iterator.py

@@ -1,5 +1,6 @@
 import os
 import sqlite3
+import traceback
 from collections.abc import Iterable
 from copy import deepcopy
 from typing import Callable, List
@@ -11,6 +12,7 @@
 from rdkit import Chem, RDLogger
 from torch.utils.data import Dataset, IterableDataset
 
+from gflownet.config import Config
 from gflownet.data.replay_buffer import ReplayBuffer
 from gflownet.envs.graph_building_env import GraphActionCategorical
 
@@ -112,9 +114,14 @@ def __init__(
  # don't want to initialize per-worker things just yet, such as where the log the worker writes
  # to. This must be done in __iter__, which is called by the DataLoader once this instance
  # has been copied into a new python process.
- self.log_dir = log_dir
+ import warnings
+
+ warnings.warn("Fix dependency on cfg.log_sampled_data")
+ self.log_dir = log_dir # if cfg.log_sampled_data else None
  self.log = SQLiteLog()
  self.log_hooks: List[Callable] = []
+ # TODO: make this a proper flag / make a separate class for logging sampled molecules to a SQLite db
+ self.log_molecule_smis = not hasattr(self.ctx, "not_a_molecule_env") and self.log_dir is not None
 
  def add_log_hook(self, hook: Callable):
  self.log_hooks.append(hook)
@@ -158,6 +165,14 @@ def __len__(self):
  return len(self.data)
 
  def __iter__(self):
+ try:
+ for x in self.iterator():
+ yield x
+ except Exception as e:
+ traceback.print_exc()
+ raise e
+
+ def iterator(self):
  worker_info = torch.utils.data.get_worker_info()
  self._wid = worker_info.id if worker_info is not None else 0
  # Now that we know we are in a worker instance, we can initialize per-worker things
@@ -189,9 +204,7 @@ def __iter__(self):
  else: # If we're not sampling the conditionals, then the idcs refer to listed preferences
  num_online = num_offline
  num_offline = 0
- cond_info = self.task.encode_conditional_information(
- steer_info=torch.stack([self.data[i] for i in idcs])
- )
+ cond_info = self.task.encode_conditional_information(torch.stack([self.data[i] for i in idcs]))
  trajs, flat_rewards = [], []
 
  # Sample some on-policy data
@@ -250,14 +263,16 @@ def __iter__(self):
  # note: we convert back into natural rewards for logging purposes
  # (allows to take averages and plot in objective space)
  # TODO: implement that per-task (in case they don't apply the same beta and log transformations)
- rewards = torch.exp(log_rewards / cond_info["beta"])
+ rewards = torch.exp(log_rewards / (cond_info["beta"] if "beta" in cond_info else 1.0))
  if num_online > 0 and self.log_dir is not None:
  self.log_generated(
  deepcopy(trajs[num_offline:]),
  deepcopy(rewards[num_offline:]),
  deepcopy(flat_rewards[num_offline:]),
  {k: v[num_offline:] for k, v in deepcopy(cond_info).items()},
  )
+
+ extra_info = {}
  if num_online > 0:
  for hook in self.log_hooks:
  extra_info.update(

src/gflownet/envs/basic_graph_ctx.py

@@ -0,0 +1,198 @@
+from typing import Dict, List, Tuple
+
+import networkx as nx
+import torch
+import torch_geometric.data as gd
+from networkx.algorithms.isomorphism import is_isomorphic as nx_is_isomorphic
+
+from gflownet.envs.graph_building_env import (
+ Graph,
+ GraphAction,
+ GraphActionType,
+ GraphBuildingEnvContext,
+ graph_without_edge,
+)
+from gflownet.utils.graphs import random_walk_probs
+
+
+def hashg(g):
+ return nx.algorithms.graph_hashing.weisfeiler_lehman_graph_hash(g, node_attr="v")
+
+
+def is_isomorphic(u, v):
+ return nx_is_isomorphic(u, v, lambda a, b: a == b, lambda a, b: a == b)
+
+
+class BasicGraphContext(GraphBuildingEnvContext):
+ """
+ A basic graph generation context.
+
+ This simple environment context is designed to be used to test implementations. It only allows for AddNode and
+ AddEdge actions, and is meant to be used within the BasicGraphTask to generate graphs of up to 7 nodes with
+ only two possible node attributes, making the state space a total of ~70k states (which is nicely enumerable
+ and allows us to compute p_theta(x) exactly for all x in the state space).
+ """
+
+ def __init__(self, max_nodes=7, num_cond_dim=0, graph_data=None, output_gid=False):
+ self.max_nodes = max_nodes
+ self.output_gid = output_gid
+
+ self.node_attr_values = {
+ "v": [0, 1], # Imagine this is as colors
+ }
+ self._num_rw_feat = 8
+
+ self.num_new_node_values = len(self.node_attr_values["v"])
+ self.num_node_attr_logits = None
+ self.num_node_dim = self.num_new_node_values + 1 + self._num_rw_feat
+ self.num_node_attrs = 1
+ self.num_edge_attr_logits = None
+ self.num_edge_attrs = 0
+ self.num_cond_dim = num_cond_dim
+ self.num_edge_dim = 1
+ self.edges_are_duplicated = True
+ self.edges_are_unordered = True
+
+ # Order in which models have to output logits
+ self.action_type_order = [
+ GraphActionType.Stop,
+ GraphActionType.AddNode,
+ GraphActionType.AddEdge,
+ ]
+ self.bck_action_type_order = [
+ GraphActionType.RemoveNode,
+ GraphActionType.RemoveEdge,
+ ]
+ self.device = torch.device("cpu")
+ self.graph_data = graph_data
+ self.hash_to_graphs: Dict[str, int] = {}
+ if graph_data is not None:
+ states_hash = [hashg(i) for i in graph_data]
+ for i, h, g in zip(range(len(graph_data)), states_hash, graph_data):
+ self.hash_to_graphs[h] = self.hash_to_graphs.get(h, list()) + [(g, i)]
+
+ def get_graph_idx(self, g, default=None):
+ h = hashg(g)
+ if h not in self.hash_to_graphs and default is not None:
+ return default
+ bucket = self.hash_to_graphs[h]
+ if len(bucket) == 1:
+ return bucket[0][1]
+ for i in bucket:
+ if is_isomorphic(i[0], g):
+ return i[1]
+ if default is not None:
+ return default
+ raise ValueError(g)
+
+ def aidx_to_GraphAction(self, g: gd.Data, action_idx: Tuple[int, int, int], fwd: bool = True):
+ """Translate an action index (e.g. from a GraphActionCategorical) to a GraphAction"""
+ act_type, act_row, act_col = [int(i) for i in action_idx]
+ if fwd:
+ t = self.action_type_order[act_type]
+ else:
+ t = self.bck_action_type_order[act_type]
+
+ if t is GraphActionType.Stop:
+ return GraphAction(t)
+ elif t is GraphActionType.AddNode:
+ return GraphAction(t, source=act_row, value=self.node_attr_values["v"][act_col])
+ elif t is GraphActionType.AddEdge:
+ a, b = g.non_edge_index[:, act_row]
+ return GraphAction(t, source=a.item(), target=b.item())
+ elif t is GraphActionType.RemoveNode:
+ return GraphAction(t, source=act_row)
+ elif t is GraphActionType.RemoveEdge:
+ a, b = g.edge_index[:, act_row * 2]
+ return GraphAction(t, source=a.item(), target=b.item())
+
+ def GraphAction_to_aidx(self, g: gd.Data, action: GraphAction) -> Tuple[int, int, int]:
+ """Translate a GraphAction to an index tuple"""
+ if action.action is GraphActionType.Stop:
+ row = col = 0
+ type_idx = self.action_type_order.index(action.action)
+ elif action.action is GraphActionType.AddNode:
+ row = action.source
+ col = self.node_attr_values["v"].index(action.value)
+ type_idx = self.action_type_order.index(action.action)
+ elif action.action is GraphActionType.AddEdge:
+ # Here we have to retrieve the index in non_edge_index of an edge (s,t)
+ # that's also possibly in the reverse order (t,s).
+ # That's definitely not too efficient, can we do better?
+ row = (
+ (g.non_edge_index.T == torch.tensor([(action.source, action.target)])).prod(1)
+ + (g.non_edge_index.T == torch.tensor([(action.target, action.source)])).prod(1)
+ ).argmax()
+ col = 0
+ type_idx = self.action_type_order.index(action.action)
+ elif action.action is GraphActionType.RemoveNode:
+ row = action.source
+ col = 0
+ type_idx = self.bck_action_type_order.index(action.action)
+ elif action.action is GraphActionType.RemoveEdge:
+ row = ((g.edge_index.T == torch.tensor([(action.source, action.target)])).prod(1)).argmax()
+ row = int(row) // 2 # edges are duplicated, but edge logits are not
+ col = 0
+ type_idx = self.bck_action_type_order.index(action.action)
+ return (type_idx, int(row), int(col))
+
+ def graph_to_Data(self, g: Graph) -> gd.Data:
+ """Convert a networkx Graph to a torch geometric Data instance"""
+ x = torch.zeros((max(1, len(g.nodes)), self.num_node_dim - self._num_rw_feat))
+ x[0, -1] = len(g.nodes) == 0
+ remove_node_mask = torch.zeros((x.shape[0], 1)) + (1 if len(g) == 0 else 0)
+ for i, n in enumerate(g.nodes):
+ ad = g.nodes[n]
+ x[i, self.node_attr_values["v"].index(ad["v"])] = 1
+ if g.degree(n) <= 1:
+ remove_node_mask[i] = 1
+
+ remove_edge_mask = torch.zeros((len(g.edges), 1))
+ for i, (u, v) in enumerate(g.edges):
+ if g.degree(u) > 1 and g.degree(v) > 1:
+ if nx.algorithms.is_connected(graph_without_edge(g, (u, v))):
+ remove_edge_mask[i] = 1
+ edge_attr = torch.zeros((len(g.edges) * 2, self.num_edge_dim))
+ edge_index = (
+ torch.tensor([e for i, j in g.edges for e in [(i, j), (j, i)]], dtype=torch.long).reshape((-1, 2)).T
+ )
+ gc = nx.complement(g)
+ non_edge_index = torch.tensor([i for i in gc.edges], dtype=torch.long).reshape((-1, 2)).T
+ gid = self.get_graph_idx(g) if self.output_gid else 0
+
+ return self._preprocess(
+ gd.Data(
+ x,
+ edge_index,
+ edge_attr,
+ non_edge_index=non_edge_index,
+ stop_mask=torch.ones((1, 1)),
+ add_node_mask=torch.ones((x.shape[0], self.num_new_node_values)) * (len(g) < self.max_nodes),
+ add_edge_mask=torch.ones((non_edge_index.shape[1], 1)),
+ remove_node_mask=remove_node_mask,
+ remove_edge_mask=remove_edge_mask,
+ gid=gid,
+ )
+ )
+
+ def _preprocess(self, g: gd.Data) -> gd.Data:
+ if self._num_rw_feat > 0:
+ g.x = torch.cat([g.x, random_walk_probs(g, self._num_rw_feat, skip_odd=True)], 1)
+ return g
+
+ def collate(self, graphs: List[gd.Data]):
+ """Batch Data instances"""
+ return gd.Batch.from_data_list(graphs, follow_batch=["edge_index", "non_edge_index"])
+
+ def mol_to_graph(self, obj: Graph) -> Graph:
+ return obj # This is already a graph
+
+ def graph_to_mol(self, g: Graph) -> Graph:
+ # idem
+ return g
+
+ def is_sane(self, g: Graph) -> bool:
+ return True
+
+ def get_object_description(self, g: Graph, is_valid: bool) -> str:
+ return str(self.get_graph_idx(g, -1))

src/gflownet/envs/graph_building_env.py

@@ -331,7 +331,7 @@ def generate_forward_trajectory(g: Graph, max_nodes: int = None) -> List[Tuple[G
  # TODO: should this be a method of GraphBuildingEnv? handle set_node_attr flags and so on?
  gn = Graph()
  # Choose an arbitrary starting point, add to the stack
- stack: List[Tuple[int, ...]] = [(np.random.randint(0, len(g.nodes)),)]
+ stack: List[Tuple[int, ...]] = [(np.random.randint(0, len(g.nodes)),)] if len(g.nodes) > 0 else []
  traj = []
  # This map keeps track of node labels in gn, since we have to start from 0
  relabeling_map: Dict[int, int] = {}
@@ -777,6 +777,7 @@ class GraphBuildingEnvContext:
  """A context class defines what the graphs are, how they map to and from data"""
 
  device: torch.device
+ num_cond_dim: int = 0
 
  def aidx_to_GraphAction(self, g: gd.Data, action_idx: Tuple[int, int, int], fwd: bool = True) -> GraphAction:
  """Translate an action index (e.g. from a GraphActionCategorical) to a GraphAction