if args.batch_size == -1:
args.batch_size = BATCH_SIZE[args.model]
set_seed(1)
print(args)
dist.init_process_group(backend='nccl')
torch.cuda.set_device(args.local_rank)
cudnn.benchmark = True
# build model
if args.model.lower() == 'dgcnn':
model = DGCNN(args.emb_dims, args.k, output_channels=40)
elif args.model.lower() == 'pointnet':
model = PointNetCls(k=40, feature_transform=args.feature_transform)
elif args.model.lower() == 'pointnet2':
model = PointNet2ClsSsg(num_classes=40)
elif args.model.lower() == 'pointconv':
model = PointConvDensityClsSsg(num_classes=40)
else:
print('Model not recognized')
exit(-1)
# load model weight
state_dict = torch.load(
BEST_WEIGHTS[args.model], map_location='cpu')
print('Loading weight {}'.format(BEST_WEIGHTS[args.model]))
try:
model.load_state_dict(state_dict)
except RuntimeError:
# eliminate 'module.' in keys
state_dict = {k[7:]: v for k, v in state_dict.items()}
def main():
# build model
model = PointNet2ClsSsg(num_classes=40)
model = nn.DataParallel(model).cuda()
model.eval()
# load model weight
print('Loading weight {}'.format(BEST_WEIGHTS['pointnet2']))
state_dict = torch.load(BEST_WEIGHTS['pointnet2'])
try:
model.load_state_dict(state_dict)
except RuntimeError:
model.module.load_state_dict(state_dict)
# load dataset
test_set = ModelNet40Attack('data/attack_data.npz', num_points=1024,
normalize=True)
test_loader = DataLoader(test_set, batch_size=1,
shuffle=True, num_workers=4,
pin_memory=True, drop_last=False)
ti = 412
data_iter = iter(test_loader)
for i in range(ti):
data = next(data_iter)
total_num = 0
success_num = 0
at_success_num = 0
i = 0
all_adv = []
all_real_label = []
all_target_label = []
for x, label, target in tqdm(test_loader):
# x, label, target = data
x, label, target = x.cuda(), label.long().cuda(), target.long().cuda()
x = x.transpose(2, 1).contiguous()
x.requires_grad = True
rx = x.clone()
x_pred = model(x)
R1 = LRP_scores(model, x, None, label, label)
x_pred = torch.argmax(x_pred, dim=-1)
# if x_pred != label:
# # all_adv.append(x_adv.transpose(1, 2).contiguous().detach().cpu().numpy())
# # all_real_label.append(label.detach().cpu().numpy())
# # all_target_label.append(target.detach().cpu().numpy())
# continue
total_num += x.shape[0]
x_adv = AOA_Attack(model, x, label, target)
pred, _, _ = model(x_adv)
pred = torch.argmax(pred, dim=-1)
success_num += (pred != label).sum().cpu().item()
logits_at = model_at(x_adv)
pred_at = torch.argmax(logits_at, dim=-1)
at_success_num += (pred_at != label).sum().cpu().item()
i += 1
if i % 20 == 0:
print("current attack success rate is", success_num / total_num)
print("current pointnet++ attack success rate is", at_success_num / total_num)
all_adv.append(x_adv.transpose(1, 2).contiguous().detach().cpu().numpy())
all_real_label.append(label.detach().cpu().numpy())
all_target_label.append(target.detach().cpu().numpy())
# if i % 20 == 0:
# break
# R0 = LRP_scores(model, x_adv, label, label)
# R1 = LRP_scores(model, x, label, label)
# utils.pc_heatmap(x_adv.transpose(2, 1)[0], R0[0].sum(-2).unsqueeze(-1))
attacked_data = np.concatenate(all_adv, axis=0) # [num_data, K, 3]
real_label = np.concatenate(all_real_label, axis=0) # [num_data]
target_label = np.concatenate(all_target_label, axis=0) # [num_data]
# save results
save_path = 'attack/results/{}_{}/AOA/{}'. \
format('mn40', 1024, 'pointnet')
if not os.path.exists(save_path):
os.makedirs(save_path)
save_name = '{}-budget_{}-iter_{}' \
'-success_{:.4f}-rank_{}.npz'. \
format('aoa', 0.5,
200, success_num/total_num, 0)
np.savez(os.path.join(save_path, save_name),
test_pc=attacked_data.astype(np.float32),
test_label=real_label.astype(np.uint8),
target_label=target_label.astype(np.uint8))
print("total attack success rate is", success_num/total_num)