def __init__(
self,
in_channels: int,
out_channels: int,
aggregators: List[str],
scalers: List[str],
deg: Tensor,
edge_dim: Optional[int] = None,
towers: int = 1,
pre_layers: int = 1,
post_layers: int = 1,
divide_input: bool = False,
act: Union[str, Callable, None] = "relu",
act_kwargs: Optional[Dict[str, Any]] = None,
**kwargs,
):
aggr = DegreeScalerAggregation(aggregators, scalers, deg)
super().__init__(aggr=aggr, node_dim=0, **kwargs)
if divide_input:
assert in_channels % towers == 0
assert out_channels % towers == 0
self.in_channels = in_channels
self.out_channels = out_channels
self.edge_dim = edge_dim
self.towers = towers
self.divide_input = divide_input
self.F_in = in_channels // towers if divide_input else in_channels
self.F_out = self.out_channels // towers
if self.edge_dim is not None:
self.edge_encoder = Linear(edge_dim, self.F_in)
self.pre_nns = ModuleList()
self.post_nns = ModuleList()
for _ in range(towers):
modules = [Linear((3 if edge_dim else 2) * self.F_in, self.F_in)]
for _ in range(pre_layers - 1):
modules += [activation_resolver(act, **(act_kwargs or {}))]
modules += [Linear(self.F_in, self.F_in)]
self.pre_nns.append(Sequential(*modules))
in_channels = (len(aggregators) * len(scalers) + 1) * self.F_in
modules = [Linear(in_channels, self.F_out)]
for _ in range(post_layers - 1):
modules += [activation_resolver(act, **(act_kwargs or {}))]
modules += [Linear(self.F_out, self.F_out)]
self.post_nns.append(Sequential(*modules))
self.lin = Linear(out_channels, out_channels)
self.reset_parameters()
def __init__(self, hidden_channels: int, out_channels: int,
num_blocks: int, num_bilinear: int, num_spherical: int,
num_radial, cutoff: float = 5.0, max_num_neighbors: int = 32,
envelope_exponent: int = 5, num_before_skip: int = 1,
num_after_skip: int = 2, num_output_layers: int = 3,
act: Union[str, Callable] = 'swish'):
super().__init__()
if num_spherical < 2:
raise ValueError("num_spherical should be greater than 1")
act = activation_resolver(act)
self.cutoff = cutoff
self.max_num_neighbors = max_num_neighbors
self.num_blocks = num_blocks
self.rbf = BesselBasisLayer(num_radial, cutoff, envelope_exponent)
self.sbf = SphericalBasisLayer(num_spherical, num_radial, cutoff,
envelope_exponent)
self.emb = EmbeddingBlock(num_radial, hidden_channels, act)
self.output_blocks = torch.nn.ModuleList([
OutputBlock(num_radial, hidden_channels, out_channels,
num_output_layers, act) for _ in range(num_blocks + 1)
])
self.interaction_blocks = torch.nn.ModuleList([
InteractionBlock(hidden_channels, num_bilinear, num_spherical,
num_radial, num_before_skip, num_after_skip, act)
for _ in range(num_blocks)
])
self.reset_parameters()
def __init__(
self,
channel_list: Optional[Union[List[int], int]] = None,
*,
in_channels: Optional[int] = None,
hidden_channels: Optional[int] = None,
out_channels: Optional[int] = None,
num_layers: Optional[int] = None,
dropout: float = 0.,
act: str = "relu",
batch_norm: bool = True,
act_first: bool = False,
act_kwargs: Optional[Dict[str, Any]] = None,
batch_norm_kwargs: Optional[Dict[str, Any]] = None,
plain_last: bool = True,
bias: bool = True,
relu_first: bool = False,
):
super().__init__()
act_first = act_first or relu_first # Backward compatibility.
batch_norm_kwargs = batch_norm_kwargs or {}
if isinstance(channel_list, int):
in_channels = channel_list
if in_channels is not None:
assert num_layers >= 1
channel_list = [hidden_channels] * (num_layers - 1)
channel_list = [in_channels] + channel_list + [out_channels]
assert isinstance(channel_list, (tuple, list))
assert len(channel_list) >= 2
self.channel_list = channel_list
self.dropout = dropout
self.act = activation_resolver(act, **(act_kwargs or {}))
self.act_first = act_first
self.plain_last = plain_last
self.lins = torch.nn.ModuleList()
iterator = zip(channel_list[:-1], channel_list[1:])
for in_channels, out_channels in iterator:
self.lins.append(Linear(in_channels, out_channels, bias=bias))
self.norms = torch.nn.ModuleList()
iterator = channel_list[1:-1] if plain_last else channel_list[1:]
for hidden_channels in iterator:
if batch_norm:
norm = BatchNorm1d(hidden_channels, **batch_norm_kwargs)
else:
norm = Identity()
self.norms.append(norm)
self.reset_parameters()
def test_activation_resolver():
assert isinstance(activation_resolver(torch.nn.ELU()), torch.nn.ELU)
assert isinstance(activation_resolver(torch.nn.ReLU()), torch.nn.ReLU)
assert isinstance(activation_resolver(torch.nn.PReLU()), torch.nn.PReLU)
assert isinstance(activation_resolver('elu'), torch.nn.ELU)
assert isinstance(activation_resolver('relu'), torch.nn.ReLU)
assert isinstance(activation_resolver('prelu'), torch.nn.PReLU)
def __init__(self, hidden_channels: int, out_channels: int,
num_blocks: int, int_emb_size: int, basis_emb_size: int,
out_emb_channels: int, num_spherical: int, num_radial: int,
cutoff: float = 5.0, max_num_neighbors: int = 32,
envelope_exponent: int = 5, num_before_skip: int = 1,
num_after_skip: int = 2, num_output_layers: int = 3,
act: Union[str, Callable] = 'swish'):
act = activation_resolver(act)
super().__init__(
hidden_channels=hidden_channels,
out_channels=out_channels,
num_blocks=num_blocks,
num_bilinear=1,
num_spherical=num_spherical,
num_radial=num_radial,
cutoff=cutoff,
max_num_neighbors=max_num_neighbors,
envelope_exponent=envelope_exponent,
num_before_skip=num_before_skip,
num_after_skip=num_after_skip,
num_output_layers=num_output_layers,
act=act,
)
# We are re-using the RBF, SBF and embedding layers of `DimeNet` and
# redefine output_block and interaction_block in DimeNet++.
# Hence, it is to be noted that in the above initalization, the
# variable `num_bilinear` does not have any purpose as it is used
# solely in the `OutputBlock` of DimeNet:
self.output_blocks = torch.nn.ModuleList([
OutputPPBlock(num_radial, hidden_channels, out_emb_channels,
out_channels, num_output_layers, act)
for _ in range(num_blocks + 1)
])
self.interaction_blocks = torch.nn.ModuleList([
InteractionPPBlock(hidden_channels, int_emb_size, basis_emb_size,
num_spherical, num_radial, num_before_skip,
num_after_skip, act) for _ in range(num_blocks)
])
self.reset_parameters()
def __init__(
self,
in_channels: int,
hidden_channels: int,
num_layers: int,
out_channels: Optional[int] = None,
dropout: float = 0.0,
act: Union[str, Callable, None] = "relu",
act_first: bool = False,
act_kwargs: Optional[Dict[str, Any]] = None,
norm: Union[str, Callable, None] = None,
norm_kwargs: Optional[Dict[str, Any]] = None,
jk: Optional[str] = None,
**kwargs,
):
super().__init__()
self.in_channels = in_channels
self.hidden_channels = hidden_channels
self.num_layers = num_layers
self.dropout = dropout
self.act = activation_resolver(act, **(act_kwargs or {}))
self.jk_mode = jk
self.act_first = act_first
self.norm = norm if isinstance(norm, str) else None
self.norm_kwargs = norm_kwargs
if out_channels is not None:
self.out_channels = out_channels
else:
self.out_channels = hidden_channels
self.convs = ModuleList()
if num_layers > 1:
self.convs.append(
self.init_conv(in_channels, hidden_channels, **kwargs))
if isinstance(in_channels, (tuple, list)):
in_channels = (hidden_channels, hidden_channels)
else:
in_channels = hidden_channels
for _ in range(num_layers - 2):
self.convs.append(
self.init_conv(in_channels, hidden_channels, **kwargs))
if isinstance(in_channels, (tuple, list)):
in_channels = (hidden_channels, hidden_channels)
else:
in_channels = hidden_channels
if out_channels is not None and jk is None:
self._is_conv_to_out = True
self.convs.append(
self.init_conv(in_channels, out_channels, **kwargs))
else:
self.convs.append(
self.init_conv(in_channels, hidden_channels, **kwargs))
self.norms = None
if norm is not None:
norm_layer = normalization_resolver(
norm,
hidden_channels,
**(norm_kwargs or {}),
)
self.norms = ModuleList()
for _ in range(num_layers - 1):
self.norms.append(copy.deepcopy(norm_layer))
if jk is not None:
self.norms.append(copy.deepcopy(norm_layer))
if jk is not None and jk != 'last':
self.jk = JumpingKnowledge(jk, hidden_channels, num_layers)
if jk is not None:
if jk == 'cat':
in_channels = num_layers * hidden_channels
else:
in_channels = hidden_channels
self.lin = Linear(in_channels, self.out_channels)
def __init__(
self,
channel_list: Optional[Union[List[int], int]] = None,
*,
in_channels: Optional[int] = None,
hidden_channels: Optional[int] = None,
out_channels: Optional[int] = None,
num_layers: Optional[int] = None,
dropout: Union[float, List[float]] = 0.,
act: Union[str, Callable, None] = "relu",
act_first: bool = False,
act_kwargs: Optional[Dict[str, Any]] = None,
norm: Union[str, Callable, None] = "batch_norm",
norm_kwargs: Optional[Dict[str, Any]] = None,
plain_last: bool = True,
bias: Union[bool, List[bool]] = True,
**kwargs,
):
super().__init__()
# Backward compatibility:
act_first = act_first or kwargs.get("relu_first", False)
batch_norm = kwargs.get("batch_norm", None)
if batch_norm is not None and isinstance(batch_norm, bool):
warnings.warn("Argument `batch_norm` is deprecated, "
"please use `norm` to specify normalization layer.")
norm = 'batch_norm' if batch_norm else None
batch_norm_kwargs = kwargs.get("batch_norm_kwargs", None)
norm_kwargs = batch_norm_kwargs or {}
if isinstance(channel_list, int):
in_channels = channel_list
if in_channels is not None:
assert num_layers >= 1
channel_list = [hidden_channels] * (num_layers - 1)
channel_list = [in_channels] + channel_list + [out_channels]
assert isinstance(channel_list, (tuple, list))
assert len(channel_list) >= 2
self.channel_list = channel_list
self.act = activation_resolver(act, **(act_kwargs or {}))
self.act_first = act_first
self.plain_last = plain_last
if isinstance(dropout, float):
dropout = [dropout] * (len(channel_list) - 1)
if plain_last:
dropout[-1] = 0.
if len(dropout) != len(channel_list) - 1:
raise ValueError(
f"Number of dropout values provided ({len(dropout)} does not "
f"match the number of layers specified "
f"({len(channel_list)-1})")
self.dropout = dropout
if isinstance(bias, bool):
bias = [bias] * (len(channel_list) - 1)
if len(bias) != len(channel_list) - 1:
raise ValueError(
f"Number of bias values provided ({len(bias)}) does not match "
f"the number of layers specified ({len(channel_list)-1})")
self.lins = torch.nn.ModuleList()
iterator = zip(channel_list[:-1], channel_list[1:], bias)
for in_channels, out_channels, _bias in iterator:
self.lins.append(Linear(in_channels, out_channels, bias=_bias))
self.norms = torch.nn.ModuleList()
iterator = channel_list[1:-1] if plain_last else channel_list[1:]
for hidden_channels in iterator:
if norm is not None:
norm_layer = normalization_resolver(
norm,
hidden_channels,
**(norm_kwargs or {}),
)
else:
norm_layer = Identity()
self.norms.append(norm_layer)
self.reset_parameters()