| 226 | |
| 227 | class PointNetSetAbstractionMsg(nn.Module): |
| 228 | def __init__(self, npoint, radius_list, nsample_list, in_channel, mlp_list): |
| 229 | super(PointNetSetAbstractionMsg, self).__init__() |
| 230 | self.npoint = npoint |
| 231 | self.radius_list = radius_list |
| 232 | self.nsample_list = nsample_list |
| 233 | self.conv_blocks = nn.ModuleList() |
| 234 | self.bn_blocks = nn.ModuleList() |
| 235 | for i in range(len(mlp_list)): |
| 236 | convs = nn.ModuleList() |
| 237 | bns = nn.ModuleList() |
| 238 | last_channel = in_channel + 3 |
| 239 | for out_channel in mlp_list[i]: |
| 240 | convs.append(nn.Conv2d(last_channel, out_channel, 1)) |
| 241 | bns.append(nn.BatchNorm2d(out_channel)) |
| 242 | last_channel = out_channel |
| 243 | self.conv_blocks.append(convs) |
| 244 | self.bn_blocks.append(bns) |
| 245 | |
| 246 | def forward(self, xyz, points): |
| 247 | """ |