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, 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, 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, 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_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, 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): 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): 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, 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(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)
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, 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(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))