def gn_helper(planes):
return nn.GroupNorm(args.group_norm, planes)
norm_layer = gn_helper
net = resnet18(num_classes = 10, norm_layer=norm_layer).to(device)
net = torch.nn.DataParallel(net)
print('Resuming from %s...' %(args.resume))
ckpt = torch.load('%s/best.pth' %(args.resume))
net.load_state_dict(ckpt['net'])
criterion = nn.CrossEntropyLoss().to(device)
optimizer = optim.SGD(net.parameters(), lr=args.lr)
_, teloader = prepare_test_data(args)
print("Random Attack...")
for i in range(args.epochs):
r = np.random.normal(0.4914, 0.2023, (32, 32))
g = np.random.normal(0.4822, 0.1994, (32, 32))
b = np.random.normal(0.4465, 0.2010, (32, 32))
img = np.stack([r, g, b], axis=0)
tensor = torch.from_numpy(img).float()
tensor = torch.clamp(tensor, 0, 1)
adapt_single_tensor(net, tensor, optimizer, criterion, args.niter, args.batch_size)
if i % 50 == 49:
print("%d%%" % ((i + 1) * 100 / 5000))
err_cls, correct_per_cls, total_per_cls = test(teloader, net, verbose=True, print_freq=0)
print("Epoch %d Test error: %.3f" % (i, err_cls))
best = 1
start_epoch = args.start_epoch
if args.resume is not None:
print('Resuming from checkpoint..')
ckpt = torch.load('%s/ckpt.pth' % (args.resume))
net.load_state_dict(ckpt['net'])
optimizer.load_state_dict(ckpt['optimizer'])
start_epoch = ckpt['epoch'] + 1
loss = torch.load('%s/loss.pth' % (args.resume))
all_err_cls, all_err_ssh = loss
for epoch in range(start_epoch, args.epochs + 1):
adjust_learning_rate(optimizer, epoch, args)
train(trloader, epoch)
err_cls = test(teloader, net)
if False:
teloader.dataset.switch_mode(False, True)
err_ssh = test(teloader, ssh)
else:
err_ssh = 0
all_err_cls.append(err_cls)
all_err_ssh.append(err_ssh)
torch.save((all_err_cls, all_err_ssh), args.outf + '/loss.pth')
plot_epochs(all_err_cls, all_err_ssh, args.outf + '/loss.pdf')
state = {
'epoch': epoch,
'args': args,
'err_cls': err_cls,
parser.add_argument('--resume', default=None)
parser.add_argument('--outf', default='.')
parser.add_argument('--none', action='store_true')
args = parser.parse_args()
my_makedir(args.outf)
import torch.backends.cudnn as cudnn
cudnn.benchmark = True
net, ext, head, ssh = build_model(args)
teset, teloader = prepare_test_data(args)
print('Resuming from %s...' % (args.resume))
ckpt = torch.load('%s/ckpt.pth' % (args.resume))
net.load_state_dict(ckpt['net'])
teloader.dataset.switch_mode(True, False)
cls_initial, cls_correct, cls_losses = test(teloader, net, verbose=True)
print('Old test error cls %.2f' % (ckpt['err_cls'] * 100))
print('New test error cls %.2f' % (cls_initial * 100))
if args.none:
rdict = {
'cls_initial': cls_initial,
'cls_correct': cls_correct,
'cls_losses': cls_losses
}
torch.save(rdict,
args.outf + '/%s_%d_none.pth' % (args.corruption, args.level))
quit()
print('Old test error ssh %.2f' % (ckpt['err_ssh'] * 100))