(self,
in_dims: Union[int, Dict[str, int]],
out_dims: Union[int, Dict[str, int]],
senders: Sequence[int],
receivers: Sequence[int],
n_layers: int = 1,
use_edge_features: bool = True,
use_global_features: bool = True,
residual: Union[bool, Dict[str, bool]] = False,
net_norm: str = 'none',
activation: str = 'relu',
dropout: float = 0,
output_normalization: bool = True,
output_activation_function: Optional[str] = None,
aggregator_funcs: Union[str, Dict[AggregationTypes, int]] = 'sum',
)
| 14 | |
| 15 | class GraphNetBlock(nn.Module): |
| 16 | def __init__(self, |
| 17 | in_dims: Union[int, Dict[str, int]], |
| 18 | out_dims: Union[int, Dict[str, int]], |
| 19 | senders: Sequence[int], |
| 20 | receivers: Sequence[int], |
| 21 | n_layers: int = 1, |
| 22 | use_edge_features: bool = True, |
| 23 | use_global_features: bool = True, |
| 24 | residual: Union[bool, Dict[str, bool]] = False, |
| 25 | net_norm: str = 'none', |
| 26 | activation: str = 'relu', |
| 27 | dropout: float = 0, |
| 28 | output_normalization: bool = True, |
| 29 | output_activation_function: Optional[str] = None, |
| 30 | aggregator_funcs: Union[str, Dict[AggregationTypes, int]] = 'sum', |
| 31 | ): |
| 32 | super().__init__() |
| 33 | if isinstance(in_dims, int): |
| 34 | in_dims: dict = {c: in_dims for c in gn_constants.GRAPH_COMPONENTS} |
| 35 | if isinstance(out_dims, int): |
| 36 | out_dims: dict = {c: out_dims for c in gn_constants.GRAPH_COMPONENTS} |
| 37 | |
| 38 | self.components = gn_constants.GRAPH_COMPONENTS |
| 39 | self.use_edge_features = use_edge_features |
| 40 | self.use_global_features = use_global_features |
| 41 | |
| 42 | if not use_edge_features: |
| 43 | in_dims[EDGES] = out_dims[EDGES] = 0 |
| 44 | self.components.remove(EDGES) |
| 45 | if not use_global_features: |
| 46 | in_dims[GLOBALS] = out_dims[GLOBALS] = 0 |
| 47 | self.components.remove(GLOBALS) |
| 48 | |
| 49 | n_feats_e = in_dims[EDGES] |
| 50 | n_feats_n = in_dims[NODES] |
| 51 | n_feats_u = in_dims[GLOBALS] |
| 52 | self._n_edges = None |
| 53 | self.update_graph_structure(senders, receivers) |
| 54 | self.residual = {c: residual for c in self.components} |
| 55 | |
| 56 | in_dims = { |
| 57 | EDGES: 2 * n_feats_n + n_feats_e + n_feats_u, |
| 58 | NODES: n_feats_n + out_dims[EDGES] + n_feats_u, |
| 59 | GLOBALS: out_dims[NODES] + out_dims[EDGES] + n_feats_u |
| 60 | } |
| 61 | |
| 62 | update_funcs = OrderedDict() # nn.ModuleDict() |
| 63 | for component in self.components: |
| 64 | c_in_dim = in_dims[component] |
| 65 | out_dim = out_dims[component] |
| 66 | if c_in_dim != out_dim: |
| 67 | self.residual[component] = False |
| 68 | |
| 69 | in_dim = in_dims[component] |
| 70 | hdim = int((in_dim + out_dim) / 2) |
| 71 | |
| 72 | update_funcs[component] = MLP( |
| 73 | input_dim=in_dim, |
nothing calls this directly
no test coverage detected