def _build_gcn(self, fibers):
# Equivariant layers
Gblock = []
fin = fibers['in']
for i in range(self.num_layers):
Gblock.append(
GSE3Res(fin,
fibers['mid'],
edge_dim=self.edge_dim,
div=self.div,
n_heads=self.n_heads,
learnable_skip=True,
skip='cat',
selfint=self.si_m,
x_ij=self.x_ij))
Gblock.append(GNormBias(fibers['mid']))
fin = fibers['mid']
Gblock.append(
GSE3Res(fibers['mid'],
fibers['out'],
edge_dim=self.edge_dim,
div=1,
n_heads=min(1, 2),
learnable_skip=True,
skip='cat',
selfint=self.si_e,
x_ij=self.x_ij))
return nn.ModuleList(Gblock)
def _output_layer(self):
return GSE3Res(self.fibers['mid'],
self.fibers['out'],
div=1,
n_heads=min(self.n_heads, 2),
selfint=self.si_e,
**self._graph_attention_common_params)
def _build_gcn(self, fibers, out_dim):
# Equivariant layers
Gblock = []
fin = fibers['in']
for i in range(self.num_layers):
Gblock.append(
GSE3Res(fin,
fibers['mid'],
edge_dim=self.edge_dim,
div=self.div,
n_heads=self.n_heads))
Gblock.append(GNormSE3(fibers['mid']))
fin = fibers['mid']
Gblock.append(
GConvSE3(fibers['mid'],
fibers['out'],
self_interaction=True,
edge_dim=self.edge_dim))
# Pooling
if self.pooling == 'avg':
Gblock.append(GAvgPooling())
elif self.pooling == 'max':
Gblock.append(GMaxPooling())
# FC layers
FCblock = []
FCblock.append(
nn.Linear(self.fibers['out'].n_features,
self.fibers['out'].n_features))
FCblock.append(nn.ReLU(inplace=True))
FCblock.append(nn.Linear(self.fibers['out'].n_features, out_dim))
return nn.ModuleList(Gblock), nn.ModuleList(FCblock)
def _middle_layer(self):
return GSE3Res(self.fibers['mid'],
self.fibers['mid'],
div=self.div,
n_heads=self.n_heads,
selfint=self.si_m,
**self._graph_attention_common_params)
def _build_gcn(self, fibers, out_dim):
# Equivariant layers
Gblock = []
fin = fibers['in']
for i in range(self.num_layers):
Gblock.append(
GSE3Res(fin,
fibers['mid'],
edge_dim=self.edge_dim,
div=self.div,
n_heads=self.n_heads))
Gblock.append(GNormSE3(fibers['mid']))
fin = fibers['mid']
Gblock.append(
GConvSE3(fibers['mid'],
fibers['out'],
self_interaction=True,
edge_dim=self.edge_dim))
Finblock = []
Finblock.append(OutEncoder())
return nn.ModuleList(Gblock), nn.ModuleList(Finblock)