Add class-resolver

Author: cthoyt

torchdrug/layers/__init__.py

@@ -3,7 +3,7 @@
 from .conv import MessagePassingBase, GraphConv, GraphAttentionConv, RelationalGraphConv, GraphIsomorphismConv, \
     NeuralFingerprintConv, ContinuousFilterConv, MessagePassing, ChebyshevConv
 from .pool import DiffPool, MinCutPool
-from .readout import MeanReadout, SumReadout, MaxReadout, Softmax, Set2Set, Sort
+from .readout import MeanReadout, SumReadout, MaxReadout, Softmax, Set2Set, Sort, readout_resolver, Readout
 from .flow import ConditionalFlow
 from .sampler import NodeSampler, EdgeSampler
 from . import distribution, functional

torchdrug/layers/readout.py

@@ -1,9 +1,14 @@
 import torch
 from torch import nn
 from torch_scatter import scatter_mean, scatter_add, scatter_max
+from class_resolver import ClassResolver
 
 
-class MeanReadout(nn.Module):
+class Readout(nn.Module):
+    """A base class for readouts."""
+
+
+class MeanReadout(Readout):
     """Mean readout operator over graphs with variadic sizes."""
 
     def forward(self, graph, input):
@@ -21,7 +26,7 @@ def forward(self, graph, input):
         return output
 
 
-class SumReadout(nn.Module):
+class SumReadout(Readout):
     """Sum readout operator over graphs with variadic sizes."""
 
     def forward(self, graph, input):
@@ -39,7 +44,7 @@ def forward(self, graph, input):
         return output
 
 
-class MaxReadout(nn.Module):
+class MaxReadout(Readout):
     """Max readout operator over graphs with variadic sizes."""
 
     def forward(self, graph, input):
@@ -57,6 +62,12 @@ def forward(self, graph, input):
         return output
 
 
+readout_resolver = ClassResolver.from_subclasses(
+    Readout,
+    default=SumReadout,
+)
+
+
 class Softmax(nn.Module):
     """Softmax operator over graphs with variadic sizes."""
 

torchdrug/models/gat.py

@@ -25,11 +25,11 @@ class GraphAttentionNetwork(nn.Module, core.Configurable):
         batch_norm (bool, optional): apply batch normalization or not
         activation (str or function, optional): activation function
         concat_hidden (bool, optional): concat hidden representations from all layers as output
-        readout (str, optional): readout function. Available functions are ``sum``, ``mean``, and ``max``.
+        readout: readout function. Available functions are ``sum``, ``mean``, and ``max``.
     """
 
     def __init__(self, input_dim, hidden_dims, edge_input_dim=None, num_head=1, negative_slope=0.2, short_cut=False,
-                 batch_norm=False, activation="relu", concat_hidden=False, readout="sum"):
+                 batch_norm=False, activation="relu", concat_hidden=False, readout: Hint[Readout] = "sum"):
         super(GraphAttentionNetwork, self).__init__()
 
         if not isinstance(hidden_dims, Sequence):
@@ -45,14 +45,7 @@ def __init__(self, input_dim, hidden_dims, edge_input_dim=None, num_head=1, nega
             self.layers.append(layers.GraphAttentionConv(self.dims[i], self.dims[i + 1], edge_input_dim, num_head,
                                                          negative_slope, batch_norm, activation))
 
-        if readout == "sum":
-            self.readout = layers.SumReadout()
-        elif readout == "mean":
-            self.readout = layers.MeanReadout()
-        elif readout == "max":
-            self.readout = layers.MaxReadout()
-        else:
-            raise ValueError("Unknown readout `%s`" % readout)
+        self.readout = readout_resolver.make(readout)
 
     def forward(self, graph, input, all_loss=None, metric=None):
         """

torchdrug/models/gcn.py

@@ -1,9 +1,11 @@
 from collections.abc import Sequence
 
 import torch
+from class_resolver import Hint
 from torch import nn
 
 from torchdrug import core, layers
+from torchdrug.layers import readout_resolver, Readout
 from torchdrug.core import Registry as R
 
 
@@ -23,11 +25,11 @@ class GraphConvolutionalNetwork(nn.Module, core.Configurable):
         batch_norm (bool, optional): apply batch normalization or not
         activation (str or function, optional): activation function
         concat_hidden (bool, optional): concat hidden representations from all layers as output
-        readout (str, optional): readout function. Available functions are ``sum``, ``mean``, and ``max``.
+        readout: readout function. Available functions are ``sum``, ``mean``, and ``max``.
     """
 
     def __init__(self, input_dim, hidden_dims, edge_input_dim=None, short_cut=False, batch_norm=False,
-                 activation="relu", concat_hidden=False, readout="sum"):
+                 activation="relu", concat_hidden=False, readout: Hint[Readout] = "sum"):
         super(GraphConvolutionalNetwork, self).__init__()
 
         if not isinstance(hidden_dims, Sequence):
@@ -42,14 +44,7 @@ def __init__(self, input_dim, hidden_dims, edge_input_dim=None, short_cut=False,
         for i in range(len(self.dims) - 1):
             self.layers.append(layers.GraphConv(self.dims[i], self.dims[i + 1], edge_input_dim, batch_norm, activation))
 
-        if readout == "sum":
-            self.readout = layers.SumReadout()
-        elif readout == "mean":
-            self.readout = layers.MeanReadout()
-        elif readout == "max":
-            self.readout = layers.MaxReadout()
-        else:
-            raise ValueError("Unknown readout `%s`" % readout)
+        self.readout = readout_resolver.make(readout)
 
     def forward(self, graph, input, all_loss=None, metric=None):
         """

torchdrug/models/gin.py

@@ -48,14 +48,7 @@ def __init__(self, input_dim=None, hidden_dims=None, edge_input_dim=None, num_ml
             self.layers.append(layers.GraphIsomorphismConv(self.dims[i], self.dims[i + 1], edge_input_dim,
                                                            layer_hidden_dims, eps, learn_eps, batch_norm, activation))
 
-        if readout == "sum":
-            self.readout = layers.SumReadout()
-        elif readout == "mean":
-            self.readout = layers.MeanReadout()
-        elif readout == "max":
-            self.readout = layers.MaxReadout()
-        else:
-            raise ValueError("Unknown readout `%s`" % readout)
+        self.readout = readout_resolver.make(readout)
 
     def forward(self, graph, input, all_loss=None, metric=None):
         """

torchdrug/models/neuralfp.py

@@ -47,12 +47,7 @@ def __init__(self, input_dim, output_dim, hidden_dims, edge_input_dim=None, shor
                                                             batch_norm, activation))
             self.linears.append(nn.Linear(self.dims[i + 1], output_dim))
 
-        if readout == "sum":
-            self.readout = layers.SumReadout()
-        elif readout == "mean":
-            self.readout = layers.MeanReadout()
-        else:
-            raise ValueError("Unknown readout `%s`" % readout)
+        self.readout = readout_resolver.make(readout)
 
     def forward(self, graph, input, all_loss=None, metric=None):
         """

torchdrug/models/schnet.py

@@ -28,7 +28,7 @@ class SchNet(nn.Module, core.Configurable):
     """
 
     def __init__(self, input_dim, hidden_dims, edge_input_dim=None, cutoff=5, num_gaussian=100, short_cut=True,
-                 batch_norm=False, activation="shifted_softplus", concat_hidden=False):
+                 batch_norm=False, activation="shifted_softplus", concat_hidden=False, readout: Hint[Readout] = "sum"):
         super(SchNet, self).__init__()
 
         if not isinstance(hidden_dims, Sequence):
@@ -44,7 +44,7 @@ def __init__(self, input_dim, hidden_dims, edge_input_dim=None, cutoff=5, num_ga
             self.layers.append(layers.ContinuousFilterConv(self.dims[i], self.dims[i + 1], edge_input_dim, None, cutoff,
                                                            num_gaussian, batch_norm, activation))
 
-        self.readout = layers.SumReadout()
+        self.readout = readout_resolver.make(readout)
 
     def forward(self, graph, input, all_loss=None, metric=None):
         """