def _build_gcn(self, fibers, out_dim):
block0 = []
fin = fibers['in']
for i in range(self.num_layers - 1):
block0.append(
GConvSE3(fin,
fibers['mid'],
self_interaction=True,
edge_dim=self.edge_dim))
block0.append(GNormSE3(fibers['mid'], num_layers=self.num_nlayers))
fin = fibers['mid']
block0.append(
GConvSE3(fibers['mid'],
fibers['out'],
self_interaction=True,
edge_dim=self.edge_dim))
block1 = [GMaxPooling()]
block2 = []
block2.append(nn.Linear(self.num_channels_out, self.num_channels_out))
block2.append(nn.ReLU(inplace=True))
block2.append(nn.Linear(self.num_channels_out, out_dim))
return nn.ModuleList(block0), nn.ModuleList(block1), nn.ModuleList(
block2)
def _build_gcn(self, fibers):
# Equivariant layers
Gblock = []
fin = fibers['in']
for i in range(self.num_layers-1):
Gblock.append(GConvSE3(fin, fibers['mid'], self_interaction=True, flavor='TFN', edge_dim=self.edge_dim))
Gblock.append(GNormBias(fibers['mid']))
fin = fibers['mid']
Gblock.append(
GConvSE3(fibers['mid'], fibers['out'], self_interaction=True, flavor='TFN', edge_dim=self.edge_dim))
return nn.ModuleList(Gblock), nn.ModuleList([])
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 _build_gcn(self, fibers):
block0 = []
fin = fibers['in']
for i in range(self.num_layers - 1):
block0.append(
GConvSE3(fin,
fibers['mid'],
self_interaction=self.use_self,
edge_dim=self.edge_dim))
block0.append(GNormSE3(fibers['mid'], num_layers=self.num_nlayers))
fin = fibers['mid']
block0.append(
GConvSE3(fibers['mid'],
fibers['out'],
self_interaction=self.use_self,
edge_dim=self.edge_dim))
return nn.ModuleList(block0)
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)
