def create_datasets_and_dataloaders(num_points=1024): assert (os.path.isdir(args.in_data_dir)) #train_data = Dataset(path=args.in_data_dir, task='train',\ # num_samples=num_points, generalize=args.generalize, robustness=args.robustness) #test_data = Dataset(path=args.in_data_dir, task='test',\ # num_samples=num_points, generalize=args.generalize, robustness=args.robustness) train_data = ModelNet40(path=args.in_data_dir, task='train',\ num_samples=num_points, generalize=args.generalize, robustness=args.robustness) test_data = ModelNet40(path=args.in_data_dir, task='test',\ num_samples=num_points, generalize=args.generalize, robustness=args.robustness) print('# points: {:d}'.format(num_points)) train_dataloader = torch.utils.data.DataLoader(train_data, batch_size=args.batch_size, shuffle=args.train, num_workers=int( args.n_workers)) test_dataloader = torch.utils.data.DataLoader(test_data, batch_size=args.batch_size, shuffle=args.train, num_workers=int( args.n_workers)) return train_data, train_dataloader, test_data, test_dataloader
def main(args): torch.backends.cudnn.deterministic = True torch.manual_seed(args.seed) torch.cuda.manual_seed_all(args.seed) np.random.seed(args.seed) boardio = SummaryWriter(log_dir='checkpoints/' + args.exp_name) _init_(args) textio = IOStream('checkpoints/' + args.exp_name + '/run.log') textio.cprint(str(args)) if args.dataset == 'modelnet40': train_loader = DataLoader(ModelNet40( num_points=args.num_points, partition='train', gaussian_noise=args.gaussian_noise, unseen=args.unseen, factor=args.factor), batch_size=args.batch_size, shuffle=True, drop_last=True) test_loader = DataLoader(ModelNet40(num_points=args.num_points, partition='test', gaussian_noise=args.gaussian_noise, unseen=args.unseen, factor=args.factor), batch_size=args.test_batch_size, shuffle=False, drop_last=False) else: raise Exception("not implemented") if args.model == 'dcp': net = DCP(args).cuda() if args.eval: if args.model_path is '': model_path = 'checkpoints' + '/' + args.exp_name + '/models/model.best.t7' else: model_path = args.model_path print(model_path) if not os.path.exists(model_path): print("can't find pretrained model") return net.load_state_dict(torch.load(model_path), strict=False) if torch.cuda.device_count() > 1: net = nn.DataParallel(net) print("Let's use", torch.cuda.device_count(), "GPUs!") else: raise Exception('Not implemented') if args.eval: test(args, net, test_loader, boardio, textio) else: train(args, net, train_loader, test_loader, boardio, textio) print('FINISH') boardio.close()
def main(): parser = hmnet_arguments() args = parser.parse_args() # 保证实验的可重复性 torch.backends.cudnn.deterministic = True torch.manual_seed(args.seed) torch.cuda.manual_seed_all(args.seed) np.random.seed(args.seed) _init_(args) if args.dataset == 'ours': train_loader = DataLoader(ModelNet40( num_points=args.n_points, num_subsampled_points=args.n_subsampled_points, partition='train', gaussian_noise=args.gaussian_noise, rot_factor=args.rot_factor, overlap=args.overlap), batch_size=args.train_batch_size, shuffle=True, drop_last=True, num_workers=args.num_workers) test_loader = DataLoader(ModelNet40( num_points=args.n_points, num_subsampled_points=args.n_subsampled_points, partition='test', gaussian_noise=args.gaussian_noise, rot_factor=args.rot_factor, overlap=args.overlap), batch_size=args.test_batch_size, shuffle=False, drop_last=False, num_workers=args.num_workers) else: raise Exception("not implemented") if args.model == 'hmnet': if args.eval: if args.model_path is '': model_path = 'checkpoints' + '/' + args.exp_name + '/models/model.best.t7' else: model_path = args.model_path if not os.path.exists(model_path): print("can't find pretrained model") return model_2 = MatchNet_2(args) net = HMNet(args, model_2) eval_model(net, test_loader) else: train(args, train_loader, test_loader) else: raise Exception('Not implemented') print('FINISH')
def test(args, io): test_loader = DataLoader(ModelNet40(partition='test', num_points=args.num_points), batch_size=args.test_batch_size, shuffle=True, drop_last=False) device = torch.device("cuda" if args.cuda else "cpu") model = Pct(args).to(device) model = nn.DataParallel(model) model.load_state_dict(torch.load(args.model_path)) model = model.eval() test_true = [] test_pred = [] for data, label in test_loader: data, label = data.to(device), label.to(device).squeeze() data = data.permute(0, 2, 1) logits = model(data) preds = logits.max(dim=1)[1] test_true.append(label.cpu().numpy()) test_pred.append(preds.detach().cpu().numpy()) test_true = np.concatenate(test_true) test_pred = np.concatenate(test_pred) test_acc = metrics.accuracy_score(test_true, test_pred) avg_per_class_acc = metrics.balanced_accuracy_score(test_true, test_pred) outstr = 'Test :: test acc: %.6f, test avg acc: %.6f'%(test_acc, avg_per_class_acc) io.cprint(outstr)
def main(): args = parse_args() # print(f"args: {args}") os.environ["HDF5_USE_FILE_LOCKING"] = "FALSE" print("===> building and loading models ...") net = plot21H() # print(net) checkpoint = torch.load(args.checkpoint, map_location=torch.device('cpu')) new_check_point = OrderedDict() for k, v in checkpoint['net'].items(): name = k[7:] # remove `module.` new_check_point[name] = v net.load_state_dict(new_check_point) net.eval() print('==> Preparing data ...') # train_set =ModelNet40(partition='train', num_points=args.num_points) test_set = ModelNet40(partition='test', num_points=args.num_points) data, label = test_set.__getitem__(args.id) plot_xyz(data, args, None, name=f"attention/Image-{args.id}.pdf" ) data = torch.tensor(data).unsqueeze(dim=0) data = data.permute(0, 2, 1) with torch.no_grad(): logits, structure_list = net(data) preds = logits.max(dim=1)[1] print(f"predict: {preds} | label: {label}") struct = structure_list[args.stage] plot_struct(struct, args)
def test(args, io): if args.dataset == 'modelnet40': test_loader = DataLoader(ModelNet40(partition='test', num_points=args.num_points), num_workers=8, batch_size=args.test_batch_size, shuffle=True, drop_last=False) elif args.dataset == 'ScanObjectNN': test_loader = DataLoader(ScanObjectNN(partition='test', num_points=args.num_points), num_workers=8, batch_size=args.test_batch_size, shuffle=True, drop_last=False) else: raise Exception("Dataset Not supported") device = torch.device("cuda" if args.cuda else "cpu") #Try to load models if args.model == 'pointnet': if args.dataset == 'modelnet40': model = PointNet(args, output_channels=40).to(device) elif args.dataset == 'ScanObjectNN': model = PointNet(args, output_channels=15).to(device) else: raise Exception("Dataset Not supported") elif args.model == 'dgcnn': if args.dataset == 'modelnet40': model = DGCNN(args, output_channels=40).to(device) elif args.dataset == 'ScanObjectNN': model = DGCNN(args, output_channels=15).to(device) else: raise Exception("Dataset Not supported") elif args.model == 'gbnet': if args.dataset == 'modelnet40': model = GBNet(args, output_channels=40).to(device) elif args.dataset == 'ScanObjectNN': model = GBNet(args, output_channels=15).to(device) else: raise Exception("Dataset Not supported") else: raise Exception("Not implemented") print(str(model)) model = nn.DataParallel(model) model.load_state_dict(torch.load(args.model_path)) model = model.eval() test_acc = 0.0 count = 0.0 test_true = [] test_pred = [] for data, label in test_loader: data, label = data.to(device), label.to(device).squeeze() data = data.permute(0, 2, 1) batch_size = data.size()[0] logits = model(data) preds = logits.max(dim=1)[1] test_true.append(label.cpu().numpy()) test_pred.append(preds.detach().cpu().numpy()) test_true = np.concatenate(test_true) test_pred = np.concatenate(test_pred) test_acc = metrics.accuracy_score(test_true, test_pred) avg_per_class_acc = metrics.balanced_accuracy_score(test_true, test_pred) outstr = 'Test :: test acc: %.6f, test avg acc: %.6f'%(test_acc, avg_per_class_acc) io.cprint(outstr)
def test(args, io): test_loader = DataLoader(ModelNet40(partition='test', num_points=args.num_points), batch_size=args.test_batch_size, shuffle=True, drop_last=False) device = torch.device("cuda" if args.cuda else "cpu") #Try to load models if args.model == 'pointnet': model = PointNet(args).to(device) elif args.model == 'dgcnn': model = DGCNN(args).to(device) else: model = PaiNet(args).to(device) #raise Exception("Not implemented") print(str(model)) # model = DGCNN(args).to(device) model = nn.DataParallel(model) if os.path.exists('checkpoints/%s/models/model_%s.t7' % (args.exp_name, args.model)): checkpoint_dict = torch.load('./checkpoints/%s/models/model_%s.t7' % (args.exp_name, args.model), map_location=device) # pretrained_dict = {} # for k, v in checkpoint_dict.items(): # if 'transform' in k: # k = k.replace('transform', 'paiIdxMatrix') # pretrained_dict[k]=v # # pretrained_dict = {k: v for k, v in pretrained_dict.items() if 'transform' in k} model.load_state_dict(checkpoint_dict, strict=True) # torch.save(model.state_dict(), 'checkpoints/%s/models/model_%s_2048.t7' % (args.exp_name, args.model)) print("Load model from './checkpoints/%s/models/model_%s_2048.t7 !'" % (args.exp_name, args.model)) model = model.eval() test_acc = 0.0 count = 0.0 test_true = [] test_pred = [] for data, label in test_loader: data, label = data.to(device), label.to(device).squeeze() data = data.permute(0, 2, 1) batch_size = data.size()[0] logits = model(data) preds = logits.max(dim=1)[1] test_true.append(label.cpu().numpy()) test_pred.append(preds.detach().cpu().numpy()) test_true = np.concatenate(test_true) test_pred = np.concatenate(test_pred) test_acc = metrics.accuracy_score(test_true, test_pred) avg_per_class_acc = metrics.balanced_accuracy_score(test_true, test_pred) outstr = 'Test :: test acc: %.6f, test avg acc: %.6f' % (test_acc, avg_per_class_acc) io.cprint(outstr)
def test(args, io): test_loader = DataLoader(ModelNet40(partition='test', num_points=args.num_points), batch_size=args.test_batch_size, shuffle=False, drop_last=False) device = torch.device("cuda" if args.cuda else "cpu") #Try to load models model = HDGN(args).to(device) model.load_state_dict(torch.load(args.model_path)) model = model.eval() test_acc = 0.0 count = 0.0 test_true = [] test_pred = [] ind_acc = [] ind_per_class_acc = [] for i in range(20): test_pred.append([]) test_true.append([]) for data, label in test_loader: data, label = data.to(device), label.to(device).squeeze() data = data.permute(0, 2, 1) batch_size = data.size()[0] logits = model(data) test_pred[i].append(logits.detach().cpu().numpy()) test_true[i].append(label.cpu().numpy()) test_pred[i] = np.concatenate(test_pred[i]) ind_acc.append( metrics.accuracy_score(np.concatenate(test_true[0]), np.argmax(test_pred[i], 1))) ind_per_class_acc.append( metrics.balanced_accuracy_score(np.concatenate(test_true[0]), np.argmax(test_pred[i], 1))) test_true = np.concatenate(test_true[0]) test_pred = np.stack(test_pred) test_pred_mean = np.mean(test_pred, axis=0) test_pred_mean = np.argmax(test_pred_mean, axis=1) test_acc = metrics.accuracy_score(test_true, test_pred_mean) avg_per_class_acc = metrics.balanced_accuracy_score( test_true, test_pred_mean) #if dist.get_rank() == 0: outstr = 'Test :: test acc: %.6f, test avg acc: %.6f' % (test_acc, avg_per_class_acc) io.cprint(outstr) avg_ind_acc = np.mean(ind_acc) avg_ind_per_class_acc = np.mean(ind_per_class_acc) std_ind_acc = np.std(ind_acc) std_ind_per_class_acc = np.std(ind_per_class_acc) outstr = 'Test :: ind test acc: %.6f std %.6f, ind test avg acc: %.6f std %.6f' % ( avg_ind_acc, std_ind_acc, avg_ind_per_class_acc, std_ind_per_class_acc) io.cprint(outstr)
def test(args, io): test_loader = DataLoader(ModelNet40(partition='test', num_points=args.num_points), batch_size=args.test_batch_size, shuffle=True, drop_last=False) device = torch.device("cuda" if args.cuda else "cpu") #Try to load models if args.model == 'pointnet': model = PointNet(args).to(device) elif args.model == 'dgcnn': model = DGCNN(args).to(device) else: raise Exception("Not implemented") print(str(model)) model = nn.DataParallel(model) checkpoint = torch.load(args.resume) model.load_state_dict(checkpoint['state_dict']) model = model.eval() test_acc = 0.0 count = 0.0 test_true = [] test_pred = [] SHAPE_NAMES = [line.rstrip() for line in \ open('data/modelnet40_ply_hdf5_2048/shape_names.txt')] NUM_CLASSES = 40 total_seen_class = [0 for _ in range(NUM_CLASSES)] total_correct_class = [0 for _ in range(NUM_CLASSES)] for data, label in test_loader: data, label = data.to(device), label.to(device).squeeze() data = data.permute(0, 2, 1) batch_size = data.size()[0] logits = model(data) preds = logits.max(dim=1)[1] test_true.append(label.cpu().numpy()) test_pred.append(preds.detach().cpu().numpy()) test_true = np.concatenate(test_true) test_pred = np.concatenate(test_pred) test_acc = metrics.accuracy_score(test_true, test_pred) avg_per_class_acc = metrics.balanced_accuracy_score(test_true, test_pred) outstr = 'Test :: test acc: %.6f, test avg acc: %.6f' % (test_acc, avg_per_class_acc) io.cprint(outstr) for i in range(test_true.shape[0]): l = test_true[i] total_seen_class[l] += 1 total_correct_class[l] += (test_pred[i] == l) class_accuracies = np.array(total_correct_class) / np.array( total_seen_class, dtype=np.float) for i, name in enumerate(SHAPE_NAMES): io.cprint('%10s:\t%0.3f' % (name, class_accuracies[i]))
def test(args, io): test_loader = DataLoader(ModelNet40(partition='test', num_points=args.num_points), num_workers=5, batch_size=args.test_batch_size, shuffle=False, drop_last=False) device = torch.device("cuda" if args.cuda else "cpu") NUM_PEPEAT = 300 NUM_VOTE = 10 # Try to load models model = GDANET().to(device) model = nn.DataParallel(model) model.load_state_dict(torch.load(args.model_path)) model = model.eval() best_acc = 0 pointscale = PointcloudScale(scale_low=2. / 3., scale_high=3. / 2., trans_low=-0.2, trans_high=0.2, trans_open=True) for i in range(NUM_PEPEAT): test_true = [] test_pred = [] for data, label in test_loader: data, label = data.to(device), label.to(device).squeeze() pred = 0 for v in range(NUM_VOTE): new_data = data batch_size = data.size()[0] if v > 0: new_data.data = pointscale(new_data.data) with torch.no_grad(): pred += F.softmax(model(new_data.permute(0, 2, 1)), dim=1) pred /= NUM_VOTE label = label.view(-1) pred_choice = pred.max(dim=1)[1] test_true.append(label.cpu().numpy()) test_pred.append(pred_choice.detach().cpu().numpy()) test_true = np.concatenate(test_true) test_pred = np.concatenate(test_pred) test_acc = metrics.accuracy_score(test_true, test_pred) if test_acc > best_acc: best_acc = test_acc outstr = 'Voting %d, test acc: %.6f,' % (i, test_acc * 100) io.cprint(outstr) final_outstr = 'Final voting result test acc: %.6f,' % (best_acc * 100) io.cprint(final_outstr)
def main(): args = parse_args() # print(f"args: {args}") os.environ["HDF5_USE_FILE_LOCKING"] = "FALSE" print('==> Preparing data ...') # train_set =ModelNet40(partition='train', num_points=args.num_points) test_set = ModelNet40(partition='test', num_points=args.num_points) data, label = test_set.__getitem__(args.id) plot_xyz(data, args, name=f"forstructure/Image-{args.id}-{args.num_points}.pdf")
def test(args, io): test_loader = DataLoader(ModelNet40(partition='test', num_points=args.num_points), batch_size=args.test_batch_size, shuffle=True, drop_last=False) device = torch.device("cuda" if args.cuda else "cpu") #Try to load models model = DGCNN(args).to(device) model = nn.DataParallel(model) model.load_state_dict(torch.load(args.model_path)) model = model.eval() test_acc = 0.0 count = 0.0 test_true = [] test_pred = [] # DUMP_DIR = 'checkpoints/' + args.exp_name + '/' + 'dump' fout = open(os.path.join(DUMP_DIR, 'pred_label.txt'), 'w') # for data, label in test_loader: data, label = data.to(device), label.to(device).squeeze() data = data.permute(0, 2, 1) batch_size = data.size()[0] logits = model(data) preds = logits.max(dim=1)[1] test_true.append(label.cpu().numpy()) test_pred.append(preds.detach().cpu().numpy()) test_true = np.concatenate(test_true) test_pred = np.concatenate(test_pred) # for i in range(list(test_true.shape)[0]): fout.write('%d, %d\n' % (test_pred[i], test_true[i])) # test_acc = metrics.accuracy_score(test_true, test_pred) avg_per_class_acc = metrics.balanced_accuracy_score(test_true, test_pred) outstr = 'Test :: test acc: %.6f, test avg acc: %.6f' % (test_acc, avg_per_class_acc) io.cprint(outstr)
def test(args, io): test_loader = DataLoader(ModelNet40(partition='test', num_points=args.num_points), batch_size=args.test_batch_size, shuffle=True, drop_last=False) device = torch.device("cuda" if args.cuda else "cpu") #Try to load models if args.model == 'pointnet': model = PointNet(args).to(device) elif args.model == 'dgcnn': model = DGCNN_cls(args).to(device) else: raise Exception("Not implemented") model = nn.DataParallel(model) model.load_state_dict(torch.load(args.model_path)) model = model.eval() test_acc = 0.0 count = 0.0 test_true = [] test_pred = [] for data, label in test_loader: data, label = data.to(device), label.to(device).squeeze() data = data.permute(0, 2, 1) batch_size = data.size()[0] logits = model(data) preds = logits.max(dim=1)[1] test_true.append(label.cpu().numpy()) test_pred.append(preds.detach().cpu().numpy()) # visualize - added by jaeha if args.visualize: xyz = data[0].cpu() ax = plt.axes(projection='3d') ax.scatter(xyz[0, :], xyz[1, :], xyz[2, :], s=1, color='blue') plt.title('True: ' + class_lists[label[0]] + ' , Pred: ' + class_lists[preds[0]]) plt.show() test_true = np.concatenate(test_true) test_pred = np.concatenate(test_pred) test_acc = metrics.accuracy_score(test_true, test_pred) avg_per_class_acc = metrics.balanced_accuracy_score(test_true, test_pred) outstr = 'Test :: test acc: %.6f, test avg acc: %.6f' % (test_acc, avg_per_class_acc) io.cprint(outstr)
def main(): args = parse_args() print(f"args: {args}") os.environ["HDF5_USE_FILE_LOCKING"] = "FALSE" torch.manual_seed(args.seed) if torch.cuda.is_available(): device = 'cuda' torch.cuda.manual_seed(args.seed) else: device = 'cpu' print(f"==> Using device: {device}") if args.msg is None: message = str(datetime.datetime.now().strftime('-%Y%m%d%H%M%S')) else: message = "-" + args.msg args.checkpoint = 'checkpoints/' + args.model + message print('==> Preparing data..') test_loader = DataLoader(ModelNet40(partition='test', num_points=args.num_points), num_workers=8, batch_size=args.batch_size, shuffle=True, drop_last=False) # Model print('==> Building model..') net = models.__dict__[args.model]() criterion = cal_loss net = net.to(device) checkpoint_path = os.path.join(args.checkpoint, 'best_checkpoint.pth') checkpoint = torch.load(checkpoint_path) # criterion = criterion.to(device) if device == 'cuda': net = torch.nn.DataParallel(net) cudnn.benchmark = True net.load_state_dict(checkpoint['net']) test_out = validate(net, test_loader, criterion, device) print(f"Vanilla out: {test_out}") print(f"===> start voting evaluation...") voting(net, test_loader, device, args)
def test(args, io): test_loader = DataLoader(ModelNet40(partition='test', num_points=args.num_points), batch_size=args.test_batch_size, shuffle=True, drop_last=False) device = torch.device("cuda" if args.cuda else "cpu") #Try to load models if args.model == 'pointnet': model = PointNet(args).to(device) elif args.model == 'dgcnn': model = DGCNN_cls(args).to(device) else: raise Exception("Not implemented") model = nn.DataParallel(model) model.load_state_dict(torch.load(args.model_path)) # model.load_state_dict(torch.load("/home/mask/xas_ws/dgcnn_pytorch/checkpoints/cls_1024/models/model.t7")) model = model.eval() test_acc = 0.0 count = 0.0 test_true = [] test_pred = [] for data, label in test_loader: data, label = data.to(device), label.to(device).squeeze() data = data.permute(0, 2, 1) batch_size = data.size()[0] logits = model(data) preds = logits.max(dim=1)[1] test_true.append(label.cpu().numpy()) test_pred.append(preds.detach().cpu().numpy()) test_true = np.concatenate(test_true) test_pred = np.concatenate(test_pred) test_acc = metrics.accuracy_score(test_true, test_pred) avg_per_class_acc = metrics.balanced_accuracy_score(test_true, test_pred) outstr = 'Test :: test acc: %.6f, test avg acc: %.6f'%(test_acc, avg_per_class_acc) io.cprint(outstr)
def main(): parser = argparse.ArgumentParser(description='Point Cloud Registration') parser.add_argument('--exp_name', type=str, default='exp', metavar='N', help='Name of the experiment') parser.add_argument('--model', type=str, default='prnet', metavar='N', choices=['prnet'], help='Model to use, [prnet]') parser.add_argument('--emb_nn', type=str, default='dgcnn', metavar='N', choices=['pointnet', 'dgcnn'], help='Embedding to use, [pointnet, dgcnn]') parser.add_argument('--attention', type=str, default='transformer', metavar='N', choices=['identity', 'transformer'], help='Head to use, [identity, transformer]') parser.add_argument('--head', type=str, default='svd', metavar='N', choices=['mlp', 'svd'], help='Head to use, [mlp, svd]') parser.add_argument('--n_emb_dims', type=int, default=512, metavar='N', help='Dimension of embeddings') parser.add_argument('--n_blocks', type=int, default=1, metavar='N', help='Num of blocks of encoder&decoder') parser.add_argument('--n_heads', type=int, default=4, metavar='N', help='Num of heads in multiheadedattention') parser.add_argument('--n_iters', type=int, default=3, metavar='N', help='Num of iters to run inference') parser.add_argument('--discount_factor', type=float, default=0.9, metavar='N', help='Discount factor to compute the loss') parser.add_argument('--n_ff_dims', type=int, default=1024, metavar='N', help='Num of dimensions of fc in transformer') parser.add_argument('--n_keypoints', type=int, default=512, metavar='N', help='Num of keypoints to use') parser.add_argument('--temp_factor', type=float, default=100, metavar='N', help='Factor to control the softmax precision') parser.add_argument( '--cat_sampler', type=str, default='gumbel_softmax', choices=['softmax', 'gumbel_softmax'], metavar='N', help='use gumbel_softmax to get the categorical sample') parser.add_argument('--dropout', type=float, default=0.0, metavar='N', help='Dropout ratio in transformer') parser.add_argument('--batch_size', type=int, default=3, metavar='batch_size', help='Size of batch)') parser.add_argument('--test_batch_size', type=int, default=12, metavar='batch_size', help='Size of batch)') parser.add_argument('--epochs', type=int, default=100, metavar='N', help='number of episode to train ') parser.add_argument('--use_sgd', type=bool, default=False, 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('--no_cuda', action='store_true', default=False, help='enables CUDA training') parser.add_argument('--seed', type=int, default=1234, metavar='S', help='random seed (default: 1)') parser.add_argument('--eval', action='store_true', default=False, help='evaluate the model') parser.add_argument('--cycle_consistency_loss', type=float, default=0.1, metavar='N', help='cycle consistency loss') parser.add_argument('--feature_alignment_loss', type=float, default=0.1, metavar='N', help='feature alignment loss') parser.add_argument('--gaussian_noise', type=bool, default=False, metavar='N', help='Wheter to add gaussian noise') parser.add_argument('--unseen', type=bool, default=False, metavar='N', help='Wheter to test on unseen category') parser.add_argument('--n_points', type=int, default=1024, metavar='N', help='Num of points to use') parser.add_argument('--n_subsampled_points', type=int, default=768, metavar='N', help='Num of subsampled points to use') parser.add_argument('--dataset', type=str, default='modelnet40', choices=['modelnet40'], metavar='N', help='dataset to use') parser.add_argument('--rot_factor', type=float, default=4, metavar='N', help='Divided factor of rotation') parser.add_argument('--model_path', type=str, default='', metavar='N', help='Pretrained model path') args = parser.parse_args() torch.backends.cudnn.deterministic = True torch.manual_seed(args.seed) torch.cuda.manual_seed_all(args.seed) np.random.seed(args.seed) _init_(args) if args.dataset == 'modelnet40': train_loader = DataLoader(ModelNet40( num_points=args.n_points, num_subsampled_points=args.n_subsampled_points, partition='train', gaussian_noise=args.gaussian_noise, unseen=args.unseen, rot_factor=args.rot_factor), batch_size=args.batch_size, shuffle=True, drop_last=True, num_workers=6) test_loader = DataLoader(ModelNet40( num_points=args.n_points, num_subsampled_points=args.n_subsampled_points, partition='test', gaussian_noise=args.gaussian_noise, unseen=args.unseen, rot_factor=args.rot_factor), batch_size=args.test_batch_size, shuffle=False, drop_last=False, num_workers=6) else: raise Exception("not implemented") if args.model == 'prnet': net = PRNet(args).cuda() if args.eval: if args.model_path == '': model_path = 'checkpoints' + '/' + args.exp_name + '/models/model.best.t7' else: model_path = args.model_path if not os.path.exists(model_path): print("can't find pretrained model") return else: raise Exception('Not implemented') if not args.eval: train(args, net, train_loader, test_loader) print('FINISH')
def main(): args = config_params() print(args) setup_seed(args.seed) if not os.path.exists(args.saved_path): os.makedirs(args.saved_path) summary_path = os.path.join(args.saved_path, 'summary') if not os.path.exists(summary_path): os.makedirs(summary_path) checkpoints_path = os.path.join(args.saved_path, 'checkpoints') if not os.path.exists(checkpoints_path): os.makedirs(checkpoints_path) train_set = ModelNet40(root=args.root, npts=args.train_npts, train=True, normal=args.normal, mode=args.mode) test_set = ModelNet40(root=args.root, npts=args.train_npts, train=False, normal=args.normal, mode=args.mode) train_loader = DataLoader(train_set, batch_size=args.batchsize, shuffle=True, num_workers=args.num_workers) test_loader = DataLoader(test_set, batch_size=args.batchsize, shuffle=False, num_workers=args.num_workers) in_dim = 6 if args.normal else 3 model = IterativeBenchmark(in_dim=in_dim, niters=args.niters, gn=args.gn) model = model.cuda() loss_fn = EMDLosspy().cuda() optimizer = torch.optim.Adam(model.parameters(), lr=args.lr) scheduler = torch.optim.lr_scheduler.MultiStepLR( optimizer, milestones=args.milestones, gamma=args.gamma, last_epoch=-1) writer = SummaryWriter(summary_path) test_min_loss, test_min_r_mse_error, test_min_rot_error = \ float('inf'), float('inf'), float('inf') for epoch in range(args.epoches): print('=' * 20, epoch + 1, '=' * 20) train_results = train_one_epoch(train_loader, model, loss_fn, optimizer) print_train_info(train_results) test_results = test_one_epoch(test_loader, model, loss_fn) print_train_info(test_results) if epoch % args.saved_frequency == 0: writer.add_scalar('Loss/train', train_results['loss'], epoch + 1) writer.add_scalar('Loss/test', test_results['loss'], epoch + 1) writer.add_scalar('RError/train', train_results['r_mse'], epoch + 1) writer.add_scalar('RError/test', test_results['r_mse'], epoch + 1) writer.add_scalar('rotError/train', train_results['r_isotropic'], epoch + 1) writer.add_scalar('rotError/test', test_results['r_isotropic'], epoch + 1) writer.add_scalar('Lr', optimizer.param_groups[0]['lr'], epoch + 1) test_loss, test_r_error, test_rot_error = \ test_results['loss'], test_results['r_mse'], test_results['r_isotropic'] if test_loss < test_min_loss: saved_path = os.path.join(checkpoints_path, "test_min_loss.pth") torch.save(model.state_dict(), saved_path) test_min_loss = test_loss if test_r_error < test_min_r_mse_error: saved_path = os.path.join(checkpoints_path, "test_min_rmse_error.pth") torch.save(model.state_dict(), saved_path) test_min_r_mse_error = test_r_error if test_rot_error < test_min_rot_error: saved_path = os.path.join(checkpoints_path, "test_min_rot_error.pth") torch.save(model.state_dict(), saved_path) test_min_rot_error = test_rot_error scheduler.step()
def train(args, io): device = torch.device("cuda" if args.cuda else "cpu") if args.pre_train == True: train_loader = DataLoader(ShapeNetPerm(partition='train', num_points=args.num_points), num_workers=1, batch_size=args.batch_size, shuffle=True, drop_last=True) test_loader = DataLoader(ShapeNetPerm(partition='test', num_points=args.num_points), num_workers=1, batch_size=args.test_batch_size, shuffle=True, drop_last=False) # model = pct_semantic(args).to(device) model = Pct_semantic(args).to(device) else: train_loader = DataLoader(ModelNet40(partition='train', num_points=args.num_points), num_workers=8, batch_size=args.batch_size, shuffle=True, drop_last=True) test_loader = DataLoader(ModelNet40(partition='test', num_points=args.num_points), num_workers=8, batch_size=args.test_batch_size, shuffle=True, drop_last=False) # model = Pct(args).to(device) model = pct_simple(args).to(device) print(str(model)) model = nn.DataParallel(model) if args.cont: model.load_state_dict(torch.load(args.model_path)) if args.load_pretrain: print('loading pretrained model') pretrain_model = pct_semantic(args).to(device) pretrain_model = nn.DataParallel(pretrain_model) pretrain_model.load_state_dict(torch.load(args.model_path)) pretrained_dict = pretrain_model.state_dict() model_dict = model.state_dict() #filter out unnecessary keys pretrained_dict = { k: v for k, v in pretrained_dict.items() if k in model_dict } #overwrite entries in the existing state dict model_dict.update(pretrained_dict) model.load_state_dict(model_dict, strict=True) print('pretrained model loaded') if args.use_sgd: print("Use SGD") opt = optim.SGD(model.parameters(), lr=args.lr * 100, momentum=args.momentum, weight_decay=5e-4) else: print("Use Adam") opt = optim.Adam(model.parameters(), lr=args.lr, weight_decay=1e-4) scheduler = CosineAnnealingLR(opt, args.epochs, eta_min=args.lr) if args.pre_train: criterion = nn.CrossEntropyLoss() else: criterion = cal_loss best_test_acc = 0 best_test_loss = 9999. for epoch in range(args.epochs): scheduler.step() train_loss = 0.0 count = 0.0 model.train() train_pred = [] train_true = [] idx = 0 total_time = 0.0 for data, label in (train_loader): data, label = data.to(device), label.to(device).squeeze() # print(data.size()) # print(label.size()) data = data.permute(0, 2, 1) batch_size = data.size()[0] opt.zero_grad() start_time = time.time() logits = model(data) # print(logits.size()) # print(label.size()) if args.pre_train: loss = criterion(logits, label) else: loss = criterion(logits, label, smoothing=(args.pre_train), pre_train=args.pre_train) loss.backward() opt.step() end_time = time.time() total_time += (end_time - start_time) preds = logits.max(dim=1)[1] count += batch_size train_loss += loss.item() * batch_size train_true.append(label.cpu().numpy()) train_pred.append(preds.detach().cpu().numpy()) idx += 1 print('train total time is', total_time) train_true = np.concatenate(train_true) train_pred = np.concatenate(train_pred) if args.pre_train: train_diff = train_true - train_pred correct = np.where(train_diff == 0)[0] # print('correct shape: {0}'.format(correct.shape)) # print('total shape: {0}'.format(train_diff.shape)) total_samples = train_diff.shape[0] * train_diff.shape[1] accuracy = correct.shape[0] / total_samples outstr = 'Train %d, loss: %.6f, train acc: %.6f, train avg acc: %.6f' % ( epoch, train_loss * 1.0 / count, accuracy, 0.) # metrics.accuracy_score( # train_true, train_pred), # metrics.balanced_accuracy_score( # train_true, train_pred)) else: outstr = 'Train %d, loss: %.6f, train acc: %.6f, train avg acc: %.6f' % ( epoch, train_loss * 1.0 / count, # accuracy, 0.) metrics.accuracy_score(train_true, train_pred), metrics.balanced_accuracy_score(train_true, train_pred)) io.cprint(outstr) #################### # Test #################### test_loss = 0.0 count = 0.0 model.eval() test_pred = [] test_true = [] total_time = 0.0 for data, label in test_loader: data, label = data.to(device), label.to(device).squeeze() data = data.permute(0, 2, 1) batch_size = data.size()[0] start_time = time.time() logits = model(data) end_time = time.time() total_time += (end_time - start_time) loss = criterion(logits, label) preds = logits.max(dim=1)[1] count += batch_size test_loss += loss.item() * batch_size test_true.append(label.cpu().numpy()) test_pred.append(preds.detach().cpu().numpy()) print('test total time is', total_time) test_true = np.concatenate(test_true) test_pred = np.concatenate(test_pred) # test_diff = test_true - test_pred # correct_test = np.where(test_diff == 0)[0] # print('correct shape: {0}'.format(correct.shape)) # print('total shape: {0}'.format(train_diff.shape)) # total_samples_test = test_diff.shape[0]*test_diff.shape[1] # accuracy_test = correct_test.shape[0]/total_samples_test # test_acc = 0. # test_acc = accuracy_test # avg_per_class_acc = 0. if args.pre_train: train_diff = train_true - train_pred correct = np.where(train_diff == 0)[0] # print('correct shape: {0}'.format(correct.shape)) # print('total shape: {0}'.format(train_diff.shape)) total_samples = train_diff.shape[0] * train_diff.shape[1] accuracy = correct.shape[0] / total_samples outstr = 'Train %d, loss: %.6f, train acc: %.6f, train avg acc: %.6f' % ( epoch, train_loss * 1.0 / count, accuracy, 0.) # metrics.accuracy_score( # train_true, train_pred), # metrics.balanced_accuracy_score( # train_true, train_pred)) else: test_acc = metrics.accuracy_score(test_true, test_pred) avg_per_class_acc = metrics.balanced_accuracy_score( test_true, test_pred) outstr = 'Test %d, loss: %.6f, test acc: %.6f, test avg acc: %.6f' % ( epoch, test_loss * 1.0 / count, test_acc, avg_per_class_acc) io.cprint(outstr) if args.pre_train: if test_acc >= best_test_acc: best_test_acc = test_acc torch.save(model.state_dict(), 'checkpoints/%s/models/model.t7' % args.exp_name) else: if test_loss <= best_test_loss: best_test_loss = test_loss print('Saving Checkpoint...') if args.cont: torch.save( model.state_dict(), 'checkpoints/%s/models/model1.t7' % args.exp_name) else: torch.save( model.state_dict(), 'checkpoints/%s/models/model.t7' % args.exp_name)
def test(args, io): #test_loader = DataLoader(ModelNet40(partition='test', num_points=args.num_points), # batch_size=args.test_batch_size, shuffle=True, drop_last=False) test_loader = DataLoader(ModelNet40(partition='test', num_points=args.num_points), batch_size=args.test_batch_size, shuffle=False, drop_last=False) device = torch.device("cuda" if args.cuda else "cpu") #Try to load models if args.model == 'pointnet': model = PointNet(args).to(device) elif args.model == 'dgcnn': model = DGCNN(args).to(device) elif args.model == 'TransformerBaseline': model = DGCNN_Transformer(args).to(device) elif args.model == 'TemporalTransformer': model = DGCNN_TemporalTransformer(args).to(device) elif args.model == 'TemporalTransformer_v2': model = DGCNN_TemporalTransformer_v2(args).to(device) elif args.model == 'TemporalTransformer_v3': model = DGCNN_TemporalTransformer_v3(args).to(device) elif args.model == 'pi': model = pi_DGCNN(args).to(device) elif args.model == 'pi2': model = pi_DGCNN_v2(args).to(device) elif args.model == 'pipoint': model = pipoint_DGCNN(args).to(device) else: raise Exception("Not implemented") print(model) print('Number of model parameters: {}'.format( sum([p.data.nelement() for p in model.parameters()]))) model = nn.DataParallel(model) if args.model_path: if os.path.isfile(args.model_path): print("=> loading checkpoint '{}'".format(args.model_path)) checkpoint = torch.load(args.model_path) print(checkpoint) if 'epoch' in checkpoint: args.start_epoch = checkpoint['epoch'] #best_prec1 = checkpoint['best_prec1'] #best_prec5 = checkpoint['best_prec5'] if 'state_dict' in checkpoint: model.load_state_dict(checkpoint['state_dict'], strict=False) else: model.load_state_dict(checkpoint, strict=False) print("=> loaded checkpoint '{}' (epoch {})".format( args.model_path, args.start_epoch)) else: print("=> no checkpoint found at '{}'".format(args.model_path)) #model.load_state_dict(torch.load(args.model_path)) model = model.eval() test_acc = 0.0 count = 0.0 test_true = [] test_pred = [] batch_time = AverageMeter() end = time.time() for i, (data, label) in enumerate(test_loader): data, label = data.to(device), label.to(device).squeeze() data = data.permute(0, 2, 1) batch_size = data.size()[0] if args.model in ["pi", "pipoint", "pi2"]: logits, atts = model(data) elif args.model in [ "TransformerBaseline", "TemporalTransformer", "TemporalTransformer_v2" ]: logits = model(data) else: logits, degree = model(data) preds = logits.max(dim=1)[1] test_true.append(label.cpu().numpy()) test_pred.append(preds.detach().cpu().numpy()) batch_time.update(time.time() - end) end = time.time() if i % 10 == 0: print('Test {}, Time {batch_time.val:.3f} ({batch_time.avg:.3f})'. format(i, batch_time=batch_time)) test_true = np.concatenate(test_true) test_pred = np.concatenate(test_pred) test_acc = metrics.accuracy_score(test_true, test_pred) avg_per_class_acc = metrics.balanced_accuracy_score(test_true, test_pred) per_class_acc = metrics.precision_score(test_true, test_pred, average=None) outstr = 'Test :: test acc: %.6f, test avg acc: %.6f' % (test_acc, avg_per_class_acc) outstr_2 = 'Test per class acc: {}' % per_class_acc io.cprint(outstr) io.cprint(outstr_2) if args.model in ["pi", "pipoint", "pi2"]: for j in range(4): io.cprint('Att {} : {}'.format(j, atts[j].mean().item()))
def train(args, io): train_dataset = ModelNet40(partition='train', num_points=args.num_points) train_loader = DataLoader(train_dataset, num_workers=0, batch_size=args.batch_size, shuffle=True, drop_last=True) test_dataset = ModelNet40(partition='test', num_points=args.num_points) test_loader = DataLoader(test_dataset, num_workers=0, batch_size=args.test_batch_size, shuffle=False, drop_last=False) device = 0 #Try to load models if args.model == 'hdgn': model = HDGN(args).to(device) else: raise Exception("Not implemented") print(str(model)) sys.stdout.flush() print("Let's use", torch.cuda.device_count(), "GPUs!") if args.use_sgd: print("Use SGD") opt = optim.SGD(model.parameters(), lr=args.lr * 100, momentum=args.momentum, weight_decay=1e-4) else: print("Use Adam") opt = optim.Adam(model.parameters(), lr=args.lr, weight_decay=1e-4) scheduler = CosineAnnealingLR(opt, args.epochs, eta_min=args.lr) criterion = cal_loss best_test_acc = 0 for epoch in range(args.epochs): scheduler.step() #################### # Train #################### train_loss = 0.0 count = 0.0 model.train() train_pred = [] train_true = [] for data, label in train_loader: data, label = data.to(device), label.to(device).squeeze(dim=-1) data = data.permute(0, 2, 1) batch_size = data.size()[0] opt.zero_grad() logits = model(data) loss = criterion(logits, label) loss.backward() opt.step() preds = logits.max(dim=1)[1] count += batch_size train_loss += loss.item() * batch_size train_true.append(label.cpu().numpy()) train_pred.append(preds.detach().cpu().numpy()) train_true = np.concatenate(train_true) train_pred = np.concatenate(train_pred) outstr = 'Train %d, loss: %.6f, train acc: %.6f, train avg acc: %.6f' % ( epoch, train_loss * 1.0 / count, metrics.accuracy_score(train_true, train_pred), metrics.balanced_accuracy_score(train_true, train_pred)) io.cprint(outstr) sys.stdout.flush() #################### # Test #################### test_loss = 0.0 count = 0.0 model.eval() test_pred = [] test_true = [] for data, label in test_loader: data, label = data.to(device), label.to(device).squeeze(dim=-1) data = data.permute(0, 2, 1) batch_size = data.size()[0] logits, wsl = model(data) loss = criterion(logits, label) + wsl preds = logits.max(dim=1)[1] count += batch_size test_loss += loss.item() * batch_size test_true.append(label.cpu().numpy()) test_pred.append(preds.detach().cpu().numpy()) test_true = np.concatenate(test_true) test_pred = np.concatenate(test_pred) test_acc = metrics.accuracy_score(test_true, test_pred) mean_test_acc = metrics.accuracy_score(test_true, test_pred) avg_per_class_acc = metrics.balanced_accuracy_score( test_true, test_pred) outstr = 'Test %d, loss: %.6f, test acc: %.6f, test avg acc: %.6f' % ( epoch, test_loss * 1.0 / count, mean_test_acc, avg_per_class_acc) io.cprint(outstr) sys.stdout.flush() if mean_test_acc >= best_test_acc: best_test_acc = mean_test_acc torch.save(model.state_dict(), 'checkpoints/%s/models/model.t7' % args.exp_name)
def train(args, io): train_loader = DataLoader(ModelNet40(partition='train', num_points=args.num_points), num_workers=8, batch_size=args.batch_size, shuffle=True, drop_last=True) test_loader = DataLoader(ModelNet40(partition='test', num_points=args.num_points), num_workers=8, batch_size=args.test_batch_size, shuffle=False, drop_last=False) device = torch.device("cuda" if args.cuda else "cpu") #Try to load models if args.model == 'pointnet': model = PointNet(args).to(device) elif args.model == 'dgcnn': model = DGCNN(args).to(device) else: model = PaiNet(args).to(device) #raise Exception("Not implemented") print(str(model)) params = sum(p.numel() for p in model.parameters() if p.requires_grad) print("Total number of parameters is: {}".format(params)) model = nn.DataParallel(model) print("Let's use", torch.cuda.device_count(), "GPUs!") # if os.path.exists('checkpoints/%s/models/model_%s.t7'% (args.exp_name, args.model)): # checkpoint_dict = torch.load('./checkpoints/%s/models/model_%s.t7'% (args.exp_name, args.model), map_location=device) # model_dict = model.state_dict() # pretrained_dict = checkpoint_dict # pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict and "one_padding" not in k} # model_dict.update(pretrained_dict) # model.load_state_dict(pretrained_dict, strict=False) # #model.load_state_dict(checkpoint_dict, strict=True) # print("Load model from './checkpoints/%s/models/model_%s.t7 !'"% (args.exp_name, args.model)) if args.use_sgd: print("Use SGD") trainables_wo_bn = [param for name, param in model.named_parameters() if param.requires_grad \ and "bn" not in name \ and "kernals" not in name] trainables_wt_bn = [ param for name, param in model.named_parameters() if param.requires_grad and ('bn' in name or "kernals" in name) ] opt = optim.SGD([{ 'params': trainables_wo_bn, 'weight_decay': 5e-5 }, { 'params': trainables_wt_bn }], lr=args.lr * 100, momentum=args.momentum) #opt = optim.SGD(model.parameters(), lr=args.lr*100, momentum=args.momentum, weight_decay=1e-4) else: print("Use Adam") opt = optim.Adam(model.parameters(), lr=args.lr, weight_decay=1e-4) scheduler = CosineAnnealingLR(opt, args.epochs, eta_min=args.lr) criterion = cal_loss best_test_acc = 0 # 0.931929 for epoch in range(args.epochs): #################### # Train #################### train_loss = 0.0 count = 0.0 model.train() train_pred = [] train_true = [] for b, (data, label) in enumerate(tqdm(train_loader, ncols=0)): data, label = data.to(device), label.to(device).squeeze() data = data.permute(0, 2, 1) batch_size = data.size()[0] opt.zero_grad() logits = model(data) loss = criterion(logits, label) loss.backward() opt.step() preds = logits.max(dim=1)[1] count += batch_size train_loss += loss.item() * batch_size train_true.append(label.cpu().numpy()) train_pred.append(preds.detach().cpu().numpy()) scheduler.step() # if epoch % 100 == 0: # scheduler = CosineAnnealingLR(opt, 50, eta_min=args.lr) train_true = np.concatenate(train_true) train_pred = np.concatenate(train_pred) outstr = 'Train %d, loss: %.6f, train acc: %.6f, train avg acc: %.6f, lr: %.6f' % ( epoch, train_loss * 1.0 / count, metrics.accuracy_score(train_true, train_pred), metrics.balanced_accuracy_score(train_true, train_pred), scheduler.get_lr()[0]) io.cprint(outstr) #################### # Test #################### test_loss = 0.0 count = 0.0 model.eval() test_pred = [] test_true = [] with torch.no_grad(): for data, label in test_loader: data, label = data.to(device), label.to(device).squeeze() data = data.permute(0, 2, 1) batch_size = data.size()[0] logits = model(data) loss = criterion(logits, label) preds = logits.max(dim=1)[1] count += batch_size test_loss += loss.item() * batch_size test_true.append(label.cpu().numpy()) test_pred.append(preds.detach().cpu().numpy()) test_true = np.concatenate(test_true) test_pred = np.concatenate(test_pred) test_acc = metrics.accuracy_score(test_true, test_pred) avg_per_class_acc = metrics.balanced_accuracy_score( test_true, test_pred) if test_acc >= best_test_acc: best_test_acc = test_acc torch.save( model.state_dict(), 'checkpoints/%s/models/model_%s.t7' % (args.exp_name, args.model)) outstr = 'Test %d, loss: %.6f, test acc: %.6f, test avg acc: %.6f, test best: %.6f' % ( epoch, test_loss * 1.0 / count, test_acc, avg_per_class_acc, best_test_acc) io.cprint(outstr)
def main(): arg_bool = lambda x: x.lower() in ['true', 't', '1'] parser = argparse.ArgumentParser(description='Point Cloud Registration') parser.add_argument('--exp_name', type=str, default='exp', metavar='N', help='Name of the experiment') parser.add_argument('--num_iter', type=int, default=4, metavar='N', help='Number of iteration inside the network') parser.add_argument('--emb_nn', type=str, default='GNN', metavar='N', help='Feature extraction method. [GNN]') parser.add_argument('--emb_dims', type=int, default=64, metavar='N', help='Dimension of embeddings') parser.add_argument('--batch_size', type=int, default=16, 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=50, metavar='N', help='number of episode to train ') parser.add_argument('--unseen', type=arg_bool, default='True', help='Test on unseen categories') parser.add_argument('--gaussian_noise', type=arg_bool, default='True', help='Wheter to add gaussian noise') parser.add_argument( '--alpha', type=float, default=0.75, metavar='N', help='Fraction of points when sampling partial point cloud') parser.add_argument('--factor', type=float, default=4, metavar='N', help='Divided factor for rotations') # 旋转角度在pi/4范围内 parser.add_argument('--pretrained', type=arg_bool, default='False', help='Load pretrained weight') # 旋转角度在pi/4范围内 args = parser.parse_args() print(args) ##### make checkpoint directory and backup ##### if not os.path.exists('checkpoints'): os.makedirs('checkpoints') if not os.path.exists('checkpoints/' + args.exp_name): os.makedirs('checkpoints/' + args.exp_name) if not os.path.exists('checkpoints/' + args.exp_name + '/' + 'models'): os.makedirs('checkpoints/' + args.exp_name + '/' + 'models') os.system('cp main.py checkpoints' + '/' + args.exp_name + '/' + 'main.py.backup') os.system('cp model.py checkpoints' + '/' + args.exp_name + '/' + 'model.py.backup') os.system('cp data.py checkpoints' + '/' + args.exp_name + '/' + 'data.py.backup') ##### make checkpoint directory and backup ##### io = IOStream('checkpoints/' + args.exp_name + '/log.txt') io.cprint(str(args)) ##### load data ##### train_loader = DataLoader(ModelNet40(partition='train', alpha=args.alpha, gaussian_noise=args.gaussian_noise, unseen=args.unseen, factor=args.factor), batch_size=args.batch_size, shuffle=True, drop_last=True, num_workers=8) test_loader = DataLoader(ModelNet40(partition='test', alpha=args.alpha, gaussian_noise=args.gaussian_noise, unseen=args.unseen, factor=args.factor), batch_size=args.test_batch_size, shuffle=False, drop_last=False, num_workers=8) ##### load model ##### net = MFGNet(GNN(args.emb_dims), args).cuda() # net.load_state_dict(torch.load('model_gaussian.t7')) # for param in net.parameters(): # print(param.name, param.size()) ##### train ##### train(args, net, train_loader, test_loader, io) io.close()
def main(): parser = argparse.ArgumentParser(description='Point Cloud Registration') parser.add_argument('--exp_name', type=str, default='exp', metavar='N', help='Name of the experiment') parser.add_argument('--model', type=str, default='dcp', metavar='N', choices=['dcp'], help='Model to use, [dcp]') parser.add_argument('--emb_nn', type=str, default='dgcnn', metavar='N', choices=['pointnet', 'dgcnn'], help='Embedding nn to use, [pointnet, dgcnn]') parser.add_argument( '--pointer', type=str, default='transformer', metavar='N', choices=['identity', 'transformer'], help='Attention-based pointer generator to use, [identity, transformer]' ) parser.add_argument('--head', type=str, default='svd', metavar='N', choices=[ 'mlp', 'svd', ], help='Head to use, [mlp, svd]') parser.add_argument('--emb_dims', type=int, default=512, metavar='N', help='Dimension of embeddings') parser.add_argument('--n_blocks', type=int, default=1, metavar='N', help='Num of blocks of encoder&decoder') parser.add_argument('--n_heads', type=int, default=16, metavar='N', help='Num of heads in multiheadedattention') parser.add_argument('--ff_dims', type=int, default=1024, metavar='N', help='Num of dimensions of fc in transformer') parser.add_argument('--dropout', type=float, default=0.0, metavar='N', help='Dropout ratio in transformer') parser.add_argument('--batch_size', type=int, default=16, metavar='batch_size', help='Size of batch)') parser.add_argument('--test_batch_size', type=int, default=10, metavar='batch_size', help='Size of batch)') parser.add_argument('--epochs', type=int, default=250, metavar='N', help='number of episode to train ') parser.add_argument('--use_sgd', action='store_true', default=False, 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('--no_cuda', action='store_true', default=False, help='enables CUDA training') parser.add_argument('--seed', type=int, default=1234, metavar='S', help='random seed (default: 1)') parser.add_argument('--eval', action='store_true', default=False, help='evaluate the model') parser.add_argument('--cycle', type=bool, default=False, metavar='N', help='Whether to use cycle consistency') parser.add_argument('--gaussian_noise', type=bool, default=False, metavar='N', help='Wheter to add gaussian noise') parser.add_argument('--unseen', type=bool, default=False, metavar='N', help='Wheter to test on unseen category') parser.add_argument('--num_points', type=int, default=1024, metavar='N', help='Num of points to use') parser.add_argument('--dataset', type=str, default='modelnet40', choices=['modelnet40', 'threedmatch'], metavar='N', help='dataset to use') parser.add_argument('--factor', type=float, default=4, metavar='N', help='Divided factor for rotations') parser.add_argument('--model_path', type=str, default='', metavar='N', help='Pretrained model path') parser.add_argument('--betas', type=float, default=(0.9, 0.999), metavar='N', nargs='+', help='Betas for adam') parser.add_argument('--same_pointclouds', type=bool, default=True, metavar='N', help='R*src + t should be exactly same as target') parser.add_argument('--debug', type=bool, default=False, metavar='N', help='saves variables in folder variables_storage') parser.add_argument('--num_itr_test', type=int, default=1, metavar='N', help='Num of net() during testing') parser.add_argument( '--loss', type=str, default='cross_entropy_corr', metavar='N', choices=['cross_entropy_corr', 'mse_transf'], help='loss function: choose one of [mse_transf or cross_entropy_corr]') parser.add_argument('--cut_plane', type=bool, default=False, metavar='N', help='generates partial data') parser.add_argument('--one_cloud', type=bool, default=False, metavar='N', help='test for one unseen cloud') parser.add_argument( '--partial', type=float, default=0.0, metavar='N', help='partial = 0.1 ==> (num_points*partial) will be removed') parser.add_argument('--pretrained', type=bool, default=False, metavar='N', help='load pretrained model') args = parser.parse_args() # for deterministic training torch.backends.cudnn.deterministic = True torch.manual_seed(args.seed) torch.cuda.manual_seed_all(args.seed) np.random.seed(args.seed) boardio = SummaryWriter(log_dir='checkpoints/' + args.exp_name) _init_(args) textio = IOStream('checkpoints/' + args.exp_name + '/run.log') textio.cprint(str(args)) # dataloading num_workers = 32 if args.dataset == 'modelnet40': train_dataset = ModelNet40(num_points=args.num_points, partition='train', gaussian_noise=args.gaussian_noise, unseen=args.unseen, factor=args.factor, same_pointclouds=args.same_pointclouds, partial=args.partial, cut_plane=args.cut_plane) train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, drop_last=True, num_workers=num_workers) test_dataset = ModelNet40(num_points=args.num_points, partition='test', gaussian_noise=args.gaussian_noise, unseen=args.unseen, factor=args.factor, same_pointclouds=args.same_pointclouds, partial=args.partial, cut_plane=args.cut_plane) test_loader = DataLoader(test_dataset, batch_size=args.test_batch_size, shuffle=False, drop_last=False, num_workers=num_workers) else: raise Exception("not implemented") # model loading if args.model == 'dcp': net = DCP(args).cuda() if args.eval: if args.model_path is '': model_path = 'checkpoints' + '/' + args.exp_name + '/models/model.best.t7' else: model_path = args.model_path print("Model loaded from ", model_path) if not os.path.exists(model_path): print("can't find pretrained model") return net.load_state_dict(torch.load(model_path), strict=False) if args.pretrained: if args.model_path == '': print( 'Please specify path to pretrained weights \n For Ex: checkpoints/partial_global_512_identical/models/model.best.t7' ) else: model_path = args.model_path print("Using pretrained weights stored at:\n{}".format(model_path)) net.load_state_dict(torch.load(model_path), strict=False) if torch.cuda.device_count() > 1: net = nn.DataParallel(net) print("Let's use", torch.cuda.device_count(), "GPUs!") else: raise Exception('Not implemented') # training and evaluation if args.eval: if args.one_cloud: # testing on a single point cloud print("one_cloud") test_bunny(args, net) else: test(args, net, test_loader, boardio, textio) else: train(args, net, train_loader, test_loader, boardio, textio) print('FINISH') boardio.close()
def train(args, io): train_loader = DataLoader(ModelNet40(partition='train', num_points=args.num_points), num_workers=8, batch_size=args.batch_size, shuffle=True, drop_last=True) test_loader = DataLoader(ModelNet40(partition='test', num_points=args.num_points), num_workers=8, batch_size=args.test_batch_size, shuffle=True, drop_last=False) device = torch.device("cuda" if args.cuda else "cpu") #Try to load models if args.model == 'pointnet': model = PointNet(args).to(device) elif args.model == 'dgcnn': model = DGCNN(args).to(device) elif args.model == 'semigcn': model = SemiGCN(args).to(device) else: raise Exception("Not implemented") print(str(model)) model = nn.DataParallel(model) print("Let's use", torch.cuda.device_count(), "GPUs!") if args.use_sgd: print("Use SGD") opt = optim.SGD(model.parameters(), lr=args.lr*100, momentum=args.momentum, weight_decay=1e-4) else: print("Use Adam") opt = optim.Adam(model.parameters(), lr=args.lr, weight_decay=1e-4) # optionally resume from a checkpoint if args.resume: if os.path.isfile(args.resume): print("=> loading checkpoint '{}'".format(args.resume)) checkpoint = torch.load(args.resume) args.start_epoch = checkpoint['epoch'] model.load_state_dict(checkpoint['state_dict']) opt.load_state_dict(checkpoint['opt']) print("=> loaded checkpoint '{}' (epoch {})" .format(args.resume, checkpoint['epoch'])) else: print("=> no checkpoint found at '{}'".format(args.resume)) #scheduler = CosineAnnealingLR(opt, args.epochs, eta_min=args.lr, last_epoch=args.start_epoch-1) scheduler = torch.optim.lr_scheduler.StepLR(opt, step_size=20, gamma=0.8)#0.7 #scheduler = torch.optim.lr_scheduler.ExponentialLR(opt, gamma=0.9825, last_epoch=args.start_epoch-1) criterion = cal_loss best_test_acc = 0 for epoch in range(args.start_epoch, args.epochs): #scheduler.step() #################### # Train #################### train_loss = 0.0 count = 0.0 model.train() train_pred = [] train_true = [] for data, label in train_loader: data, label = data.to(device), label.to(device).squeeze() data = data.permute(0, 2, 1) batch_size = data.size()[0] opt.zero_grad() logits = model(data) loss = criterion(logits, label) loss.backward() opt.step() preds = logits.max(dim=1)[1] count += batch_size train_loss += loss.item() * batch_size train_true.append(label.cpu().numpy()) train_pred.append(preds.detach().cpu().numpy()) scheduler.step() train_true = np.concatenate(train_true) train_pred = np.concatenate(train_pred) outstr = 'Train %d, loss: %.6f, train acc: %.6f, train avg acc: %.6f' % (epoch, train_loss*1.0/count, metrics.accuracy_score( train_true, train_pred), metrics.balanced_accuracy_score( train_true, train_pred)) io.cprint(outstr) if epoch%10 == 0: # save running checkpoint per 10 epoch torch.save({'epoch': epoch + 1, 'arch': args.model, 'state_dict': model.state_dict(), 'opt' : opt.state_dict()}, 'checkpoints/%s/models/checkpoint_latest.pth.tar' % args.exp_name) #################### # Test #################### test_loss = 0.0 count = 0.0 model.eval() test_pred = [] test_true = [] for data, label in test_loader: data, label = data.to(device), label.to(device).squeeze() data = data.permute(0, 2, 1) batch_size = data.size()[0] logits = model(data) loss = criterion(logits, label) preds = logits.max(dim=1)[1] count += batch_size test_loss += loss.item() * batch_size test_true.append(label.cpu().numpy()) test_pred.append(preds.detach().cpu().numpy()) test_true = np.concatenate(test_true) test_pred = np.concatenate(test_pred) test_acc = metrics.accuracy_score(test_true, test_pred) avg_per_class_acc = metrics.balanced_accuracy_score(test_true, test_pred) outstr = 'Test %d, loss: %.6f, test acc: %.6f, test avg acc: %.6f' % (epoch, test_loss*1.0/count, test_acc, avg_per_class_acc) io.cprint(outstr) if test_acc >= best_test_acc: best_test_acc = test_acc torch.save({'epoch': epoch + 1, 'arch': args.model, 'state_dict': model.state_dict(), 'opt' : opt.state_dict()}, 'checkpoints/%s/models/checkpoint_best.pth.tar' % args.exp_name)
def train(args, io): train_loader = DataLoader(ModelNet40(partition='train', num_points=args.num_points), num_workers=8, batch_size=args.batch_size, shuffle=True, drop_last=True) test_loader = DataLoader(ModelNet40(partition='test', num_points=args.num_points), num_workers=8, batch_size=args.test_batch_size, shuffle=True, drop_last=False) device = torch.device("cuda" if args.cuda else "cpu") model = Pct(args).to(device) print(str(model)) model = nn.DataParallel(model) if args.use_sgd: print("Use SGD") opt = optim.SGD(model.parameters(), lr=args.lr*100, momentum=args.momentum, weight_decay=1e-4) else: print("Use Adam") opt = optim.Adam(model.parameters(), lr=args.lr, weight_decay=1e-4) scheduler = CosineAnnealingLR(opt, args.epochs, eta_min=args.lr) criterion = cal_loss best_test_acc = 0 for epoch in range(args.epochs): scheduler.step() train_loss = 0.0 count = 0.0 model.train() train_pred = [] train_true = [] idx = 0 total_time = 0.0 for data, label in (train_loader): data, label = data.to(device), label.to(device).squeeze() data = data.permute(0, 2, 1) batch_size = data.size()[0] opt.zero_grad() start_time = time.time() logits = model(data) loss = criterion(logits, label) loss.backward() opt.step() end_time = time.time() total_time += (end_time - start_time) preds = logits.max(dim=1)[1] count += batch_size train_loss += loss.item() * batch_size train_true.append(label.cpu().numpy()) train_pred.append(preds.detach().cpu().numpy()) idx += 1 print ('train total time is',total_time) train_true = np.concatenate(train_true) train_pred = np.concatenate(train_pred) outstr = 'Train %d, loss: %.6f, train acc: %.6f, train avg acc: %.6f' % (epoch, train_loss*1.0/count, metrics.accuracy_score( train_true, train_pred), metrics.balanced_accuracy_score( train_true, train_pred)) io.cprint(outstr) #################### # Test #################### test_loss = 0.0 count = 0.0 model.eval() test_pred = [] test_true = [] total_time = 0.0 for data, label in test_loader: data, label = data.to(device), label.to(device).squeeze() data = data.permute(0, 2, 1) batch_size = data.size()[0] start_time = time.time() logits = model(data) end_time = time.time() total_time += (end_time - start_time) loss = criterion(logits, label) preds = logits.max(dim=1)[1] count += batch_size test_loss += loss.item() * batch_size test_true.append(label.cpu().numpy()) test_pred.append(preds.detach().cpu().numpy()) print ('test total time is', total_time) test_true = np.concatenate(test_true) test_pred = np.concatenate(test_pred) test_acc = metrics.accuracy_score(test_true, test_pred) avg_per_class_acc = metrics.balanced_accuracy_score(test_true, test_pred) outstr = 'Test %d, loss: %.6f, test acc: %.6f, test avg acc: %.6f' % (epoch, test_loss*1.0/count, test_acc, avg_per_class_acc) io.cprint(outstr) if test_acc >= best_test_acc: best_test_acc = test_acc torch.save(model.state_dict(), 'checkpoints/%s/models/model.t7' % args.exp_name)
def train(args, io): train_loader = DataLoader(ModelNet40(partition='train', num_points=args.num_points), num_workers=8, batch_size=args.batch_size, shuffle=True, drop_last=False) test_loader = DataLoader(ModelNet40(partition='test', num_points=args.num_points), num_workers=8, batch_size=args.test_batch_size, shuffle=False, drop_last=False) device = torch.device("cuda" if args.cuda else "cpu") #Try to load models if args.model == 'pointnet': model = PointNet(args).to(device) elif args.model == 'dgcnn': model = DGCNN(args).to(device) elif args.model == 'TransformerBaseline': model = DGCNN_Transformer(args).to(device) elif args.model == 'TemporalTransformer': model = DGCNN_TemporalTransformer(args).to(device) elif args.model == 'TemporalTransformer_v2': model = DGCNN_TemporalTransformer_v2(args).to(device) elif args.model == 'TemporalTransformer_v3': model = DGCNN_TemporalTransformer_v3(args).to(device) elif args.model == 'pi': model = pi_DGCNN(args).to(device) elif args.model == 'pi2': model = pi_DGCNN_v2(args).to(device) elif args.model == 'pipoint': model = pipoint_DGCNN(args).to(device) else: raise Exception("Not implemented") print(str(model)) #model = nn.DataParallel(model, device_ids=list(range(3))) model = nn.DataParallel(model) print("Let's use", torch.cuda.device_count(), "GPUs!") # get the number of model parameters print('Number of model parameters: {}'.format( sum([p.data.nelement() for p in model.parameters()]))) if args.use_sgd: print("Use SGD") optimizer = optim.SGD(model.parameters(), lr=args.lr * 100, momentum=args.momentum, weight_decay=1e-4) else: print("Use Adam") optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=1e-4) scheduler = CosineAnnealingLR(optimizer, args.epochs, eta_min=args.lr) criterion = cal_loss best_test_acc = 0 if args.model_path: if os.path.isfile(args.model_path): print("=> loading checkpoint '{}'".format(args.model_path)) checkpoint = torch.load(args.model_path) print(checkpoint) if 'epoch' in checkpoint: args.start_epoch = checkpoint['epoch'] #best_prec1 = checkpoint['best_prec1'] #best_prec5 = checkpoint['best_prec5'] if 'state_dict' in checkpoint: model.load_state_dict(checkpoint['state_dict']) else: model.load_state_dict(checkpoint, strict=False) if 'optimizer' in checkpoint: optimizer.load_state_dict(checkpoint['optimizer']) print("=> loaded checkpoint '{}' (epoch {})".format( args.model_path, args.start_epoch)) else: print("=> no checkpoint found at '{}'".format(args.resume)) end = time.time() for epoch in range(args.start_epoch, args.epochs): scheduler.step() #################### # Train #################### train_loss = 0.0 count = 0.0 model.train() train_pred = [] train_true = [] batch_time = AverageMeter() data_time = AverageMeter() end = time.time() for i, (data, label) in enumerate(train_loader): data_time.update(time.time() - end) data, label = data.to(device), label.to(device).squeeze() data = data.permute(0, 2, 1) batch_size = data.size()[0] optimizer.zero_grad() if args.model in ["pi", "pipoint", "pi2"]: logits, atts = model(data) elif args.model in [ "TransformerBaseline", "TemporalTransformer", "TemporalTransformer_v2", "TemporalTransformer_v3" ]: logits = model(data) else: logits, degree = model(data) ''' if args.visualize == True: print(args.visualize) import matplotlib.pyplot as plt #cmap = plt.cm.get_cmap("hsv", 30) cmap = plt.cm.get_cmap("binary", 40) cmap = np.array([cmap(i) for i in range(40)])[:,:3] obj = degree[7,:,:3].cpu().numpy() obj_degree = degree[7,:,3:].squeeze() obj_degree = obj_degree.cpu().numpy().astype(int) obj_max = np.max(obj_degree) obj_min = np.min(obj_degree) print(obj_max) print(obj_min) for i in range(obj_min, obj_max): print("{} : {}".format(i, sum(obj_degree == i))) gt = cmap[obj_degree-obj_min, :] showpoints(obj, gt, gt, ballradius=3) ''' loss = criterion(logits, label, smoothing=False) loss.backward() optimizer.step() preds = logits.max(dim=1)[1] count += batch_size train_loss += loss.item() * batch_size train_true.append(label.cpu().numpy()) train_pred.append(preds.detach().cpu().numpy()) batch_time.update(time.time() - end) end = time.time() if i % 10 == 0: print_str = 'Train {}, loss {}, Time {batch_time.val:.3f} ({batch_time.avg:.3f}), Data {data_time.val:.3f} ({data_time.avg:.3f})'.format( epoch, train_loss * 1.0 / count, batch_time=batch_time, data_time=data_time) print(print_str) train_true = np.concatenate(train_true) train_pred = np.concatenate(train_pred) outstr = 'Train %d, loss: %.6f, train acc: %.6f, train avg acc: %.6f' % ( epoch, train_loss * 1.0 / count, metrics.accuracy_score(train_true, train_pred), metrics.balanced_accuracy_score(train_true, train_pred)) #outstr = 'Train {}, Time {batch_time.val:.3f} ({batch_time.avg:.3f}), Data {data_time.val:.3f} ({data_time.avg:.3f}), loss: {}, train acc: {}, train avg acc: {}'.format(epoch, batch_time=batch_time, data_time=data_time, train_loss*1.0/count, metrics.accuracy_score(train_true, train_pred), metrics.balanced_accuracy_score(train_true, train_pred)) io.cprint(outstr) #################### # Test #################### with torch.no_grad(): test_loss = 0.0 count = 0.0 model.eval() test_pred = [] test_true = [] batch_time = AverageMeter() losses = AverageMeter() end = time.time() for j, (data, label) in enumerate(test_loader): data, label = data.to(device), label.to(device).squeeze() data = data.permute(0, 2, 1) batch_size = data.size()[0] if args.model in ["pi", "pipoint", "pi2"]: logits, atts = model(data) elif args.model in [ "TransformerBaseline", "TemporalTransformer", "TemporalTransformer_v2", "TemporalTransformer_v3" ]: logits = model(data) else: logits, degree = model(data) loss = criterion(logits, label, smoothing=False) preds = logits.max(dim=1)[1] count += batch_size test_loss += loss.item() * batch_size test_true.append(label.cpu().numpy()) test_pred.append(preds.detach().cpu().numpy()) batch_time.update(time.time() - end) end = time.time() if j % 10 == 0: print( 'Test {}, Loss {}, Time {batch_time.val:.3f} ({batch_time.avg:.3f})' .format(j, test_loss * 1.0 / count, batch_time=batch_time)) test_true = np.concatenate(test_true) test_pred = np.concatenate(test_pred) test_acc = metrics.accuracy_score(test_true, test_pred) avg_per_class_acc = metrics.balanced_accuracy_score( test_true, test_pred) #per_class_acc = metrics.precision_score(test_true, test_pred, average=None) #outstr = 'Test {}, loss: {}, train acc: {}, train avg acc: {}'.format(epoch, batch_time=batch_time, data_time=data_time, test_loss*1.0/count, test_acc, avg_per_class_acc) outstr = 'Test %d, loss: %.6f, test acc: %.6f, test avg acc: %.6f' % ( epoch, test_loss * 1.0 / count, test_acc, avg_per_class_acc) #outstr_2 = 'Test per class acc: {}'%per_class_acc io.cprint(outstr) #io.cprint(outstr_2) if args.model in ["pi", "pipoint", "pi2"]: for j in range(4): io.cprint('Att {} : {}'.format(j, atts[j].mean().item())) is_best = test_acc >= best_test_acc if is_best: best_test_acc = test_acc save_checkpoint( { 'epoch': epoch + 1, 'state_dict': model.state_dict(), 'optimizer': optimizer.state_dict(), }, is_best, args.exp_name)
def train(args, io): train_loader = DataLoader(ModelNet40(partition='train', num_points=args.num_points), num_workers=8, batch_size=args.batch_size, shuffle=True, drop_last=True) test_loader = DataLoader(ModelNet40(partition='test', num_points=args.num_points), num_workers=8, batch_size=args.test_batch_size, shuffle=True, drop_last=False) device = torch.device("cuda" if args.cuda else "cpu") #Try to load models if args.model == 'pointnet': model = PointNet(args).to(device) elif args.model == 'dgcnn': model = DGCNN_cls(args).to(device) else: raise Exception("Not implemented") print(str(model)) model = nn.DataParallel(model) print("Let's use", torch.cuda.device_count(), "GPUs!") if args.use_sgd: print("Use SGD") opt = optim.SGD(model.parameters(), lr=args.lr*100, momentum=args.momentum, weight_decay=1e-4) else: print("Use Adam") opt = optim.Adam(model.parameters(), lr=args.lr, weight_decay=1e-4) if args.scheduler == 'cos': scheduler = CosineAnnealingLR(opt, args.epochs, eta_min=1e-3) elif args.scheduler == 'step': scheduler = StepLR(opt, step_size=20, gamma=0.7) criterion = cal_loss #criterion = cal_loss_v2 best_test_acc = 0 for epoch in range(args.epochs): #################### # Train #################### train_loss = 0.0 count = 0.0 model.train() train_pred = [] train_true = [] for data, label in train_loader: data, label = data.to(device), label.to(device).squeeze() data = data.permute(0, 2, 1) batch_size = data.size()[0] opt.zero_grad() logits = model(data) loss = criterion(logits, label) loss.backward() opt.step() preds = logits.max(dim=1)[1] count += batch_size train_loss += loss.item() * batch_size train_true.append(label.cpu().numpy()) train_pred.append(preds.detach().cpu().numpy()) if args.scheduler == 'cos': scheduler.step() elif args.scheduler == 'step': if opt.param_groups[0]['lr'] > 1e-5: scheduler.step() if opt.param_groups[0]['lr'] < 1e-5: for param_group in opt.param_groups: param_group['lr'] = 1e-5 train_true = np.concatenate(train_true) train_pred = np.concatenate(train_pred) outstr = 'Train %d, loss: %.6f, train acc: %.6f, train avg acc: %.6f' % (epoch, train_loss*1.0/count, metrics.accuracy_score( train_true, train_pred), metrics.balanced_accuracy_score( train_true, train_pred)) io.cprint(outstr) #################### # Test #################### test_loss = 0.0 count = 0.0 model.eval() test_pred = [] test_true = [] for data, label in test_loader: data, label = data.to(device), label.to(device).squeeze() data = data.permute(0, 2, 1) batch_size = data.size()[0] logits = model(data) loss = criterion(logits, label) preds = logits.max(dim=1)[1] count += batch_size test_loss += loss.item() * batch_size test_true.append(label.cpu().numpy()) test_pred.append(preds.detach().cpu().numpy()) test_true = np.concatenate(test_true) test_pred = np.concatenate(test_pred) test_acc = metrics.accuracy_score(test_true, test_pred) avg_per_class_acc = metrics.balanced_accuracy_score(test_true, test_pred) outstr = 'Test %d, loss: %.6f, test acc: %.6f, test avg acc: %.6f' % (epoch, test_loss*1.0/count, test_acc, avg_per_class_acc) io.cprint(outstr) if test_acc >= best_test_acc: best_test_acc = test_acc torch.save(model.state_dict(), 'checkpoints/%s/models/model.t7' % args.exp_name)
import torch from torch.autograd import Variable import torch.nn as nn import torch.utils.data as data from config import get_test_config from data import ModelNet40 from models import MeshNet from utils import point_wise_L1_loss, get_unit_diamond_vertices, axis_aligned_miou, point_wise_mse_loss #, stochastic_loss root_path = '/content/drive/MyDrive/DL_diamond_cutting/MeshNet/' cfg = get_test_config(root_path) os.environ['CUDA_VISIBLE_DEVICES'] = cfg['cuda_devices'] use_gpu = torch.cuda.is_available() data_set = ModelNet40(cfg=cfg['dataset'], root_path=root_path, part='test') data_loader = data.DataLoader(data_set, batch_size=1, num_workers=4, shuffle=False, pin_memory=False) def test_model(model): criterion = nn.L1Loss() running_loss = 0.0 running_l1_loss = 0.0 running_scale_loss = 0.0 running_center_loss = 0.0 running_rotation_loss = 0.0
def main(): parser = argparse.ArgumentParser(description='Point Cloud Registration') parser.add_argument('--exp_name', type=str, default='exp', metavar='N', help='Name of the experiment') parser.add_argument('--model', type=str, default='dcp', metavar='N', choices=['dcp'], help='Model to use, [dcp]') parser.add_argument('--emb_nn', type=str, default='pointnet', metavar='N', choices=['pointnet', 'dgcnn'], help='Embedding nn to use, [pointnet, dgcnn]') parser.add_argument( '--pointer', type=str, default='transformer', metavar='N', choices=['identity', 'transformer'], help='Attention-based pointer generator to use, [identity, transformer]' ) parser.add_argument('--head', type=str, default='svd', metavar='N', choices=[ 'mlp', 'svd', ], help='Head to use, [mlp, svd]') parser.add_argument('--emb_dims', type=int, default=512, metavar='N', help='Dimension of embeddings') parser.add_argument('--n_blocks', type=int, default=1, metavar='N', help='Num of blocks of encoder&decoder') parser.add_argument('--n_heads', type=int, default=4, metavar='N', help='Num of heads in multiheadedattention') parser.add_argument('--ff_dims', type=int, default=1024, metavar='N', help='Num of dimensions of fc in transformer') parser.add_argument('--dropout', type=float, default=0.0, metavar='N', help='Dropout ratio in transformer') parser.add_argument('--batch_size', type=int, default=32, metavar='batch_size', help='Size of batch)') parser.add_argument('--test_batch_size', type=int, default=10, metavar='batch_size', help='Size of batch)') parser.add_argument('--epochs', type=int, default=250, metavar='N', help='number of episode to train ') parser.add_argument('--use_sgd', action='store_true', default=False, 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('--no_cuda', action='store_true', default=False, help='enables CUDA training') parser.add_argument('--seed', type=int, default=1234, metavar='S', help='random seed (default: 1)') parser.add_argument('--eval', action='store_true', default=False, help='evaluate the model') parser.add_argument('--cycle', type=bool, default=False, metavar='N', help='Whether to use cycle consistency') parser.add_argument('--gaussian_noise', type=bool, default=False, metavar='N', help='Wheter to add gaussian noise') parser.add_argument('--unseen', type=bool, default=False, metavar='N', help='Wheter to test on unseen category') parser.add_argument('--num_points', type=int, default=1024, metavar='N', help='Num of points to use') parser.add_argument('--dataset', type=str, default='modelnet40', choices=['modelnet40'], metavar='N', help='dataset to use') parser.add_argument('--factor', type=float, default=4, metavar='N', help='Divided factor for rotations') parser.add_argument('--model_path', type=str, default='pretrained/dcp_v1.t7', metavar='N', help='Pretrained model path') args = parser.parse_args() torch.backends.cudnn.deterministic = True torch.manual_seed(args.seed) torch.cuda.manual_seed_all(args.seed) np.random.seed(args.seed) boardio = SummaryWriter(log_dir='checkpoints/' + args.exp_name) _init_(args) textio = IOStream('checkpoints/' + args.exp_name + '/run.log') textio.cprint(str(args)) if args.dataset == 'modelnet40': train_loader = DataLoader(ModelNet40( num_points=args.num_points, partition='train', gaussian_noise=args.gaussian_noise, unseen=args.unseen, factor=args.factor), batch_size=args.batch_size, shuffle=True, drop_last=True) test_loader = DataLoader(ModelNet40(num_points=args.num_points, partition='test', gaussian_noise=args.gaussian_noise, unseen=args.unseen, factor=args.factor), batch_size=args.test_batch_size, shuffle=False, drop_last=False) else: raise Exception("not implemented") if args.model == 'dcp': net = DCP(args).cuda() print("Model Parameters") count_parameters(net) if args.eval: if args.model_path is '': model_path = 'checkpoints' + '/' + args.exp_name + '/models/model.best.t7' else: model_path = args.model_path print(model_path) if not os.path.exists(model_path): print("can't find pretrained model") return net.load_state_dict(torch.load(model_path), strict=False) if torch.cuda.device_count() > 1: net = nn.DataParallel(net) print("Let's use", torch.cuda.device_count(), "GPUs!") else: raise Exception('Not implemented') start = torch.cuda.Event(enable_timing=True) end = torch.cuda.Event(enable_timing=True) if args.eval: start.record() test(args, net, test_loader, boardio, textio) end.record() torch.cuda.synchronize() print("Time to test: ", start.elapsed_time(end)) else: train(args, net, train_loader, test_loader, boardio, textio) print('FINISH') boardio.close()
def test(args, io): test_loader = DataLoader(ModelNet40(partition='test', num_points=args.num_points), batch_size=args.test_batch_size, shuffle=False, drop_last=False) device = torch.device("cuda" if args.cuda else "cpu") #Try to load models if args.model == 'pointnet': model = PointNet(args).to(device) elif args.model == 'dgcnn': model = DGCNN(args).to(device) else: model = PaiNet(args).to(device) #raise Exception("Not implemented") print(str(model)) # model = DGCNN(args).to(device) model = nn.DataParallel(model) if os.path.exists('checkpoints/%s/models/model_%s.t7' % (args.exp_name, args.model)): checkpoint_dict = torch.load('./checkpoints/%s/models/model_%s.t7' % (args.exp_name, args.model), map_location=device) model.load_state_dict(checkpoint_dict, strict=True) print("Load model from './checkpoints/%s/models/model_%s_2048.t7 !'" % (args.exp_name, args.model)) model = model.eval() test_acc = 0.0 count = 0.0 test_true = [] test_pred = [] feats = [] for data, label in test_loader: data, label = data.to(device), label.to(device).squeeze() data = data.permute(0, 2, 1) batch_size = data.size()[0] logits, feat = model(data) preds = logits.max(dim=1)[1] test_true.append(label.cpu().numpy()) test_pred.append(preds.detach().cpu().numpy()) feats.append(feat.detach().cpu()) feats = torch.cat(feats) distances = torch.norm((feats - feats[6:7]), dim=1) _, index = torch.topk(-distances, k=5) test = ModelNet40(1024, 'test') points = test.data[index] vis = o3d.visualization.Visualizer() vis.create_window(visible=True) pcd = o3d.geometry.PointCloud() for i, data in enumerate(points): pcd.points = o3d.utility.Vector3dVector(data) data[:] = 0.2 pcd.colors = o3d.utility.Vector3dVector(data) vis.add_geometry(pcd) ctr = vis.get_view_control() ctr.rotate(-260, 100) vis.update_geometry() vis.poll_events() vis.update_renderer() time.sleep(1) vis.capture_screen_image('d{}.png'.format(i)) vis.remove_geometry(pcd) test_true = np.concatenate(test_true) test_pred = np.concatenate(test_pred) test_acc = metrics.accuracy_score(test_true, test_pred) avg_per_class_acc = metrics.balanced_accuracy_score(test_true, test_pred) outstr = 'Test :: test acc: %.6f, test avg acc: %.6f' % (test_acc, avg_per_class_acc) io.cprint(outstr)