'--save_model', type=str, default='./save_model/' ) parser.add_argument( '--base_model', type=str, default='pointnet' ) cfg = parser.parse_args() print(cfg) model_dict = { 'pointnet': (PointNet(k=40), F.nll_loss, accuracy, torch.optim.Adam, PointNetTrainer), } train_transforms = transforms.Compose([ data_utils.PointcloudToTensor(), data_utils.PointcloudRotatebyAngle(np.pi/4), data_utils.PointcloudJitter(), data_utils.PointcloudScaleAndTranslate(), data_utils.PointcloudScale(), data_utils.PointcloudTranslate(), data_utils.PointcloudRandomInputDropout() ]) test_transforms = transforms.Compose([ data_utils.PointcloudToTensor() ]) device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') print(device)
parser.add_argument('--dataset', type=str, default='data/') parser.add_argument('--workers', type=int, default=4) parser.add_argument('--save_model', type=str, default='./save_model/') parser.add_argument('--base_model', type=str, default='pointnet') cfg = parser.parse_args() print(cfg) model_dict = { 'pointnet': (PointNet(k=40), F.nll_loss, accuracy, torch.optim.Adam, PointNetTrainer) } train_transforms = transforms.Compose([ data_utils.PointcloudToTensor(), data_utils.PointcloudRotatebyAngle(np.pi / 4), data_utils.PointcloudJitter(), data_utils.PointcloudScaleAndTranslate(), data_utils.PointcloudScale(), data_utils.PointcloudTranslate(), data_utils.PointcloudRandomInputDropout() ]) test_transforms = transforms.Compose([data_utils.PointcloudToTensor()]) device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') print(device) if __name__ == "__main__": ds_train = ModelNet40Dataset(num_points=2500,