|
- """
- @Author: Dinghao Yang
- @Contact: dinghaoyang@gmail,cin
- @Time: 2020/9/28 7:29 PM
-
- """
-
- import torch
- import argparse
- from ptflops import get_model_complexity_info
- import sys
- import numpy as np
- from model_test import DGCNN_cls, DGCNN_cls_LLE, DGCNN_cls_MP
- import torchsummary
-
- if __name__ == "__main__":
- # Training settings
- parser = argparse.ArgumentParser(description='Point Cloud Recognition')
- parser.add_argument('--exp_name', type=str, default='exp', metavar='N',
- help='Name of the experiment')
- parser.add_argument('--model', type=str, default='dgcnn', metavar='N',
- choices=['pointnet', 'dgcnn'],
- help='Model to use, [pointnet, dgcnn]')
- parser.add_argument('--dataset', type=str, default='modelnet40', metavar='N',
- choices=['modelnet40'])
- parser.add_argument('--batch_size', type=int, default=64, metavar='batch_size',
- help='Size of batch)')
- parser.add_argument('--test_batch_size', type=int, default=16, metavar='batch_size',
- help='Size of batch)')
- parser.add_argument('--epochs', type=int, default=300, metavar='N',
- help='number of episode to train ')
- parser.add_argument('--use_sgd', type=bool, default=True,
- help='Use SGD')
- parser.add_argument('--lr', type=float, default=0.001, metavar='LR',
- help='learning rate (default: 0.001, 0.1 if using sgd)')
- parser.add_argument('--momentum', type=float, default=0.9, metavar='M',
- help='SGD momentum (default: 0.9)')
- parser.add_argument('--scheduler', type=str, default='cos', metavar='N',
- choices=['cos', 'step'],
- help='Scheduler to use, [cos, step]')
- parser.add_argument('--no_cuda', type=bool, default=False,
- help='enables CUDA training')
- parser.add_argument('--seed', type=int, default=1, metavar='S',
- help='random seed (default: 1)')
- parser.add_argument('--eval', type=bool, default=False,
- help='evaluate the model')
- parser.add_argument('--num_points', type=int, default=1024,
- help='num of points to use')
- parser.add_argument('--dropout', type=float, default=0.5,
- help='initial dropout rate')
- parser.add_argument('--emb_dims', type=int, default=2048, metavar='N',
- help='Dimension of embeddings')
- parser.add_argument('--k', type=int, default=20, metavar='N',
- help='Num of nearest neighbors to use')
- parser.add_argument('--model_path', type=str, default='', metavar='N',
- help='Pretrained model path')
- parser.add_argument('--hyper_times', type=int, default=1, metavar='N',
- help='The time of model size')
- args = parser.parse_args()
-
- with torch.cuda.device(0):
- model = DGCNN_cls(args).cuda()
- torchsummary.summary(model, (32, 3, 1024))
- # print('# Parameters:', sum(param.numel() for param in DGCNN_cls.parameters() if param.requires_grad))
- # macs, params = get_model_complexity_info(net, (32, 3, 1024), as_strings=True,
- # print_per_layer_stat=True, verbose=True)
- # print('{:<30} {:<8}'.format('Computational complexity: ', macs))
- # print('{:<30} {:<8}'.format('Number of parameters: ', params))
|