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, 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, channels, classes, imagesize, **kwargs):
super(S2ConvNet, self).__init__()
f1 = 20
f2 = 40
f_output = classes
b_in = imagesize[0]/2
b_l1 = 10
b_l2 = 6
grid_s2 = s2cnn.s2_near_identity_grid()
grid_so3 = s2cnn.so3_near_identity_grid()
self.conv1 = s2cnn.S2Convolution(
nfeature_in=1,
nfeature_out=f1,
b_in=b_in,
b_out=b_l1,
grid=grid_s2)
self.conv2 = s2cnn.SO3Convolution(
nfeature_in=f1,
nfeature_out=f2,
b_in=b_l1,
b_out=b_l2,
grid=grid_so3)
self.out_layer = torch.nn.Linear(f2, 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, 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, f_in, f_hidden, bandwidth):
"""
Initialize the ConvLSTM cell
:param dtype: torch.cuda.FloatTensor or torch.FloatTensor
Whether or not to use cuda.
"""
super().__init__()
self.input_features = f_in
self.hidden_features = f_hidden
self.bandwidth = bandwidth
grid = s2_near_identity_grid(n_alpha=2 * bandwidth)
self.reset_gate = S2Convolution(f_in + f_hidden, f_hidden, bandwidth,
bandwidth, grid)
self.update_gate = S2Convolution(f_in + f_hidden, f_hidden, bandwidth,
bandwidth, grid)
self.output_gate = S2Convolution(f_in + f_hidden, f_hidden, bandwidth,
bandwidth, grid)
def __init__(self, nclasses):
super().__init__()
self.features = [2, 100, 100, nclasses]
self.bandwidths = [64, 16, 10]
assert len(self.bandwidths) == len(self.features) - 1
sequence = []
# S2 layer
grid = s2_near_identity_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_near_identity_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 __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, bandwidths, features, use_equatorial_grid):
super().__init__()
self.bandwidths = bandwidths
self.features = features
assert len(self.bandwidths) == len(self.features)
sequence = []
# S2 layer
grid_s2 = s2_equatorial_grid(
max_beta=0, n_alpha=2 * self.bandwidths[0],
n_beta=1) if use_equatorial_grid else s2_near_identity_grid(
max_beta=np.pi / 8, n_alpha=8, n_beta=3)
sequence.append(
S2Convolution(self.features[0], self.features[1],
self.bandwidths[0], self.bandwidths[1], grid_s2))
sequence.append(nn.BatchNorm3d(self.features[-1], affine=True))
sequence.append(nn.ReLU())
self.sequential = nn.Sequential(*sequence)
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):
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) #, max_beta=0.2)
sequence.append(
S2Convolution(self.features[0], self.features[1],
self.bandwidths[0], self.bandwidths[1], grid))
sequence.append(nn.GroupNorm(self.groups[0], self.features[1]))
sequence.append(nonlinear())
# SO3 layers
for l in range(1, len(self.features) - 1, 2):
f_in = self.features[l]
f_mid = self.features[l + 1]
f_out = self.features[l + 2]
b_in = self.bandwidths[l]
b_mid = self.bandwidths[l + 1]
b_out = self.bandwidths[l + 2]
n_group_mid = self.groups[l]
n_group_out = self.groups[l + 1]
# sequence.append(so3_residual_block(f_in, f_mid, f_out, b_in, b_mid, b_out, n_group_mid, n_group_out))
sequence.append(
so3_residual_block(f_in, f_mid, f_out, b_in, b_mid, b_out,
n_group_mid, n_group_out))
self.sequential = nn.Sequential(*sequence)
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):
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))
assume_unique=True,
invert=True)
validation_data = data.sel(sample=validation_set)
train_data = data.sel(sample=train_set)
# Build the data generators
generator = DataGenerator(dlwp, train_data, batch_size=batch_size)
if validation_data is not None:
val_generator = DataGenerator(dlwp, validation_data, batch_size=batch_size)
else:
val_generator = None
#%% Compile the model structure with some generator data information
# Convolutional feed-forward network
s2_grid = s2_near_identity_grid(max_beta=0.2, n_alpha=12, n_beta=1)
truncation = 12
layers = (('S2Convolution', (3, 16, 36, truncation, s2_grid), {
'mean_gamma': True,
'activation': 'tanh'
}), ('S2Convolution', (16, 16, truncation, truncation, s2_grid), {
'mean_gamma': True,
'activation': 'tanh'
}), ('TorchReshape', ((-1, 16 * (2 * truncation)**2), ),
None), ('Linear', (16 * (2 * truncation)**2, generator.n_features), None),
('TorchReshape', ((-1, ) + generator.convolution_shape, ), None))
dlwp.build_model(layers, loss='MSELoss', optimizer='Adam')
#%% Load the scaling and validating data
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):
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)
