def __init__(self, f_in, f_mid, f_out, b_in, b_mid, b_out, n_group_mid,
n_group_out):
super().__init__()
grid1 = so3_near_identity_grid(n_alpha=2 * b_in, n_beta=2,
n_gamma=2) #max_beta=0, max_gamma=0,
grid2 = so3_near_identity_grid(n_alpha=2 * b_mid, n_beta=2,
n_gamma=2) #max_beta=0, max_gamma=0,
self.left = nn.Sequential(
SO3Convolution(f_in, f_mid, b_in, b_mid, grid1),
nn.GroupNorm(n_group_mid, f_mid),
SO3Convolution(f_mid, f_out, b_mid, b_out, grid2),
nn.GroupNorm(n_group_out, f_out), nonlinear())
self.shortcut = SO3Shortcut(f_in, f_out, b_in, b_out)
def __init__(self, features=None, bandwidths=None):
super().__init__()
if bandwidths is None:
bandwidths = [24, 24]
if features is None:
features = [3, 100]
self.features = features
self.bandwidths = bandwidths
# define the number of layers.
self.num_layers = 1
assert len(self.bandwidths) == len(self.features)
sequence = []
# S2 layer
# =====================If we could chose another gridding method=====================
grid = so3_equatorial_grid(max_beta=0,
max_gamma=0,
n_alpha=2 * bandwidths[0],
n_beta=1,
n_gamma=1)
sequence.append(
SO3Convolution(features[0], features[1], bandwidths[0],
bandwidths[1], grid))
sequence.append(nn.BatchNorm3d(self.features[-1], affine=True))
sequence.append(nn.ReLU())
self.sequential = nn.Sequential(*sequence)
self.bn1 = nn.BatchNorm1d(self.features[-1])
def __init__(self, f_in, f_out, b_in, b_out):
super(VolterraBlock, self).__init__()
self.s2_seq = S2Convolution(f_in, f_out, b_in, b_out, s2_near_identity_grid())
self.so3_seq = SO3Convolution(f_out, f_out, b_out, b_out, so3_near_identity_grid())
self.bn = nn.BatchNorm3d(f_out, affine=True)
self.relu = nn.ReLU()
def __init__(self, nclasses):
super().__init__()
self.features = [hyper.R_IN, 40, 40, nclasses]
self.bandwidths = [hyper.BANDWIDTH_IN, 32, 32, hyper.BANDWIDTH_OUT]
self.linear1 = nn.Linear(nclasses + hyper.N_CATS, 50)
self.linear2 = nn.Linear(50, 50)
sequence = []
# S2 layer
grid = s2_equatorial_grid(max_beta=0, n_alpha=2 * self.bandwidths[0], n_beta=1)
sequence.append(S2Convolution(self.features[0], self.features[1], self.bandwidths[0], self.bandwidths[1], grid))
# SO3 layers
for l in range(1, len(self.features) - 1):
nfeature_in = self.features[l]
nfeature_out = self.features[l + 1]
b_in = self.bandwidths[l]
b_out = self.bandwidths[l + 1]
sequence.append(nn.BatchNorm3d(nfeature_in, affine=True))
sequence.append(nn.ReLU())
grid = so3_equatorial_grid(max_beta=0, max_gamma=0, n_alpha=2 * b_in, n_beta=1, n_gamma=1)
sequence.append(SO3Convolution(nfeature_in, nfeature_out, b_in, b_out, grid))
sequence.append(nn.BatchNorm3d(self.features[-1], affine=True))
sequence.append(nn.ReLU())
self.sequential = nn.Sequential(*sequence)
def __init__(self, para_dict):
super(S2ConvNet, self).__init__()
self.para_dict = para_dict
self.batch_size = self.para_dict['batchsize']
self.num_points = self.para_dict['num_points']
self.f1 = self.para_dict['f1']
self.f2 = self.para_dict['f2']
self.f_output = self.para_dict['f_output']
self.b_in = self.para_dict['b_in']
self.b_l1 = self.para_dict['b_l1']
self.b_l2 = self.para_dict['b_l2']
# self.kernel_size = self.para_dict['kernel_size']
grid_s2 = s2_near_identity_grid()
grid_so3 = so3_near_identity_grid()
self.conv1 = S2Convolution(nfeature_in=1,
nfeature_out=self.f1,
b_in=self.b_in,
b_out=self.b_l1,
grid=grid_s2)
self.conv2 = SO3Convolution(nfeature_in=self.f1,
nfeature_out=self.f2,
b_in=self.b_l1,
b_out=self.b_l2,
grid=grid_so3)
self.maxPool = nn.MaxPool1d(kernel_size=self.num_points)
self.out_layer = nn.Linear(self.f2, self.f_output)
def __init__(self):
super(S2ConvNet_original, self).__init__()
f1 = 20
f2 = 40
f_output = 10
b_in = 30
b_l1 = 10
b_l2 = 6
grid_s2 = s2_near_identity_grid()
grid_so3 = so3_near_identity_grid()
self.conv1 = S2Convolution(nfeature_in=1,
nfeature_out=f1,
b_in=b_in,
b_out=b_l1,
grid=grid_s2)
self.conv2 = SO3Convolution(nfeature_in=f1,
nfeature_out=f2,
b_in=b_l1,
b_out=b_l2,
grid=grid_so3)
self.out_layer = nn.Linear(f2, f_output)
def __init__(self, nclasses, bw_in):
super().__init__()
self.features = [1, 100, 100, nclasses]
self.bandwidths = [bw_in, 16, 10]
assert len(self.bandwidths) == len(self.features) - 1
sequence = []
# S2 layer
grid = s2_equatorial_grid(max_beta=0, n_alpha=2 * self.bandwidths[0], n_beta=1)
sequence.append(S2Convolution(self.features[0], self.features[1], self.bandwidths[0], self.bandwidths[1], grid))
# SO3 layers
for l in range(1, len(self.features) - 2):
nfeature_in = self.features[l]
nfeature_out = self.features[l + 1]
b_in = self.bandwidths[l]
b_out = self.bandwidths[l + 1]
sequence.append(nn.BatchNorm3d(nfeature_in, affine=True))
sequence.append(nn.ReLU())
grid = so3_equatorial_grid(max_beta=0, max_gamma=0, n_alpha=2 * b_in, n_beta=1, n_gamma=1)
sequence.append(SO3Convolution(nfeature_in, nfeature_out, b_in, b_out, grid))
sequence.append(nn.BatchNorm3d(self.features[-2], affine=True))
sequence.append(nn.ReLU())
self.sequential = nn.Sequential(*sequence)
# Output layer
output_features = self.features[-2]
self.out_layer = nn.Linear(output_features, self.features[-1])
def _make_SO3_layer(self, l):
sequence = []
nfeature_in = self.features[l] * l
nfeature_out = self.features[l + 1]
b_in = self.bandwidths[l]
b_out = self.bandwidths[l + 1]
sequence.append(
SO3Convolution(nfeature_in, nfeature_out, b_in, b_out, grid))
#sequence.append(nn.Dropout(p=0.5))
return nn.Sequential(*sequence)
def __init__(self):
super().__init__()
self.features = [3, 25, 25, 25, 36]
self.bandwidths = [32, 32, 32, 32, 32]
# S2 layer
grid = s2_equatorial_grid(max_beta=0,
n_alpha=2 * self.bandwidths[0],
n_beta=1)
self.s2_conv = S2Convolution(self.features[0], self.features[1],
self.bandwidths[0], self.bandwidths[1],
grid)
self.bn1 = nn.BatchNorm3d(self.features[1], affine=True)
self.relu1 = nn.ReLU()
# SO3 layers
b_in = 32
b_out = 32
grid3 = so3_equatorial_grid(max_beta=0,
max_gamma=0,
n_alpha=2 * 32,
n_beta=1,
n_gamma=1)
self.so3_layer1 = SO3Convolution(self.features[1], self.features[2],
b_in, b_out, grid3)
self.bn2 = nn.BatchNorm3d(self.features[1] + self.features[2],
affine=True)
self.relu2 = nn.ReLU()
self.so3_layer2 = SO3Convolution(self.features[1] + self.features[2],
self.features[3], b_in, b_out, grid3)
self.bn3 = nn.BatchNorm3d(self.features[1] + self.features[2] +
self.features[3],
affine=True)
self.relu3 = nn.ReLU()
self.so3_layer3 = SO3Convolution(
self.features[1] + self.features[2] + self.features[3],
self.features[4], b_in, b_out, grid3)
self.bn4 = nn.BatchNorm3d(self.features[4], affine=True)
self.relu4 = nn.ReLU()
def __init__(self, f0, f1, f2, b_in, b_l):
super(VolterraBlock, self).__init__()
self.s2_seq = S2Convolution(f0, f1, b_in, b_l, s2_near_identity_grid())
self.so3_seq = SO3Convolution(f1, f2, b_l, b_l,
so3_near_identity_grid())
f = (f1 + f2) // 2
self.conv3d_0 = nn.Conv3d(f1, f, kernel_size=1)
self.conv3d_1 = nn.Conv3d(f2, f, kernel_size=1)
self.bn = nn.BatchNorm3d(f, affine=True)
self.relu = nn.ReLU()
def main():
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# load image
x = imread("earth128.jpg").astype(np.float32).transpose((2, 0, 1)) / 255
b = 64
x = torch.tensor(x, dtype=torch.float, device=device)
x = x.view(1, 3, 2 * b, 2 * b)
# equivariant transformation
s2_grid = s2_near_identity_grid(max_beta=0.2, n_alpha=12, n_beta=1)
s2_conv = S2Convolution(3, 50, b_in=b, b_out=b, grid=s2_grid)
s2_conv.to(device)
so3_grid = so3_near_identity_grid(max_beta=0.2, n_alpha=12, n_beta=1)
so3_conv = SO3Convolution(50, 1, b_in=b, b_out=b, grid=so3_grid)
so3_conv.to(device)
def phi(x):
x = s2_conv(x)
x = torch.nn.functional.softplus(x)
x = so3_conv(x)
return x
# test equivariance
abc = (0.5, 1, 0) # rotation angles
y1 = phi(s2_rotation(x, *abc))
y2 = so3_rotation(phi(x), *abc)
print((y1 - y2).std().item(), y1.std().item())
plt.figure(figsize=(12, 8))
plt.subplot(2, 3, 1)
plot(x, "x : signal on the sphere")
plt.subplot(2, 3, 2)
plot(phi(x), "phi(x) : convolutions", True)
plt.subplot(2, 3, 3)
plot(so3_rotation(phi(x), *abc), "R(phi(x))", True)
plt.subplot(2, 3, 4)
plot(s2_rotation(x, *abc), "R(x) : rotation using fft")
plt.subplot(2, 3, 5)
plot(phi(s2_rotation(x, *abc)), "phi(R(x))", True)
plt.tight_layout()
plt.savefig("fig.jpeg")
def __init__(self, nclasses, cfg):
super().__init__()
self.features = [1, 50, 50, 50, 256, 128, nclasses]
self.bandwidths = [cfg.bw, cfg.bw, cfg.bw, cfg.bw]
sequence = []
# S2 layer
grid = s2_equatorial_grid(max_beta=0,
n_alpha=2 * self.bandwidths[0],
n_beta=1)
sequence.append(
S2HConvolution(self.features[0], self.features[1],
self.bandwidths[0], self.bandwidths[1], grid))
# SO3 layers
for l in range(1, len(self.features) - 4):
nfeature_in = self.features[l]
nfeature_out = self.features[l + 1]
b_in = self.bandwidths[l]
b_out = self.bandwidths[l + 1]
sequence.append(nn.BatchNorm3d(nfeature_in, affine=True))
sequence.append(nn.ReLU())
grid = so3_equatorial_grid(max_beta=0,
max_gamma=0,
n_alpha=2 * b_in,
n_beta=1,
n_gamma=1)
sequence.append(
SO3Convolution(nfeature_in, nfeature_out, b_in, b_out, grid))
sequence.append(nn.BatchNorm3d(self.features[3], affine=True))
sequence.append(nn.ReLU())
self.sequential = nn.Sequential(*sequence)
fcs = []
# Output layer
for i in range(3, len(self.features) - 1):
ch_in = self.features[i]
ch_out = self.features[i + 1]
fcs.append(nn.Linear(ch_in, ch_out))
if i < len(self.features) - 2:
fcs.append(nn.BatchNorm1d(ch_out))
fcs.append(nn.ReLU())
self.out_layer = nn.Sequential(*fcs)
def __init__(self, f_in, f_out, b_in, b_out):
super().__init__()
assert b_out <= b_in
self.tmp = [f_in, f_out, b_in, b_out]
if (f_in != f_out) or (b_in != b_out):
self.conv = nn.Sequential(
SO3Convolution(nfeature_in=f_in,
nfeature_out=f_out,
b_in=b_in,
b_out=b_out,
grid=((0, 0, 0), )),
nn.BatchNorm3d(f_out, affine=True))
else:
self.conv = None
def __init__(self, f_in, f_out, b_in, b_out):
super(VolterraBlock, self).__init__()
# s2 sequence
seq = []
seq.append(
S2Convolution(f_in, f_out, b_in, b_out, s2_near_identity_grid()))
seq.append(nn.BatchNorm3d(f_out, affine=True))
seq.append(nn.ReLU())
self.s2_seq = nn.Sequential(*seq)
# so3 sequence
seq = []
seq.append(
SO3Convolution(f_out, f_out, b_out, b_out,
so3_near_identity_grid()))
seq.append(nn.BatchNorm3d(f_out, affine=True))
seq.append(nn.ReLU())
self.so3_seq = nn.Sequential(*seq)
def __init__(self):
super().__init__()
self.features = constant.ENCODER_FEATURES
self.bandwidths = constant.ENCODER_BANDWIDTH
assert len(self.bandwidths) == len(self.features)
sequence = []
# S2 layer
grid = s2_near_identity_grid(n_alpha=2 * self.bandwidths[0], n_beta=2)
sequence.append(
S2Convolution(self.features[0], self.features[1],
self.bandwidths[0], self.bandwidths[1], grid))
sequence.append(nn.BatchNorm3d(self.features[1], affine=True))
sequence.append(nn.ReLU())
# SO3 layers
for l in range(1, len(self.features) - 1):
nfeature_in = self.features[l]
nfeature_out = self.features[l + 1]
b_in = self.bandwidths[l]
b_out = self.bandwidths[l + 1]
grid = so3_near_identity_grid(n_alpha=2 * b_in,
n_beta=2,
n_gamma=2)
sequence.append(
SO3Convolution(nfeature_in, nfeature_out, b_in, b_out, grid))
sequence.append(nn.BatchNorm3d(nfeature_out, affine=True))
sequence.append(nn.ReLU())
self.sequential = nn.Sequential(*sequence)
self.input_size = constant.REPRESENTATION_SIZE
self.hidden_size = constant.ENCODER_HIDDEN_SIZE
self.n_layers = constant.ENCODER_LAYERS
self.rnn = nn.GRU(input_size=self.input_size,
hidden_size=self.hidden_size,
num_layers=self.n_layers,
batch_first=True)
def __init__(self, nclasses, cfg):
super().__init__()
self.features = [1, 40, 40, nclasses]
self.bandwidths = [
cfg.bw,
] * len(self.features)
self.linear1 = nn.Linear(nclasses + 16, 50)
self.linear2 = nn.Linear(50, 50)
sequence = []
# S2 layer
grid = s2_equatorial_grid(max_beta=0,
n_alpha=2 * self.bandwidths[0],
n_beta=1)
sequence.append(
S2HConvolution(self.features[0], self.features[1],
self.bandwidths[0], self.bandwidths[1], grid))
# SO3 layers
for l in range(1, len(self.features) - 1):
nfeature_in = self.features[l]
nfeature_out = self.features[l + 1]
b_in = self.bandwidths[l]
b_out = self.bandwidths[l + 1]
sequence.append(nn.BatchNorm3d(nfeature_in, affine=True))
sequence.append(nn.ReLU())
grid = so3_equatorial_grid(max_beta=0,
max_gamma=0,
n_alpha=2 * b_in,
n_beta=1,
n_gamma=1)
sequence.append(
SO3Convolution(nfeature_in, nfeature_out, b_in, b_out, grid))
sequence.append(nn.BatchNorm3d(self.features[-1], affine=True))
sequence.append(nn.ReLU())
sequence.append(GammaMean(2 * self.bandwidths[-1]))
self.sequential = nn.Sequential(*sequence)
def __init__(self, out_channels=64, dropout_ratio=None):
super().__init__()
self.features = [3, 100, 100, out_channels]
self.bandwidths = [init_bandwidth, 16, 10]
assert len(self.bandwidths) == len(self.features) - 1
sequence = []
# S2 layer
# =====================If we could chose another gridding method=====================
grid = s2_equatorial_grid(max_beta=0, n_alpha=2 * self.bandwidths[0], n_beta=1)
sequence.append(S2Convolution(self.features[0], self.features[1], self.bandwidths[0], self.bandwidths[1], grid))
# SO3 layers
for l in range(1, len(self.features) - 2):
nfeature_in = self.features[l]
nfeature_out = self.features[l + 1]
b_in = self.bandwidths[l]
b_out = self.bandwidths[l + 1]
sequence.append(nn.BatchNorm3d(nfeature_in, affine=True))
sequence.append(nn.ReLU())
grid = so3_equatorial_grid(max_beta=0, max_gamma=0, n_alpha=2 * b_in, n_beta=1, n_gamma=1)
sequence.append(SO3Convolution(nfeature_in, nfeature_out, b_in, b_out, grid))
sequence.append(nn.BatchNorm3d(self.features[-2], affine=True))
sequence.append(nn.ReLU())
self.sequential = nn.Sequential(*sequence)
# Output layer
output_features = self.features[-2]
self.out_layer = nn.Linear(output_features, self.features[-1])
self.bn = nn.BatchNorm1d(self.features[-1])
if dropout_ratio != 0:
self.dropout = nn.Dropout(p=dropout_ratio)
print("Spherical models use dropout.")
else:
self.dropout = None
def __init__(self):
super().__init__()
self.features = constant.ENCODER_FEATURES
self.bandwidths = constant.ENCODER_BANDWIDTH
self.groups = constant.ENCODER_GROUPS
assert len(self.bandwidths) == len(self.features)
sequence = []
# S2 layer
grid = s2_near_identity_grid(n_alpha=2 * self.bandwidths[0], n_beta=2)
sequence.append(
S2Convolution(self.features[0], self.features[1],
self.bandwidths[0], self.bandwidths[1], grid))
sequence.append(nn.BatchNorm3d(self.features[1], affine=True))
sequence.append(nonlinear())
# SO3 layers
for l in range(1, len(self.features) - 1):
f_in = self.features[l]
f_out = self.features[l + 1]
b_in = self.bandwidths[l]
b_out = self.bandwidths[l + 1]
n_group = self.groups[l]
grid = so3_near_identity_grid(n_alpha=2 * b_in,
n_beta=2,
n_gamma=2)
sequence.append(SO3Convolution(f_in, f_out, b_in, b_out, grid))
sequence.append(nn.BatchNorm3d(f_out, affine=True))
sequence.append(nonlinear())
self.sequential = nn.Sequential(*sequence)
def __init__(self, params):
super(SphericalGMMNet, self).__init__()
self.params = params
self.num_grids = self.params['num_grids']
self.batch_size = self.params['batch_size']
self.num_points = self.params['num_points']
self.density_radius = self.params['density_radius']
self.feature_out1 = self.params['feature_out1']
self.feature_out2 = self.params['feature_out2']
self.feature_out3 = self.params['feature_out3']
self.feature_out4 = self.params['feature_out4']
self.feature_out5 = self.params['feature_out5']
self.num_classes = self.params['num_classes']
self.num_so3_layers = self.params['num_so3_layers']
self.bandwidth_0 = self.params['bandwidth_0']
self.bandwidth_out1 = self.params['bandwidth_out1']
self.bandwidth_out2 = self.params['bandwidth_out2']
self.bandwidth_out3 = self.params['bandwidth_out3']
self.bandwidth_out4 = self.params['bandwidth_out4']
self.bandwidth_out5 = self.params['bandwidth_out5']
grid_s2 = s2_near_identity_grid()
grid_so3 = so3_near_identity_grid()
# s2 conv [Learn Pattern] -----------------------------------------------
self.conv0_0 = S2Convolution(nfeature_in=1,
nfeature_out=self.feature_out1,
b_in=self.bandwidth_0,
b_out=self.bandwidth_out1,
grid=grid_s2)
self.conv0_1 = S2Convolution(nfeature_in=1,
nfeature_out=self.feature_out1,
b_in=self.bandwidth_0,
b_out=self.bandwidth_out1,
grid=grid_s2)
self.conv0_2 = S2Convolution(nfeature_in=1,
nfeature_out=self.feature_out1,
b_in=self.bandwidth_0,
b_out=self.bandwidth_out1,
grid=grid_s2)
self.bn0_0 = nn.BatchNorm3d(num_features=self.feature_out1)
self.bn0_1 = nn.BatchNorm3d(num_features=self.feature_out1)
self.bn0_2 = nn.BatchNorm3d(num_features=self.feature_out1)
# so3 conv (1) [Rotation Invariant] -----------------------------------------------
self.conv1_0 = SO3Convolution(nfeature_in=self.feature_out1,
nfeature_out=self.feature_out2,
b_in=self.bandwidth_out1,
b_out=self.bandwidth_out2,
grid=grid_so3)
self.conv1_1 = SO3Convolution(nfeature_in=self.feature_out1,
nfeature_out=self.feature_out2,
b_in=self.bandwidth_out1,
b_out=self.bandwidth_out2,
grid=grid_so3)
self.conv1_2 = SO3Convolution(nfeature_in=self.feature_out1,
nfeature_out=self.feature_out2,
b_in=self.bandwidth_out1,
b_out=self.bandwidth_out2,
grid=grid_so3)
self.bn1_0 = nn.BatchNorm3d(num_features=self.feature_out2)
self.bn1_1 = nn.BatchNorm3d(num_features=self.feature_out2)
self.bn1_2 = nn.BatchNorm3d(num_features=self.feature_out2)
# so3 conv (2) [Rotation Invariant] -----------------------------------------------
self.conv2_0 = SO3Convolution(nfeature_in=self.feature_out2,
nfeature_out=self.feature_out3,
b_in=self.bandwidth_out2,
b_out=self.bandwidth_out3,
grid=grid_so3)
self.conv2_1 = SO3Convolution(nfeature_in=self.feature_out2,
nfeature_out=self.feature_out3,
b_in=self.bandwidth_out2,
b_out=self.bandwidth_out3,
grid=grid_so3)
self.conv2_2 = SO3Convolution(nfeature_in=self.feature_out2,
nfeature_out=self.feature_out3,
b_in=self.bandwidth_out2,
b_out=self.bandwidth_out3,
grid=grid_so3)
self.bn2_0 = nn.BatchNorm3d(num_features=self.feature_out3)
self.bn2_1 = nn.BatchNorm3d(num_features=self.feature_out3)
self.bn2_2 = nn.BatchNorm3d(num_features=self.feature_out3)
# so3 conv (3) [Rotation Invariant] -----------------------------------------------
self.conv3_0 = SO3Convolution(nfeature_in=self.feature_out3,
nfeature_out=self.feature_out4,
b_in=self.bandwidth_out3,
b_out=self.bandwidth_out4,
grid=grid_so3)
self.conv3_1 = SO3Convolution(nfeature_in=self.feature_out3,
nfeature_out=self.feature_out4,
b_in=self.bandwidth_out3,
b_out=self.bandwidth_out4,
grid=grid_so3)
self.conv3_2 = SO3Convolution(nfeature_in=self.feature_out3,
nfeature_out=self.feature_out4,
b_in=self.bandwidth_out3,
b_out=self.bandwidth_out4,
grid=grid_so3)
self.bn3_0 = nn.BatchNorm3d(num_features=self.feature_out4)
self.bn3_1 = nn.BatchNorm3d(num_features=self.feature_out4)
self.bn3_2 = nn.BatchNorm3d(num_features=self.feature_out4)
self.weights = nn.Parameter(
nn.init.uniform_(torch.Tensor(self.feature_out4, self.num_grids)))
self.out_layer = nn.Sequential(
nn.Linear(self.feature_out4, int(self.feature_out4 / 2)),
nn.ReLU(), nn.Linear(int(self.feature_out4 / 2), 10))
def __init__(self, params):
super(S2ConvNet, self).__init__()
grid_s2 = s2_equatorial_grid()
if (params.so3grid):
grid_so3 = so3_equatorial_grid()
else:
grid_so3 = so3_near_identity_grid()
self.conv1 = S2Convolution(nfeature_in=34,
nfeature_out=params.s2_1,
b_in=13,
b_out=10,
grid=grid_s2)
self.conv2 = SO3Convolution(nfeature_in=params.s2_1,
nfeature_out=params.so3_2,
b_in=10,
b_out=8,
grid=grid_so3)
self.conv3 = SO3Convolution(nfeature_in=params.so3_2,
nfeature_out=params.so3_3,
b_in=8,
b_out=5,
grid=grid_so3)
self.conv4 = SO3Convolution(nfeature_in=params.so3_3,
nfeature_out=params.so3_4,
b_in=5,
b_out=3,
grid=grid_so3)
self.conv5 = SO3Convolution(nfeature_in=params.so3_4,
nfeature_out=params.so3_5,
b_in=3,
b_out=2,
grid=grid_so3)
last_entry = params.so3_3
if (params.if_so3_4) and (params.if_so3_5):
last_entry = params.so3_5
elif (params.if_so3_4):
last_entry = params.so3_4
self.fc_layer = nn.Linear(last_entry, params.fc1)
self.fc_layer_2 = nn.Linear(params.fc1, params.fc2)
self.fc_layer_3 = nn.Linear(params.fc2, params.fc3)
self.fc_layer_4 = nn.Linear(params.fc3, params.fc4)
self.fc_layer_5 = nn.Linear(params.fc4, params.fc5)
self.norm_layer_2d_1 = nn.BatchNorm2d(34)
self.norm_1d_1 = nn.BatchNorm1d(params.fc1)
self.norm_1d_2 = nn.BatchNorm1d(params.fc2)
self.norm_1d_3 = nn.BatchNorm1d(params.fc3)
self.norm_1d_4 = nn.BatchNorm1d(params.fc4)
self.norm_1d_5 = nn.BatchNorm1d(params.fc5)
self.norm_1d_6 = nn.BatchNorm1d(1)
last_fc_entry = params.fc3
if params.if_fc_4 and params.if_fc_5:
last_fc_entry = params.fc5
elif params.if_fc_4:
last_fc_entry = params.fc4
#print(last_fc_entry, params.if_fc_4, params.if_fc_5, params.fc3,params.fc4,params.fc5, "Aoba=================")
self.fc_layer_6 = nn.Linear(last_fc_entry, 1)
self.do1 = nn.Dropout(params.do1r)
self.do2 = nn.Dropout(params.do2r)
self.do3 = nn.Dropout(params.do3r)
self.do4 = nn.Dropout(params.do4r)
self.do5 = nn.Dropout(params.do5r)
self.if_so3_4 = params.if_so3_4
self.if_so3_5 = params.if_so3_5
self.if_fc_4 = params.if_fc_4
self.if_fc_5 = params.if_fc_5
def __init__(self):
super(S2Model, self).__init__()
self.leaky_alpha = 0.1
grid_s2 = s2_near_identity_grid(max_beta=np.pi / 64,
n_alpha=4,
n_beta=2)
self.layer0 = nn.Sequential(
S2Convolution(3, 16, 128, 64, grid_s2),
nn.GroupNorm(1, 16),
nn.LeakyReLU(self.leaky_alpha, inplace=True),
)
grid_so3 = so3_near_identity_grid(max_beta=np.pi / 32,
max_gamma=0,
n_alpha=4,
n_beta=2,
n_gamma=1)
self.layer1 = nn.Sequential(
SO3Convolution(16, 16, 64, 32, grid_so3),
nn.GroupNorm(1, 16),
nn.LeakyReLU(self.leaky_alpha, inplace=True),
# SO3Convolution(16, 32, 32, 32, grid_so3),
# nn.GroupNorm(2, 32),
# nn.LeakyReLU(self.leaky_alpha, inplace=True)
)
grid_so3 = so3_near_identity_grid(max_beta=np.pi / 16,
max_gamma=0,
n_alpha=4,
n_beta=2,
n_gamma=1)
self.layer2 = nn.Sequential(
SO3Convolution(16, 32, 32, 16, grid_so3),
nn.GroupNorm(2, 32),
nn.LeakyReLU(self.leaky_alpha, inplace=True),
# SO3Convolution(32, 64, 16, 16, grid_so3),
# nn.GroupNorm(4, 64),
# nn.LeakyReLU(self.leaky_alpha, inplace=True),
)
grid_so3 = so3_near_identity_grid(max_beta=np.pi / 8,
max_gamma=0,
n_alpha=4,
n_beta=2,
n_gamma=1)
self.layer3 = nn.Sequential(
SO3Convolution(32, 64, 16, 8, grid_so3),
nn.GroupNorm(4, 64),
nn.LeakyReLU(self.leaky_alpha, inplace=True),
# SO3Convolution(64, 128, 8, 8, grid_so3),
# nn.GroupNorm(8, 128),
# nn.LeakyReLU(self.leaky_alpha, inplace=True),
)
grid_so3 = so3_near_identity_grid(max_beta=np.pi / 16,
max_gamma=0,
n_alpha=4,
n_beta=2,
n_gamma=1)
self.layer4 = nn.Sequential(
SO3Convolution(64, 128, 8, 8, grid_so3), nn.GroupNorm(8, 128),
nn.LeakyReLU(self.leaky_alpha, inplace=True))
# self.score_layers = nn.Sequential(
# nn.Conv2d(128, 64, 3, stride=2, padding=1, bias=False),
# nn.BatchNorm2d(64),
# nn.MaxPool2d(2, 2),
# nn.LeakyReLU(self.leaky_alpha, inplace=True),
# nn.Conv2d(64, 32, 3, stride=2, padding=1, bias=False),
# nn.BatchNorm2d(32),
# nn.MaxPool2d(2, 2),
# nn.AdaptiveAvgPool2d(1)
# )
self.fc = nn.Sequential(nn.Linear(32768, 32, bias=False),
nn.GroupNorm(1, 32), nn.Linear(32, 1))
def __init__(self, time_channel):
super(S2ConvNet, self).__init__()
# Two types of grid used to sample on Wigner D matrix
# From hyperparameter search, the equatorial grid and near identity grid does not
# change performance a lot
grid_dict = {
's2_eq': s2_equatorial_grid,
's2_ni': s2_near_identity_grid,
"so3_eq": so3_equatorial_grid,
'so3_ni': so3_near_identity_grid
}
s2_grid_type = 's2_eq'
grid_s2 = grid_dict[s2_grid_type]()
so3_grid_type = 'so3_eq'
grid_so3 = grid_dict[so3_grid_type]()
# number of neuron in each layer
s2_1 = 20
so3_2 = 40
so3_3 = 60
so3_4 = 120
so3_5 = 200
so3_6 = 256
so3_numlayers = 'six'
# neuron of FC layer
fc1 = 256
fc2 = 190
fc3 = 128
fc4 = 64
fc5 = 12
fc_numlayers = 'five_fc'
do1r = 0.2000025346505366
do2r = 0.2000020575803182
do3r = 0.2000023878002161
do4r = 0.19999793704295525
do5r = 0.11299997708596853
do1r = min(max(do1r, 0.0), 1.0)
do2r = min(max(do2r, 0.0), 1.0)
do3r = min(max(do3r, 0.0), 1.0)
do4r = min(max(do4r, 0.0), 1.0)
do5r = min(max(do5r, 0.0), 1.0)
layer_dict = {
'four': so3_4,
'five': so3_5,
'six': so3_6,
'three_fc': fc3,
'four_fc': fc4,
'five_fc': fc5
}
last_entry = layer_dict[so3_numlayers]
last_fc_entry = layer_dict[fc_numlayers]
# Bandwidth: a concept controlling the size of output feature map
# the feature map size is (2*bandwidth, 2*bandwidth, 2*bandwidth)
bw = np.linspace(time_channel[1] / 2, 2, 7).astype(int)
self.conv1 = S2Convolution(nfeature_in=time_channel[0],
nfeature_out=s2_1,
b_in=bw[0],
b_out=bw[1],
grid=grid_s2)
self.conv2 = SO3Convolution(nfeature_in=s2_1,
nfeature_out=so3_2,
b_in=bw[1],
b_out=bw[2],
grid=grid_so3)
self.conv3 = SO3Convolution(nfeature_in=so3_2,
nfeature_out=so3_3,
b_in=bw[2],
b_out=bw[3],
grid=grid_so3)
self.conv4 = SO3Convolution(nfeature_in=so3_3,
nfeature_out=so3_4,
b_in=bw[3],
b_out=bw[4],
grid=grid_so3)
self.conv5 = SO3Convolution(nfeature_in=so3_4,
nfeature_out=so3_5,
b_in=bw[4],
b_out=bw[5],
grid=grid_so3)
self.conv6 = SO3Convolution(nfeature_in=so3_5,
nfeature_out=so3_6,
b_in=bw[5],
b_out=bw[6],
grid=grid_so3)
self.fc_layer = nn.Linear(so3_6, fc1)
# self.fc_layer_0 = nn.Linear(fcn1, fc0)
# self.fc_layer_1 = nn.Linear(fc0, fc1)
self.fc_layer_2 = nn.Linear(fc1, fc2)
self.fc_layer_3 = nn.Linear(fc2, fc3)
self.fc_layer_4 = nn.Linear(fc3, fc4)
self.fc_layer_5 = nn.Linear(fc4, fc5)
self.norm_layer_2d_1 = nn.BatchNorm3d(s2_1)
self.norm_layer_2d_2 = nn.BatchNorm3d(so3_2)
self.norm_layer_2d_3 = nn.BatchNorm3d(so3_3)
self.norm_layer_2d_4 = nn.BatchNorm3d(so3_4)
self.norm_layer_2d_5 = nn.BatchNorm3d(so3_5)
self.norm_layer_2d_6 = nn.BatchNorm3d(so3_6)
self.norm_1d_1 = nn.BatchNorm1d(fc1)
self.norm_1d_2 = nn.BatchNorm1d(fc2)
self.norm_1d_3 = nn.BatchNorm1d(fc3)
self.norm_1d_4 = nn.BatchNorm1d(fc4)
self.norm_1d_5 = nn.BatchNorm1d(fc5)
self.norm_1d_6 = nn.BatchNorm1d(1)
self.fc_layer_6 = nn.Linear(last_fc_entry, 1)
self.do1 = nn.Dropout(do1r)
self.do2 = nn.Dropout(do2r)
self.do3 = nn.Dropout(do3r)
self.do4 = nn.Dropout(do4r)
self.do5 = nn.Dropout(do5r)
self.sdo1 = nn.Dropout(do1r)
self.sdo2 = nn.Dropout(do2r)
self.sdo3 = nn.Dropout(do3r)
self.sdo4 = nn.Dropout(do4r)
self.sdo5 = nn.Dropout(do5r)
self.so3_numlayers = so3_numlayers
self.fc_numlayers = fc_numlayers
def __init__(self, bandwidth=30):
super(S2ConvNet_deep, self).__init__()
grid_s2 = s2_near_identity_grid(n_alpha=6,
max_beta=np.pi / 16,
n_beta=1)
grid_so3_1 = so3_near_identity_grid(n_alpha=6,
max_beta=np.pi / 16,
n_beta=1,
max_gamma=2 * np.pi,
n_gamma=6)
grid_so3_2 = so3_near_identity_grid(n_alpha=6,
max_beta=np.pi / 8,
n_beta=1,
max_gamma=2 * np.pi,
n_gamma=6)
grid_so3_3 = so3_near_identity_grid(n_alpha=6,
max_beta=np.pi / 4,
n_beta=1,
max_gamma=2 * np.pi,
n_gamma=6)
grid_so3_4 = so3_near_identity_grid(n_alpha=6,
max_beta=np.pi / 2,
n_beta=1,
max_gamma=2 * np.pi,
n_gamma=6)
self.convolutional = nn.Sequential(
S2Convolution(nfeature_in=1,
nfeature_out=8,
b_in=bandwidth,
b_out=bandwidth,
grid=grid_s2), nn.ReLU(inplace=False),
SO3Convolution(nfeature_in=8,
nfeature_out=16,
b_in=bandwidth,
b_out=bandwidth // 2,
grid=grid_so3_1), nn.ReLU(inplace=False),
SO3Convolution(nfeature_in=16,
nfeature_out=16,
b_in=bandwidth // 2,
b_out=bandwidth // 2,
grid=grid_so3_2), nn.ReLU(inplace=False),
SO3Convolution(nfeature_in=16,
nfeature_out=24,
b_in=bandwidth // 2,
b_out=bandwidth // 4,
grid=grid_so3_2), nn.ReLU(inplace=False),
SO3Convolution(nfeature_in=24,
nfeature_out=24,
b_in=bandwidth // 4,
b_out=bandwidth // 4,
grid=grid_so3_3), nn.ReLU(inplace=False),
SO3Convolution(nfeature_in=24,
nfeature_out=32,
b_in=bandwidth // 4,
b_out=bandwidth // 8,
grid=grid_so3_3), nn.ReLU(inplace=False),
SO3Convolution(nfeature_in=32,
nfeature_out=64,
b_in=bandwidth // 8,
b_out=bandwidth // 8,
grid=grid_so3_4), nn.ReLU(inplace=False))
self.linear = nn.Sequential(
# linear 1
nn.BatchNorm1d(64),
nn.Linear(in_features=64, out_features=64),
nn.ReLU(inplace=False),
# linear 2
nn.BatchNorm1d(64),
nn.Linear(in_features=64, out_features=32),
nn.ReLU(inplace=False),
# linear 3
nn.BatchNorm1d(32),
nn.Linear(in_features=32, out_features=10))
def __init__(self, n_features, bandwidth=100):
super().__init__(
) # call the initialization function of father class (nn.Module)
self.features = [n_features, 10, 20, 60, 100, 200]
#self.bandwidths = [bandwidth, 50, 25, 20, 10, 5]
self.bandwidths = [bandwidth, 50, 40, 30, 20, 5]
#self.bandwidths = [bandwidth, 100, 50, 25, 20, 10]
assert len(self.bandwidths) == len(self.features)
grid_s2 = s2_near_identity_grid(n_alpha=6,
max_beta=np.pi / 160,
n_beta=1)
grid_so3_1 = so3_near_identity_grid(n_alpha=6,
max_beta=np.pi / 16,
n_beta=1,
max_gamma=2 * np.pi,
n_gamma=6)
grid_so3_2 = so3_near_identity_grid(n_alpha=6,
max_beta=np.pi / 8,
n_beta=1,
max_gamma=2 * np.pi,
n_gamma=6)
grid_so3_3 = so3_near_identity_grid(n_alpha=6,
max_beta=np.pi / 4,
n_beta=1,
max_gamma=2 * np.pi,
n_gamma=6)
grid_so3_4 = so3_near_identity_grid(n_alpha=6,
max_beta=np.pi / 2,
n_beta=1,
max_gamma=2 * np.pi,
n_gamma=6)
# grid_so3_4 = so3_near_identity_grid(n_alpha=6, max_beta=np.pi/ 2, n_beta=1, max_gamma=2*np.pi, n_gamma=6)
self.convolutional = nn.Sequential(
S2Convolution(nfeature_in=self.features[0],
nfeature_out=self.features[1],
b_in=self.bandwidths[0],
b_out=self.bandwidths[1],
grid=grid_s2),
nn.PReLU(),
nn.BatchNorm3d(self.features[1], affine=True),
SO3Convolution(nfeature_in=self.features[1],
nfeature_out=self.features[2],
b_in=self.bandwidths[1],
b_out=self.bandwidths[2],
grid=grid_so3_1),
nn.PReLU(),
nn.BatchNorm3d(self.features[2], affine=True),
SO3Convolution(nfeature_in=self.features[2],
nfeature_out=self.features[3],
b_in=self.bandwidths[2],
b_out=self.bandwidths[3],
grid=grid_so3_2),
nn.PReLU(),
nn.BatchNorm3d(self.features[3], affine=True),
SO3Convolution(nfeature_in=self.features[3],
nfeature_out=self.features[4],
b_in=self.bandwidths[3],
b_out=self.bandwidths[4],
grid=grid_so3_3),
nn.BatchNorm3d(self.features[4], affine=True),
nn.PReLU(),
SO3Convolution(nfeature_in=self.features[4],
nfeature_out=self.features[5],
b_in=self.bandwidths[4],
b_out=self.bandwidths[5],
grid=grid_so3_4),
nn.BatchNorm3d(self.features[5], affine=True),
nn.PReLU(),
)
self.linear = nn.Sequential(
# linear 1
nn.BatchNorm1d(self.features[5]),
nn.Linear(in_features=self.features[5], out_features=512),
nn.PReLU(),
nn.Dropout(p=0.4),
# linear 2
nn.BatchNorm1d(512),
nn.Linear(in_features=512, out_features=256),
nn.PReLU(),
nn.Dropout(p=0.4),
# linear 3
nn.BatchNorm1d(256),
nn.Linear(in_features=256, out_features=256),
)
def __init__(self, time_channel):
super(S2ConvNet, self).__init__()
grid_dict = {
's2_eq': s2_equatorial_grid,
's2_ni': s2_near_identity_grid,
"so3_eq": so3_equatorial_grid,
'so3_ni': so3_near_identity_grid
}
s2_grid_type = 's2_eq'
grid_s2 = grid_dict[s2_grid_type]()
so3_grid_type = 'so3_eq'
grid_so3 = grid_dict[so3_grid_type]()
s2_1 = 64
so3_2 = 96
so3_3 = 128
so3_4 = 160
so3_5 = 200
so3_6 = 256
so3_numlayers = 'six'
bias_injection = False
np.random.seed(None)
bias_seed = np.random.randint(100000)
#287
#72715
#83675
#bias_seed = 287
print('Bias Seed', bias_seed)
global SEED
SEED = bias_seed
#@nevergrad@ bias_seed = NG_G{777, 534}
bias_seed = abs(int(bias_seed))
np.random.seed(bias_seed)
bias1 = np.random.randn(s2_1)
bias2 = np.random.randn(so3_2)
bias3 = np.random.randn(so3_3)
bias4 = np.random.randn(so3_4)
bias5 = np.random.randn(so3_5)
bias6 = np.random.randn(so3_6)
# fcn1 = 1024
# fc0 = 512
fc1 = 256
fc2 = 190
fc3 = 128
fc4 = 64
fc5 = 12
fc_numlayers = 'five_fc'
do1r = 0.2000025346505366
#@nevergrad@ do1r = NG_G{0.5, 0.3}
do2r = 0.2000020575803182
#@nevergrad@ do2r = NG_G{0.5, 0.3}
do3r = 0.2000023878002161
#@nevergrad@ do3r = NG_G{0.5, 0.3}
do4r = 0.19999793704295525
#@nevergrad@ do4r = NG_G{0.5, 0.3}
do5r = 0.11299997708596853
#@nevergrad@ do5r = NG_G{0.5, 0.3}
do1r = min(max(do1r, 0.0), 1.0)
do2r = min(max(do2r, 0.0), 1.0)
do3r = min(max(do3r, 0.0), 1.0)
do4r = min(max(do4r, 0.0), 1.0)
do5r = min(max(do5r, 0.0), 1.0)
layer_dict = {
'four': so3_4,
'five': so3_5,
'six': so3_6,
'three_fc': fc3,
'four_fc': fc4,
'five_fc': fc5
}
last_entry = layer_dict[so3_numlayers]
last_fc_entry = layer_dict[fc_numlayers]
self.conv1 = S2Convolution(nfeature_in=time_channel,
nfeature_out=s2_1,
b_in=20,
b_out=17,
grid=grid_s2,
use_bias=bias_injection,
input_bias=bias1)
self.conv2 = SO3Convolution(nfeature_in=s2_1,
nfeature_out=so3_2,
b_in=17,
b_out=15,
grid=grid_so3,
use_bias=bias_injection,
input_bias=bias2)
self.conv3 = SO3Convolution(nfeature_in=so3_2,
nfeature_out=so3_3,
b_in=15,
b_out=12,
grid=grid_so3,
use_bias=bias_injection,
input_bias=bias3)
self.conv4 = SO3Convolution(nfeature_in=so3_3,
nfeature_out=so3_4,
b_in=12,
b_out=10,
grid=grid_so3,
use_bias=bias_injection,
input_bias=bias4)
self.conv5 = SO3Convolution(nfeature_in=so3_4,
nfeature_out=so3_5,
b_in=10,
b_out=7,
grid=grid_so3,
use_bias=bias_injection,
input_bias=bias5)
self.conv6 = SO3Convolution(nfeature_in=so3_5,
nfeature_out=so3_6,
b_in=7,
b_out=3,
grid=grid_so3,
use_bias=bias_injection,
input_bias=bias6)
print(last_entry)
self.fc_layer = nn.Linear(so3_6, fc1)
# self.fc_layer_0 = nn.Linear(fcn1, fc0)
# self.fc_layer_1 = nn.Linear(fc0, fc1)
self.fc_layer_2 = nn.Linear(fc1, fc2)
self.fc_layer_3 = nn.Linear(fc2, fc3)
self.fc_layer_4 = nn.Linear(fc3, fc4)
self.fc_layer_5 = nn.Linear(fc4, fc5)
self.norm_layer_2d_1 = nn.BatchNorm3d(s2_1)
self.norm_layer_2d_2 = nn.BatchNorm3d(so3_2)
self.norm_layer_2d_3 = nn.BatchNorm3d(so3_3)
self.norm_layer_2d_4 = nn.BatchNorm3d(so3_4)
self.norm_layer_2d_5 = nn.BatchNorm3d(so3_5)
self.norm_layer_2d_6 = nn.BatchNorm3d(so3_6)
self.norm_1d_1 = nn.BatchNorm1d(fc1)
self.norm_1d_2 = nn.BatchNorm1d(fc2)
self.norm_1d_3 = nn.BatchNorm1d(fc3)
self.norm_1d_4 = nn.BatchNorm1d(fc4)
self.norm_1d_5 = nn.BatchNorm1d(fc5)
self.norm_1d_6 = nn.BatchNorm1d(1)
self.fc_layer_6 = nn.Linear(last_fc_entry, 1)
self.do1 = nn.Dropout(do1r)
self.do2 = nn.Dropout(do2r)
self.do3 = nn.Dropout(do3r)
self.do4 = nn.Dropout(do4r)
self.do5 = nn.Dropout(do5r)
self.sdo1 = nn.Dropout(do1r)
self.sdo2 = nn.Dropout(do2r)
self.sdo3 = nn.Dropout(do3r)
self.sdo4 = nn.Dropout(do4r)
self.sdo5 = nn.Dropout(do5r)
self.so3_numlayers = so3_numlayers
self.fc_numlayers = fc_numlayers
def __init__(self, final_output_dim, sconv_dims, bandwidths
#BANDWIDTH_IN,
#BANDWIDTH_OUT,
#R_IN=64,
#sconv_intermed_dims
):
super().__init__()
self.per_point_dims = [32, 32]
#self.features = [R_IN, 40, 40, dim_output]
self.features = list(sconv_dims)
self.bandwidths = bandwidths #[BANDWIDTH_IN, 32, 32, BANDWIDTH_OUT]
#self.linear1 = nn.Linear(dim_output, 50)
#self.linear2 = nn.Linear(50, 50)
print('Building PRIN-based AE Encoder')
print('\tFeatures:', self.features)
print('\tInput Dim:', final_output_dim)
print('\tBandwidths:', self.bandwidths)
### SConv module ###
sequence = []
# S2 layer
grid = s2_equatorial_grid(max_beta=0,
n_alpha=2 * self.bandwidths[0],
n_beta=1)
sequence.append(
S2Convolution(self.features[0], self.features[1],
self.bandwidths[0], self.bandwidths[1], grid))
# SO3 layers
for l in range(1, len(self.features) - 1):
nfeature_in = self.features[l]
nfeature_out = self.features[l + 1]
b_in = self.bandwidths[l]
b_out = self.bandwidths[l + 1]
sequence.append(nn.BatchNorm3d(nfeature_in, affine=True))
sequence.append(nn.ReLU())
grid = so3_equatorial_grid(max_beta=0,
max_gamma=0,
n_alpha=2 * b_in,
n_beta=1,
n_gamma=1)
sequence.append(
SO3Convolution(nfeature_in, nfeature_out, b_in, b_out, grid))
# Final BN and non-linearity
sequence.append(nn.BatchNorm3d(self.features[-1], affine=True))
sequence.append(nn.ReLU())
self.sequential = nn.Sequential(*sequence)
# TODO maybe just operate on the voxel and pool that directly, without converting to a PC first?
# Per-point processor
self.per_point_processor = nn.Sequential(
nn.Conv1d(self.features[-1], self.per_point_dims[0], 1,
bias=False), nn.BatchNorm1d(self.per_point_dims[0]),
nn.ReLU(False),
nn.Conv1d(self.per_point_dims[0],
self.per_point_dims[1],
1,
bias=False), nn.BatchNorm1d(self.per_point_dims[1]),
nn.ReLU(False))
# Final pooled processor
self.pooled_processor_dims = [
2 * self.per_point_dims[-1], final_output_dim, final_output_dim
]
self.pooled_processor = nn.Sequential(
nn.Linear(self.pooled_processor_dims[0],
self.pooled_processor_dims[1]),
nn.BatchNorm1d(self.pooled_processor_dims[1]),
nn.ReLU(),
nn.Linear(self.pooled_processor_dims[1],
self.pooled_processor_dims[2]),
)
#self.protlin = nn.Linear(2*self.features[-1], final_output_dim)
self.protlin = nn.Sequential(
nn.Linear(2 * self.features[-1], 2 * final_output_dim),
nn.BatchNorm1d(2 * final_output_dim), nn.ReLU(),
nn.Linear(final_output_dim * 2, final_output_dim))
s2_grid = s2_equatorial_grid(max_beta=0, n_alpha=64, n_beta=1)
s2_conv = S2Convolution(nfeature_in=12,
nfeature_out=15,
b_in=64,
b_out=32,
grid=s2_grid)
s2_conv.to(device)
so3_grid = so3_equatorial_grid(max_beta=0,
max_gamma=0,
n_alpha=64,
n_beta=1,
n_gamma=1)
so3_conv = SO3Convolution(nfeature_in=15,
nfeature_out=21,
b_in=32,
b_out=24,
grid=so3_grid)
so3_conv.to(device)
def phi(x):
x = s2_conv(x)
x = torch.nn.functional.relu(x)
x = so3_conv(x)
return x
def rot(x, angle):
# rotate the signal around the Z axis
n = round(x.size(3) * angle / 360)
def __init__(self):
super(Model, self).__init__()
self.leaky_alpha = 0.1
# S2 layer
grid = s2_near_identity_grid(max_beta=np.pi / 64, n_alpha=4, n_beta=2)
self.layer0 = nn.Sequential(
S2Convolution(3, 16, 128, 64, grid),
nn.GroupNorm(1, 16),
nn.LeakyReLU(self.leaky_alpha, inplace=True),
)
self.flow_layer0 = nn.Sequential(
S2Convolution(2, 16, 128, 64, grid),
nn.GroupNorm(1, 16),
nn.LeakyReLU(self.leaky_alpha, inplace=True),
)
grid = so3_near_identity_grid(max_beta=np.pi / 32,
max_gamma=0,
n_alpha=4,
n_beta=2,
n_gamma=1)
self.layer1, self.flow_layer1 = (nn.Sequential(
SO3Convolution(16, 16, 64, 32, grid),
nn.GroupNorm(1, 16),
nn.LeakyReLU(self.leaky_alpha, inplace=True),
SO3Convolution(16, 32, 32, 32, grid),
nn.GroupNorm(2, 32),
nn.LeakyReLU(self.leaky_alpha, inplace=True),
) for _ in range(2))
grid = so3_near_identity_grid(max_beta=np.pi / 16,
max_gamma=0,
n_alpha=4,
n_beta=2,
n_gamma=1)
self.layer2, self.flow_layer2 = (nn.Sequential(
SO3Convolution(32, 32, 32, 16, grid),
nn.GroupNorm(2, 32),
nn.LeakyReLU(self.leaky_alpha, inplace=True),
SO3Convolution(32, 64, 16, 16, grid),
nn.GroupNorm(4, 64),
nn.LeakyReLU(self.leaky_alpha, inplace=True),
) for _ in range(2))
grid = so3_near_identity_grid(max_beta=np.pi / 8,
max_gamma=0,
n_alpha=4,
n_beta=2,
n_gamma=1)
self.layer3, self.flow_layer3 = (nn.Sequential(
SO3Convolution(64, 64, 16, 8, grid),
nn.GroupNorm(4, 64),
nn.LeakyReLU(self.leaky_alpha, inplace=True),
SO3Convolution(64, 128, 8, 8, grid),
nn.GroupNorm(8, 128),
nn.LeakyReLU(self.leaky_alpha, inplace=True),
) for _ in range(2))
grid = so3_near_identity_grid(max_beta=np.pi / 16,
max_gamma=0,
n_alpha=4,
n_beta=2,
n_gamma=1)
self.layer4 = nn.Sequential(
SO3Convolution(256, 128, 8, 8, grid),
nn.GroupNorm(8, 128),
nn.LeakyReLU(self.leaky_alpha, inplace=True),
)
self.weight_layer = nn.Sequential(
nn.Conv2d(129, 1, kernel_size=1, stride=1, bias=False), )
self.refine_layer = nn.Sequential(
nn.Conv2d(129, 2, kernel_size=1, stride=1, bias=False), )
self.motion_layer1 = nn.Sequential(
nn.Conv2d(256, 32, 3, stride=2, padding=1, bias=True),
nn.LeakyReLU(self.leaky_alpha, inplace=True),
nn.Conv2d(32, 8, 3, stride=2, padding=1, bias=False),
)
self.motion_layer2 = nn.Linear(128, 2, bias=False)
self.control_layer = nn.Sequential(
nn.Conv2d(128, 129, 1, 1, 0, bias=False), nn.Sigmoid())
self.softmax = nn.Softmax(dim=-1)