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 = models[args.model]().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 choose_dataset(args): if args.dataset == 'modelnet40': transform = None transform = transforms.Compose( [RotatePointCloud(), JitterPointCloud()]) dataset = ModelNet40(args, transform) return dataset
def choose_dataset(args): if args.dataset == 'modelnet40': transform = transforms.Compose( [RotatePointCloud(), JitterPointCloud()]) dataset = ModelNet40(args, transform) elif args.dataset == 'modelnet_normal_resampled': if args.normal: transform = transforms.Compose( [RotatePointCloud_Normal, JitterPointCloud]) else: transform = transforms.Compose( [RotatePointCloud(), JitterPointCloud()]) dataset = ModelNet_Normal_Resampled(args, transform) return dataset
# load model weight if args.model_path: model.load_state_dict(torch.load(args.model_path)) else: model.load_state_dict(torch.load(BEST_WEIGHTS[args.model])) # prepare data if args.mode == 'target': test_set = ModelNet40Attack(args.data_root, num_points=args.num_points, normalize=args.normalize_pc) else: test_set = ModelNet40(args.data_root, num_points=args.num_points, normalize=args.normalize_pc, partition='test', augmentation=False) test_loader = DataLoader(test_set, batch_size=args.batch_size, shuffle=False, num_workers=8, pin_memory=True, drop_last=False) # test if args.mode == 'normal': test_normal() else: test_target()
def train_and_eval(): device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # Dataset train_dataset = ModelNet40(n=N, split='train') test_dataset = ModelNet40(n=N, split='test') train_data_loader = Data.DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=NUM_WORKERS) test_data_loader = Data.DataLoader(test_dataset, batch_size=BATCH_SIZE, num_workers=NUM_WORKERS) # Network Model categories = len(train_dataset.code2category) model = ClassificationPointNet(categories, input_transform=INPUT_TRANSFORM, feature_transform=FEATURE_TRANSFORM) model.to(device) # Optimizer optimizer = optim.Adam(model.parameters(), lr=0.001, betas=(0.9, 0.999)) scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.5) train_batch_num = math.ceil(len(train_dataset) / BATCH_SIZE) test_batch_num = math.ceil(len(test_dataset) / BATCH_SIZE) max_overall_accuracy, counterpart_avg_class_accuracy = 0.0, 0.0 counterpart_overall_accuracy, max_avg_class_accuracy = 0.0, 0.0 epoch1, epoch2 = 0, 0 # records the epoch of best performance for overall and avg. class accuracy for epoch in range(1, EPOCHS + 1): for i, (point_clouds, labels) in enumerate(train_data_loader): # a batch of data point_clouds = point_clouds.transpose(2, 1) labels = labels.view(-1) point_clouds, labels = point_clouds.to(device), labels.to(device) # forward-propagation optimizer.zero_grad() model.train() preds, transformer1, transformer2 = model(point_clouds) loss = F.nll_loss(preds, labels) if FEATURE_TRANSFORM: loss += feature_transform_regularization(transformer2) * 0.001 # back-propagation and gradient descent loss.backward() optimizer.step() # metrics on train predictions = preds.data.max(1)[1] comparison = predictions.data.eq(labels).cpu() correct = torch.sum(comparison).item() if (i + 1) % 50 == 0: print( "Training epoch {} iteration {}/{} ==> Loss: {}, Accuracy: {}" .format(epoch, i + 1, train_batch_num, loss.item(), correct / len(labels))) scheduler.step() test_total_correct = 0 correct_nums_per_cat = np.zeros((categories, )) with torch.no_grad(): for i, (data, target) in enumerate(test_data_loader): data = data.transpose(2, 1) target = target.view(-1) data, target = data.to(device), target.to(device) model.eval() preds, _, _ = model(data) predictions = preds.data.max(1)[1] comparison = predictions.data.eq(target).cpu() correct = torch.sum(comparison).item() test_total_correct += correct for index, cat in enumerate(target): correct_nums_per_cat[cat] += comparison[index] overall_accuracy = test_total_correct / len(test_dataset) avg_class_accuracy = np.mean(correct_nums_per_cat / np.array(test_dataset.categories_nums)) print("Evaluation on testing dataset ==> Overall Accuracy: {}, " "Avg. class Accuracy: {}".format(overall_accuracy, avg_class_accuracy)) if overall_accuracy > max_overall_accuracy: max_overall_accuracy = overall_accuracy counterpart_avg_class_accuracy = avg_class_accuracy torch.save(model.state_dict(), "max_overall_accuracy_epoch_{}.pth".format(epoch)) epoch1 = epoch if avg_class_accuracy > max_avg_class_accuracy: max_avg_class_accuracy = avg_class_accuracy counterpart_overall_accuracy = overall_accuracy torch.save(model.state_dict(), "max_avg_accuracy_epoch_{}.pth".format(epoch)) epoch2 = epoch print( "Max overall accuracy {} in epoch {}, and the avg class accuracy is {}" .format(max_overall_accuracy, epoch1, counterpart_avg_class_accuracy)) print("Max avg accuracy {} in epoch {}, and the overall accuracy is {}". format(max_avg_class_accuracy, epoch2, counterpart_overall_accuracy))
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 = models[args.model]().to(device) 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): train_loss = 0.0 count = 0.0 # numbers of data model.train() train_pred = [] train_true = [] idx = 0 # iterations 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) scheduler.step()
model = PointNet2ClsSsg(num_classes=40) elif args.model.lower() == 'pointconv': model = PointConvDensityClsSsg(num_classes=40) else: print('Model not recognized') exit(-1) model = nn.DataParallel(model).cuda() # use Adam optimizer, cosine lr decay opt = optim.Adam(model.parameters(), lr=args.lr, weight_decay=1e-4) scheduler = CosineAnnealingLR(opt, T_max=args.epochs, eta_min=1e-5) # prepare data train_set = ModelNet40(args.data_root, num_points=args.num_points, normalize=True, partition='train') train_loader = DataLoader(train_set, batch_size=args.batch_size, shuffle=True, num_workers=8, pin_memory=True, drop_last=True) test_set = ModelNet40(args.data_root, num_points=args.num_points, normalize=True, partition='test') test_loader = DataLoader(test_set, batch_size=args.batch_size * 2, shuffle=False,