testset_cfg = cfg.dataset.test
print(cfg)
output_dir = cfg.classifier.path
os.makedirs(output_dir, exist_ok=True)
checkpoint_path = os.path.join(output_dir, 'classifier.pt')
output_filename = os.path.join(output_dir, 'metrics.txt')
logging.basicConfig(level=logging.INFO, filename=output_filename, filemode='w')
# set device and random seed
set_device(cfg)
set_random_seed(cfg)
cudnn.benchmark = True
# set classifier
net = get_classifier(cfg, cfg.classifier)
state_dict = torch.load(checkpoint_path)['state_dict']
net.load_state_dict(state_dict)
for p in net.parameters():
p.requires_grad_(False)
net.cuda()
net.eval()
# load GAN
gan = StyleGANGeneratorModel()
state_dict = torch.load('./stylegan/pretrain/stylegan_imagenet.pth')
var_name = 'truncation.truncation'
state_dict[var_name] = gan.state_dict()[var_name]
gan.load_state_dict(state_dict)
# reduce memory consumption
gan = gan.synthesis
train=False,
transform=transform.default)
except:
testset = dataset(root=dataset_cfg.path, transform=transform.default)
testloader = DataLoader(testset,
batch_size=dataset_cfg.batch_size,
num_workers=0,
shuffle=False)
progress_bar = tqdm(testloader)
n_batch = len(testloader)
chunk_size = n_batch // args.world_size
start_ind = args.rank * chunk_size
end_ind = (args.rank + 1) * chunk_size
# load classifier
predict = get_classifier(cfg, cfg.classifier)
state_dict = torch.load(cfg.classifier.ckpt)
try:
predict.load_state_dict(state_dict)
except:
predict.load_state_dict(state_dict["state_dict"])
for p in predict.parameters():
p.requires_grad_(False)
predict = torch.nn.Sequential(transform.classifier_preprocess_layer,
predict).cuda()
predict.eval()
# create attacker
attacker = PGDAttack(predict=predict,
eps=args.eps / 255.0,
eps_iter=1 / 255.0,
