File size: 1,435 Bytes
b30c1d8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 |
import pytorch_lightning as pl
import torch
import torch.nn as nn
import torch.nn.functional as F
from pointnet2_ops.pointnet2_modules import PointnetSAModule, PointnetSAModuleMSG
from pointnet2.models.pointnet2_ssg_cls import PointNet2ClassificationSSG
class PointNet2ClassificationMSG(PointNet2ClassificationSSG):
def _build_model(self):
super()._build_model()
self.SA_modules = nn.ModuleList()
self.SA_modules.append(
PointnetSAModuleMSG(
npoint=512,
radii=[0.1, 0.2, 0.4],
nsamples=[16, 32, 128],
mlps=[[3, 32, 32, 64], [3, 64, 64, 128], [3, 64, 96, 128]],
use_xyz=self.hparams["model.use_xyz"],
)
)
input_channels = 64 + 128 + 128
self.SA_modules.append(
PointnetSAModuleMSG(
npoint=128,
radii=[0.2, 0.4, 0.8],
nsamples=[32, 64, 128],
mlps=[
[input_channels, 64, 64, 128],
[input_channels, 128, 128, 256],
[input_channels, 128, 128, 256],
],
use_xyz=self.hparams["model.use_xyz"],
)
)
self.SA_modules.append(
PointnetSAModule(
mlp=[128 + 256 + 256, 256, 512, 1024],
use_xyz=self.hparams["model.use_xyz"],
)
)
|