Spaces:
Sleeping
Sleeping
File size: 3,756 Bytes
59a9ccf |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 |
import torch
import torch.nn as nn
#from equivariant_attention.modules import get_basis_and_r, GSE3Res, GNormBias
#from equivariant_attention.modules import GConvSE3, GNormSE3
#from equivariant_attention.fibers import Fiber
from util_module import init_lecun_normal_param
from se3_transformer.model import SE3Transformer
from se3_transformer.model.fiber import Fiber
class SE3TransformerWrapper(nn.Module):
"""SE(3) equivariant GCN with attention"""
def __init__(self, num_layers=2, num_channels=32, num_degrees=3, n_heads=4, div=4,
l0_in_features=32, l0_out_features=32,
l1_in_features=3, l1_out_features=2,
num_edge_features=32):
super().__init__()
# Build the network
self.l1_in = l1_in_features
#
fiber_edge = Fiber({0: num_edge_features})
if l1_out_features > 0:
if l1_in_features > 0:
fiber_in = Fiber({0: l0_in_features, 1: l1_in_features})
fiber_hidden = Fiber.create(num_degrees, num_channels)
fiber_out = Fiber({0: l0_out_features, 1: l1_out_features})
else:
fiber_in = Fiber({0: l0_in_features})
fiber_hidden = Fiber.create(num_degrees, num_channels)
fiber_out = Fiber({0: l0_out_features, 1: l1_out_features})
else:
if l1_in_features > 0:
fiber_in = Fiber({0: l0_in_features, 1: l1_in_features})
fiber_hidden = Fiber.create(num_degrees, num_channels)
fiber_out = Fiber({0: l0_out_features})
else:
fiber_in = Fiber({0: l0_in_features})
fiber_hidden = Fiber.create(num_degrees, num_channels)
fiber_out = Fiber({0: l0_out_features})
self.se3 = SE3Transformer(num_layers=num_layers,
fiber_in=fiber_in,
fiber_hidden=fiber_hidden,
fiber_out = fiber_out,
num_heads=n_heads,
channels_div=div,
fiber_edge=fiber_edge,
use_layer_norm=True)
#use_layer_norm=False)
self.reset_parameter()
def reset_parameter(self):
# make sure linear layer before ReLu are initialized with kaiming_normal_
for n, p in self.se3.named_parameters():
if "bias" in n:
nn.init.zeros_(p)
elif len(p.shape) == 1:
continue
else:
if "radial_func" not in n:
p = init_lecun_normal_param(p)
else:
if "net.6" in n:
nn.init.zeros_(p)
else:
nn.init.kaiming_normal_(p, nonlinearity='relu')
# make last layers to be zero-initialized
#self.se3.graph_modules[-1].to_kernel_self['0'] = init_lecun_normal_param(self.se3.graph_modules[-1].to_kernel_self['0'])
#self.se3.graph_modules[-1].to_kernel_self['1'] = init_lecun_normal_param(self.se3.graph_modules[-1].to_kernel_self['1'])
nn.init.zeros_(self.se3.graph_modules[-1].to_kernel_self['0'])
nn.init.zeros_(self.se3.graph_modules[-1].to_kernel_self['1'])
def forward(self, G, type_0_features, type_1_features=None, edge_features=None):
if self.l1_in > 0:
node_features = {'0': type_0_features, '1': type_1_features}
else:
node_features = {'0': type_0_features}
edge_features = {'0': edge_features}
return self.se3(G, node_features, edge_features)
|