(args, G, T, E, iden, itr, lr=2e-2, iter_times=1500, num_seeds=5)
| 44 | |
| 45 | |
| 46 | def inversion(args, G, T, E, iden, itr, lr=2e-2, iter_times=1500, num_seeds=5): |
| 47 | save_img_dir = os.path.join(args.save_dir, 'all_imgs') |
| 48 | success_dir = os.path.join(args.save_dir, 'success_imgs') |
| 49 | os.makedirs(save_img_dir, exist_ok=True) |
| 50 | os.makedirs(success_dir, exist_ok=True) |
| 51 | |
| 52 | bs = iden.shape[0] |
| 53 | iden = iden.view(-1).long().cuda() |
| 54 | |
| 55 | G.eval() |
| 56 | T.eval() |
| 57 | E.eval() |
| 58 | |
| 59 | flag = torch.zeros(bs) |
| 60 | no = torch.zeros(bs) # index for saving all success attack images |
| 61 | |
| 62 | res = [] |
| 63 | res5 = [] |
| 64 | seed_acc = torch.zeros((bs, 5)) |
| 65 | |
| 66 | aug_list = augmentation.container.ImageSequential( |
| 67 | augmentation.RandomResizedCrop((64, 64), scale=(0.8, 1.0), ratio=(1.0, 1.0)), |
| 68 | augmentation.ColorJitter(brightness=0.2, contrast=0.2), |
| 69 | augmentation.RandomHorizontalFlip(), |
| 70 | augmentation.RandomRotation(5), |
| 71 | ) |
| 72 | |
| 73 | for random_seed in range(num_seeds): |
| 74 | tf = time.time() |
| 75 | r_idx = random_seed |
| 76 | |
| 77 | set_random_seed(random_seed) |
| 78 | |
| 79 | z = utils.sample_z( |
| 80 | bs, args.gen_dim_z, device, args.gen_distribution |
| 81 | ) |
| 82 | z.requires_grad = True |
| 83 | |
| 84 | optimizer = torch.optim.Adam([z], lr=lr) |
| 85 | |
| 86 | for i in range(iter_times): |
| 87 | |
| 88 | fake = G(z, iden) |
| 89 | |
| 90 | out1 = T(aug_list(fake))[-1] |
| 91 | out2 = T(aug_list(fake))[-1] |
| 92 | |
| 93 | if z.grad is not None: |
| 94 | z.grad.data.zero_() |
| 95 | |
| 96 | if args.inv_loss_type == 'ce': |
| 97 | inv_loss = L.cross_entropy_loss(out1, iden) + L.cross_entropy_loss(out2, iden) |
| 98 | elif args.inv_loss_type == 'margin': |
| 99 | inv_loss = L.max_margin_loss(out1, iden) + L.max_margin_loss(out2, iden) |
| 100 | elif args.inv_loss_type == 'poincare': |
| 101 | inv_loss = L.poincare_loss(out1, iden) + L.poincare_loss(out2, iden) |
| 102 | |
| 103 | optimizer.zero_grad() |
no test coverage detected