sphvlad with bn
| 15 | |
| 16 | |
| 17 | class SphereVLAD(nn.Module): |
| 18 | ''' |
| 19 | sphvlad with bn |
| 20 | ''' |
| 21 | def __init__(self, config): |
| 22 | super(SphereVLAD, self).__init__() |
| 23 | |
| 24 | bandwidth = (int)(config.DATA.SPH_IM.IMAGE_SIZE[0]/2.0) |
| 25 | |
| 26 | grid_s2 = s2_near_identity_grid(n_alpha=6, max_beta=np.pi/16, n_beta=1) |
| 27 | grid_so3_1 = so3_near_identity_grid( |
| 28 | n_alpha=6, max_beta=np.pi / 16, n_beta=1, max_gamma=2*np.pi, n_gamma=6) |
| 29 | grid_so3_2 = so3_near_identity_grid( |
| 30 | n_alpha=6, max_beta=np.pi / 8, n_beta=1, max_gamma=2*np.pi, n_gamma=6) |
| 31 | grid_so3_3 = so3_near_identity_grid( |
| 32 | n_alpha=6, max_beta=np.pi / 4, n_beta=1, max_gamma=2*np.pi, n_gamma=6) |
| 33 | grid_so3_4 = so3_near_identity_grid( |
| 34 | n_alpha=6, max_beta=np.pi / 2, n_beta=1, max_gamma=2*np.pi, n_gamma=6) |
| 35 | |
| 36 | self.conv1 = nn.Sequential( |
| 37 | S2Convolution(nfeature_in=3, nfeature_out=8, |
| 38 | b_in=bandwidth, b_out=bandwidth//2, grid=grid_s2), |
| 39 | nn.BatchNorm3d(8), |
| 40 | nn.ReLU()) |
| 41 | |
| 42 | self.conv2 = nn.Sequential( |
| 43 | SO3Convolution(nfeature_in=8, nfeature_out=8, |
| 44 | b_in=bandwidth//2, b_out=bandwidth//4, grid=grid_so3_1), |
| 45 | nn.BatchNorm3d(8), |
| 46 | nn.ReLU()) |
| 47 | |
| 48 | self.conv3 = nn.Sequential( |
| 49 | SO3Convolution(nfeature_in=8, nfeature_out=16, |
| 50 | b_in=bandwidth//4, b_out=bandwidth//8, grid=grid_so3_2), |
| 51 | nn.BatchNorm3d(16), |
| 52 | nn.ReLU()) |
| 53 | |
| 54 | self.conv4 = nn.Sequential( |
| 55 | SO3Convolution(nfeature_in=16, nfeature_out=16, |
| 56 | b_in=bandwidth//8, b_out=bandwidth//8, grid=grid_so3_3), |
| 57 | nn.BatchNorm3d(16), |
| 58 | nn.ReLU()) |
| 59 | |
| 60 | self.vlad = NetVLAD(num_clusters=config.MODEL.NETVLAD.CLUSTER_NUM, |
| 61 | dim=config.MODEL.NETVLAD.FEATURE_DIM, |
| 62 | normalize_input=config.MODEL.NETVLAD.NORMALIZE_INPUT, |
| 63 | output_dim=config.MODEL.NETVLAD.OUTPUT_DIM, |
| 64 | gate=config.MODEL.NETVLAD.GATE) |
| 65 | |
| 66 | |
| 67 | def forward(self, x): |
| 68 | # encoder |
| 69 | x = self.conv1(x) |
| 70 | x = self.conv2(x) |
| 71 | x = self.conv3(x) |
| 72 | x = self.conv4(x) |
| 73 | # reshape |
| 74 | x = x.view(x.shape[0], x.shape[1], x.shape[2], x.shape[3]*x.shape[4]) |
nothing calls this directly
no outgoing calls
no test coverage detected