| 174 | |
| 175 | |
| 176 | class Model(nn.Module): |
| 177 | def __init__(self, in_channels=3, num_classes=40, scale=0.001): |
| 178 | super().__init__() |
| 179 | self.mat_diff_loss_scale = scale |
| 180 | self.backbone = PointNetEncoder(global_feat=True, feature_transform=True, in_channels=in_channels) |
| 181 | self.cls_head = nn.Sequential( |
| 182 | nn.Linear(1024, 512), |
| 183 | nn.BatchNorm1d(512), |
| 184 | nn.ReLU(inplace=True), |
| 185 | nn.Linear(512, 256), |
| 186 | nn.Dropout(p=0.4), |
| 187 | nn.BatchNorm1d(256), |
| 188 | nn.ReLU(inplace=True), |
| 189 | nn.Linear(256, num_classes) |
| 190 | ) |
| 191 | |
| 192 | def forward(self, x, gts): |
| 193 | x, trans, trans_feat = self.backbone(x) |
| 194 | x = self.cls_head(x) |
| 195 | x = F.log_softmax(x, dim=1) |
| 196 | loss = F.nll_loss(x, gts) |
| 197 | mat_diff_loss = feature_transform_reguliarzer(trans_feat) |
| 198 | total_loss = loss + mat_diff_loss * self.mat_diff_loss_scale |
| 199 | return total_loss, x |
| 200 | |
| 201 | |
| 202 | """ |