content: import pytorch_lightning as pl
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  import torch.optim.lr_scheduler as lr_sched
  from pointnet2_ops.pointnet2_modules import PointnetFPModule, PointnetSAModule
  from torch.utils.data import DataLoader, DistributedSampler
  from torchvision import transforms
  
  import pointnet2.data.data_utils as d_utils
  from pointnet2.data.ModelNet40Loader import ModelNet40Cls
  
  
  def set_bn_momentum_default(bn_momentum):
      def fn(m):
          if isinstance(m, (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d)):
              m.momentum = bn_momentum
  
      return fn
  
  
  class BNMomentumScheduler(lr_sched.LambdaLR):
      def __init__(self, model, bn_lambda, last_epoch=-1, setter=set_bn_momentum_default):
          if not isinstance(model, nn.Module):
              raise RuntimeError(
                  "Class '{}' is not a PyTorch nn Module".format(type(model)._name_)
              )
  
          self.model = model
          self.setter = setter
          self.lmbd = bn_lambda
  
          self.step(last_epoch + 1)
          self.last_epoch = last_epoch
  
      def step(self, epoch=None):
          if epoch is None:
              epoch = self.last_epoch + 1
  
          self.last_epoch = epoch
          self.model.apply(self.setter(self.lmbd(epoch)))
  
      def state_dict(self):
          return dict(last_epoch=self.last_epoch)
  
      def load_state_dict(self, state):
          self.last_epoch = state["last_epoch"]
          self.step(self.last_epoch)
  
  
  lr_clip = 1e-5
  bnm_clip = 1e-2
  
  
  class PointNet2ClassificationSSG(pl.LightningModule):
      def __init__(self, hparams):
          super().__init__()
  
          self.hparams = hparams
  
          self._build_model()
  
      def _build_model(self):
          self.SA_modules = nn.ModuleList()
          self.SA_modules.append(
              PointnetSAModule(
                  npoint=512,
                  radius=0.2,
                  nsample=64,
                  mlp=[3, 64, 64, 128],
                  use_xyz=self.hparams["model.use_xyz"],
              )
          )
          self.SA_modules.append(
              PointnetSAModule(
                  npoint=128,
                  radius=0.4,
                  nsample=64,
                  mlp=[128, 128, 128, 256],
                  use_xyz=self.hparams["model.use_xyz"],
              )
          )
          self.SA_modules.append(
              PointnetSAModule(
                  mlp=[256, 256, 512, 1024], use_xyz=self.hparams["model.use_xyz"]
              )
          )
  
          self.fc_layer = nn.Sequential(
              nn.Linear(1024, 512, bias=False),
              nn.BatchNorm1d(512),
              nn.ReLU(True),
              nn.Linear(512, 256, bias=False),
              nn.BatchNorm1d(256),
              nn.ReLU(True),
              nn.Dropout(0.5),
              nn.Linear(256, 40),
          )
  
      def _break_up_pc(self, pc):
          xyz = pc[..., 0:3].contiguous()
          features = pc[..., 3:].transpose(1, 2).contiguous() if pc.size(-1) > 3 else None
  
          return xyz, features
  
      def forward(self, pointcloud):
          r"""
              Forward pass of the network
  
              Parameters
              ----------
              pointcloud: Variable(torch.cuda.FloatTensor)
                  (B, N, 3 + input_channels) tensor
                  Point cloud to run predicts on
                  Each point in the point-cloud MUST
                  be formated as (x, y, z, features...)
          """
          xyz, features = self._break_up_pc(pointcloud)
  
          for module in self.SA_modules:
              xyz, features = module(xyz, features)
  
          return self.fc_layer(features.squeeze(-1))
  
      def training_step(self, batch, batch_idx):
          pc, labels = batch
  
          logits = self.forward(pc)
          loss = F.cross_entropy(logits, labels)
          with torch.no_grad():
              acc = (torch.argmax(logits, dim=1) == labels).float().mean()
  
          log = dict(train_loss=loss, train_acc=acc)
  
          return dict(loss=loss, log=log, progress_bar=dict(train_acc=acc))
  
      def validation_step(self, batch, batch_idx):
          pc, labels = batch
  
          logits = self.forward(pc)
          loss = F.cross_entropy(logits, labels)
          acc = (torch.argmax(logits, dim=1) == labels).float().mean()
  
          return dict(val_loss=loss, val_acc=acc)
  
      def validation_end(self, outputs):
          reduced_outputs = {}
          for k in outputs[0]:
              for o in outputs:
                  reduced_outputs[k] = reduced_outputs.get(k, []) + [o[k]]
  
          for k in reduced_outputs:
              reduced_outputs[k] = torch.stack(reduced_outputs[k]).mean()
  
          reduced_outputs.update(
              dict(log=reduced_outputs.copy(), progress_bar=reduced_outputs.copy())
          )
  
          return reduced_outputs
  
      def configure_optimizers(self):
          lr_lbmd = lambda _: max(
              self.hparams["optimizer.lr_decay"]
              ** (
                  int(
                      self.global_step
                      * self.hparams["batch_size"]
                      / self.hparams["optimizer.decay_step"]
                  )
              ),
              lr_clip / self.hparams["optimizer.lr"],
          )
          bn_lbmd = lambda _: max(
              self.hparams["optimizer.bn_momentum"]
              * self.hparams["optimizer.bnm_decay"]
              ** (
                  int(
                      self.global_step
                      * self.hparams["batch_size"]
                      / self.hparams["optimizer.decay_step"]
                  )
              ),
              bnm_clip,
          )
  
          optimizer = torch.optim.Adam(
              self.parameters(),
              lr=self.hparams["optimizer.lr"],
              weight_decay=self.hparams["optimizer.weight_decay"],
          )
          lr_scheduler = lr_sched.LambdaLR(optimizer, lr_lambda=lr_lbmd)
          bnm_scheduler = BNMomentumScheduler(self, bn_lambda=bn_lbmd)
  
          return [optimizer], [lr_scheduler, bnm_scheduler]
  
      def prepare_data(self):
          train_transforms = transforms.Compose(
              [
                  d_utils.PointcloudToTensor(),
                  d_utils.PointcloudScale(),
                  d_utils.PointcloudRotate(),
                  d_utils.PointcloudRotatePerturbation(),
                  d_utils.PointcloudTranslate(),
                  d_utils.PointcloudJitter(),
                  d_utils.PointcloudRandomInputDropout(),
              ]
          )
  
          self.train_dset = ModelNet40Cls(
              self.hparams["num_points"], transforms=train_transforms, train=True
          )
          self.val_dset = ModelNet40Cls(
              self.hparams["num_points"], transforms=None, train=False
          )
  
      def _build_dataloader(self, dset, mode):
          return DataLoader(
              dset,
              batch_size=self.hparams["batch_size"],
              shuffle=mode == "train",
              num_workers=4,
              pin_memory=True,
              drop_last=mode == "train",
          )
  
      def train_dataloader(self):
          return self._build_dataloader(self.train_dset, mode="train")
  
      def val_dataloader(self):
          return self._build_dataloader(self.val_dset, mode="val")
  