| 243 | |
| 244 | class PointNetSetAbstractionMsg(nn.Module): |
| 245 | def __init__(self, npoint, radius_list, nsample_list, in_channel, mlp_list): |
| 246 | super(PointNetSetAbstractionMsg, self).__init__() |
| 247 | self.npoint = npoint |
| 248 | self.radius_list = radius_list |
| 249 | self.nsample_list = nsample_list |
| 250 | self.conv_blocks = nn.ModuleList() |
| 251 | self.bn_blocks = nn.ModuleList() |
| 252 | for i in range(len(mlp_list)): |
| 253 | convs = nn.ModuleList() |
| 254 | bns = nn.ModuleList() |
| 255 | last_channel = in_channel + 3 |
| 256 | for out_channel in mlp_list[i]: |
| 257 | convs.append(nn.Conv2d(last_channel, out_channel, 1)) |
| 258 | bns.append(nn.BatchNorm2d(out_channel)) |
| 259 | last_channel = out_channel |
| 260 | self.conv_blocks.append(convs) |
| 261 | self.bn_blocks.append(bns) |
| 262 | |
| 263 | def forward(self, xyz, points): |
| 264 | """ |