import torch
[docs]
class SE3Net(torch.nn.Module):
"""Rotational equivariant neural network
Parameters
----------
num_classes : int
size : int
activation : str
"""
def __init__(self, num_classes, size, activation="relu"):
super(SE3Net, self).__init__()
features = [(1,), (2, 2, 2, 1), (4, 4, 4, 4), (6, 4, 4, 0), (64,)]
self.num_features = len(features)
num_radial = size // 2 + 1
max_radius = size // 2
radii = torch.linspace(0, max_radius, steps=num_radial, dtype=torch.float64)
activation = activation
kwargs = {
"radii": radii,
"activation": (torch.nn.functional.relu, torch.sigmoid),
}
layers = []
for i in range(len(features) - 1):
layers.append(PointGatedBlock(features[i], features[i + 1], **kwargs))
self.layers = torch.nn.ModuleList(layers)
with torch_default_dtype(torch.float64):
self.layers.extend(
[
AvgSpacial(),
torch.nn.Dropout(p=0.2),
torch.nn.Linear(64, num_classes),
]
)
[docs]
def forward(self, inputs, difference_mat):
output = inputs
for i in range(self.num_features - 1):
conv = self.layers[i]
output = conv(output, difference_mat)
for i in range(self.num_features - 1, len(self.layers)):
layer = self.layers[i]
output = layer(output)
return output
[docs]
class torch_default_dtype:
def __init__(self, dtype):
self.saved_dtype = None
self.dtype = dtype
def __enter__(self):
self.saved_dtype = torch.get_default_dtype()
torch.set_default_dtype(self.dtype)
def __exit__(self, exc_type, exc_value, traceback):
torch.set_default_dtype(self.saved_dtype)
[docs]
class AvgSpacial(torch.nn.Module):
[docs]
def forward(self, inp):
return inp.view(inp.size(0), inp.size(1), -1).mean(-1)