Initial commit

2026-06-27 16:00:07 +08:00 · 2017-12-26 18:43:17 -05:00
commit dc4e2b0db3
42 changed files with 3486 additions and 0 deletions
@@ -0,0 +1,143 @@
+import torch
+import torch.optim as optim
+import torch.optim.lr_scheduler as lr_sched
+import torch.nn as nn
+import numpy as np
+from torch.utils.data import DataLoader
+from torch.autograd import Variable
+from torchvision import transforms
+import os
+import tensorboard_logger as tb_log
+
+from models import PointnetCls as Pointnet
+from models.PointnetCls import model_fn_decorator
+from data import ModelNet40Cls
+import utils.pytorch_utils as pt_utils
+import utils.data_utils as d_utils
+import argparse
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Arg parser")
+    parser.add_argument(
+        "-batch_size", type=int, default=128, help="Batch size [default: 128]")
+    parser.add_argument(
+        "-num_points",
+        type=int,
+        default=1024,
+        help="Number of points to train with [default: 1024]")
+    parser.add_argument(
+        "-weight_decay", type=float, default=1e-5, help="L2 regularization coeff")
+    parser.add_argument(
+        "-lr",
+        type=float,
+        default=1e-2,
+        help="Initial learning rate [default: 1e-2]")
+    parser.add_argument(
+        "-lr_decay",
+        type=float,
+        default=0.7,
+        help="Learning rate decay gamma [default: 0.7]")
+    parser.add_argument(
+        "-decay_step",
+        type=int,
+        default=20,
+        help="Learning rate decay step [default: 20]")
+    parser.add_argument(
+        "-bn_momentum",
+        type=float,
+        default=0.5,
+        help="Initial batch norm momentum [default: 0.5]")
+    parser.add_argument(
+        "-bnm_decay",
+        type=float,
+        default=0.5,
+        help="Batch norm momentum decay gamma [default: 0.5]")
+    parser.add_argument(
+        "-checkpoint", type=str, default=None, help="Checkpoint to start from")
+    parser.add_argument(
+        "-epochs", type=int, default=200, help="Number of epochs to train for")
+    parser.add_argument(
+        "-run_name",
+        type=str,
+        default="cls_run_1",
+        help="Name for run in tensorboard_logger")
+
+    return parser.parse_args()
+
+lr_clip = 1e-5
+bnm_clip = 1e-2
+
+if __name__ == "__main__":
+    args = parse_args()
+
+    BASE_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'data')
+
+    transforms = transforms.Compose([
+        d_utils.PointcloudToTensor(),
+        d_utils.PointcloudRotate(x_axis=True),
+        d_utils.PointcloudScale(),
+        d_utils.PointcloudTranslate(),
+        d_utils.PointcloudJitter()
+    ])
+
+    test_set = ModelNet40Cls(
+        args.num_points, BASE_DIR, transforms=transforms, train=False)
+    test_loader = DataLoader(
+        test_set,
+        batch_size=args.batch_size,
+        shuffle=True,
+        num_workers=2,
+        pin_memory=True)
+
+    train_set = ModelNet40Cls(args.num_points, BASE_DIR, transforms=transforms)
+    train_loader = DataLoader(
+        train_set,
+        batch_size=args.batch_size,
+        shuffle=True,
+        num_workers=2,
+        pin_memory=True)
+
+    tb_log.configure('runs/{}'.format(args.run_name))
+
+    model = Pointnet()
+    model.cuda()
+    optimizer = optim.Adam(
+        model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
+    lr_lbmd = lambda e: max(args.lr_decay**(e // args.decay_step), lr_clip / args.lr)
+    bn_lbmd = lambda e: max(args.bn_momentum * args.bnm_decay**(e // args.decay_step), bnm_clip)
+
+    if args.checkpoint is not None:
+        start_epoch, best_prec = pt_utils.load_checkpoint(
+            model, optimizer, filename=args.checkpoint.split(".")[0])
+
+        lr_scheduler = lr_sched.LambdaLR(
+            optimizer, lr_lambda=lr_lbmd, last_epoch=start_epoch)
+        bnm_scheduler = pt_utils.BNMomentumScheduler(
+            model, bn_lambda=bn_lbmd, last_epoch=start_epoch)
+    else:
+        lr_scheduler = lr_sched.LambdaLR(optimizer, lr_lambda=lr_lbmd)
+        bnm_scheduler = pt_utils.BNMomentumScheduler(model, bn_lambda=bn_lbmd)
+
+        best_prec = 0.0
+        start_epoch = 1
+
+    model_fn = model_fn_decorator(nn.CrossEntropyLoss())
+
+    trainer = pt_utils.Trainer(
+        model,
+        model_fn,
+        optimizer,
+        checkpoint_name="cls_checkpoint",
+        best_name="cls_best",
+        lr_scheduler=lr_scheduler,
+        bnm_scheduler=bnm_scheduler)
+
+    trainer.train(
+        start_epoch,
+        args.epochs,
+        train_loader,
+        test_loader,
+        best_prec=best_prec)
+
+    if start_epoch == args.epochs:
+        _ = trainer.eval_epoch(start_epoch, test_loader)