| 234 | return x, x1 # [batch, 1024, 3], [batch, 1024, 1024] |
| 235 | |
| 236 | class DGCNN_New(nn.Module): |
| 237 | |
| 238 | def __init__(self, feat_dim): |
| 239 | super(DGCNN_New, self).__init__() |
| 240 | self.n_knn = 20 |
| 241 | num_part = feat_dim |
| 242 | |
| 243 | print("feat_dim is: ", feat_dim) |
| 244 | |
| 245 | pooling = 'mean' |
| 246 | |
| 247 | self.conv1 = NonEquivariantLinearLeakyReLU(2, 64 // 3) |
| 248 | self.conv2 = NonEquivariantLinearLeakyReLU(64 // 3, 64 // 3) |
| 249 | self.conv3 = NonEquivariantLinearLeakyReLU(64 // 3 * 2, 64 // 3) |
| 250 | self.conv4 = NonEquivariantLinearLeakyReLU(64 // 3, 64 // 3) |
| 251 | self.conv5 = NonEquivariantLinearLeakyReLU(64 // 3 * 2, 64 // 3) |
| 252 | self.VnInv = NonEquivariantStdFeature(2 * feat_dim, dim=3, normalize_frame=False) |
| 253 | |
| 254 | if pooling == 'max': |
| 255 | self.pool1 = NonEquivariantMaxPool(64 // 3) |
| 256 | self.pool2 = NonEquivariantMaxPool(64 // 3) |
| 257 | self.pool3 = NonEquivariantMaxPool(64 // 3) |
| 258 | self.pool4 = NonEquivariantMaxPool(2 * feat_dim) |
| 259 | elif pooling == 'mean': |
| 260 | self.pool1 = mean_pool |
| 261 | self.pool2 = mean_pool |
| 262 | self.pool3 = mean_pool |
| 263 | self.pool4 = mean_pool |
| 264 | |
| 265 | self.conv6 = NonEquivariantLinearLeakyReLU(64 // 3 * 3, feat_dim, dim=4) |
| 266 | self.linear0 = nn.Linear(3, 2 * feat_dim) |
| 267 | |
| 268 | def forward(self, x): |
| 269 | |
| 270 | # x: (batch_size, 3, num_points) |
| 271 | # l: (batch_size, 1, 16) |
| 272 | |
| 273 | batch_size = x.size(0) |
| 274 | num_points = x.size(2) |
| 275 | l = x[:, 0, 0:16].reshape(batch_size, 1, 16) |
| 276 | print("!!! the shape of l is: ", l.shape) |
| 277 | |
| 278 | print("!!! the shape of x is: ", x.shape) |
| 279 | x = x.unsqueeze(1) # (32, 1, 3, 1024) |
| 280 | print("!!! the shape of x is: ", x.shape) |
| 281 | |
| 282 | x = get_graph_feature(x, k=self.n_knn) # (32, 2, 3, 1024, 20) |
| 283 | |
| 284 | print("!!! the shape of x is: ", x.shape) |
| 285 | x = self.conv1(x) # (32, 21, 3, 1024, 20) |
| 286 | print("!!! the shape of x is: ", x.shape) |
| 287 | x = self.conv2(x) # (32, 21, 3, 1024, 20) |
| 288 | x1 = self.pool1(x) # (32, 21, 3, 1024) |
| 289 | |
| 290 | x = get_graph_feature(x1, k=self.n_knn) |
| 291 | x = self.conv3(x) |
| 292 | x = self.conv4(x) |
| 293 | x2 = self.pool2(x) |