def main(args):
save_folder = '%s_%s' % (args.dataset, args.affix)
log_folder = os.path.join(args.log_root, save_folder)
model_folder = os.path.join(args.model_root, save_folder)
makedirs(log_folder)
makedirs(model_folder)
setattr(args, 'log_folder', log_folder)
setattr(args, 'model_folder', model_folder)
logger = create_logger(log_folder, args.todo, 'info')
print_args(args, logger)
model = Model(i_c=1, n_c=10)
attack = FastGradientSignUntargeted(model,
args.epsilon,
args.alpha,
min_val=0,
max_val=1,
max_iters=args.k,
_type=args.perturbation_type)
if torch.cuda.is_available():
model.cuda()
trainer = Trainer(args, logger, attack)
if args.todo == 'train':
tr_dataset = tv.datasets.MNIST(args.data_root,
train=True,
transform=tv.transforms.ToTensor(),
download=True)
tr_loader = DataLoader(tr_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=4)
# evaluation during training
te_dataset = tv.datasets.MNIST(args.data_root,
train=False,
transform=tv.transforms.ToTensor(),
download=True)
te_loader = DataLoader(te_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=4)
trainer.train(model, tr_loader, te_loader, args.adv_train)
elif args.todo == 'test':
pass
else:
raise NotImplementedError
def main():
log_folder = './results/plots/robustness/revision/luna/bn_adv/'
makedirs(log_folder)
# model = models.resnet50(pretrained=False)
model = resnet50dsbn()
num_classes = 1
model.fc = nn.Linear(model.fc.in_features, num_classes)
if torch.cuda.is_available():
model = model.cuda()
todo = 'test'
if todo == 'test': # set 'valid' fold for knee and luna dataset and set 'test' fold for CXR dataset
eps = np.linspace(0, 0.01, num=21)
for i in range(len(eps)):
epsilon = eps[i]
alpha = epsilon / 2
attack = FastGradientSignUntargeted(model,
epsilon,
alpha,
min_val=0,
max_val=1,
max_iters=10,
_type='linf')
trainer = Trainer(attack, log_folder)
te_dataset = patd.PatchDataset(path_to_images='../luna16/IMGs/',
fold='test',
transform=tv.transforms.Compose([
tv.transforms.Resize(64),
tv.transforms.ToTensor()
]))
te_loader = DataLoader(te_dataset,
batch_size=1,
shuffle=False,
num_workers=1)
checkpoint = torch.load(
'./checkpoint/luna/luna_linf_/checkpoint_best.pth')
model.load_state_dict(checkpoint)
trainer.test(model,
te_loader,
i,
adv_test=True,
use_pseudo_label=False,
if_AUC=True)
# print("std acc: %.4f, adv_acc: %.4f" % (std_acc * 100, adv_acc * 100))
else:
raise NotImplementedError
def main(args):
save_folder = '%s_%s' % (args.dataset, args.affix)
log_folder = os.path.join(args.log_root, save_folder)
model_folder = os.path.join(args.model_root, save_folder)
makedirs(log_folder)
makedirs(model_folder)
setattr(args, 'log_folder', log_folder)
setattr(args, 'model_folder', model_folder)
logger = create_logger(log_folder, args.todo, 'info')
print_args(args, logger)
model = WideResNet(depth=34, num_classes=10, widen_factor=10, dropRate=0.0)
attack = FastGradientSignUntargeted(model,
args.epsilon,
args.alpha,
min_val=0,
max_val=1,
max_iters=args.k,
_type=args.perturbation_type)
if torch.cuda.is_available():
model.cuda()
trainer = Trainer(args, logger, attack)
if args.todo == 'train':
transform_train = tv.transforms.Compose([
tv.transforms.ToTensor(),
tv.transforms.Lambda(lambda x: F.pad(x.unsqueeze(0),
(4,4,4,4), mode='constant', value=0).squeeze()),
tv.transforms.ToPILImage(),
tv.transforms.RandomCrop(32),
tv.transforms.RandomHorizontalFlip(),
tv.transforms.ToTensor(),
])
tr_dataset = tv.datasets.CIFAR10(args.data_root,
train=True,
transform=transform_train,
download=True)
tr_loader = DataLoader(tr_dataset, batch_size=args.batch_size, shuffle=True, num_workers=4)
# evaluation during training
te_dataset = tv.datasets.CIFAR10(args.data_root,
train=False,
transform=tv.transforms.ToTensor(),
download=True)
te_loader = DataLoader(te_dataset, batch_size=args.batch_size, shuffle=False, num_workers=4)
trainer.train(model, tr_loader, te_loader, args.adv_train)
elif args.todo == 'test':
pass
else:
raise NotImplementedError
def main(args):
save_folder = '%s_%s' % (args.dataset, args.affix)
log_folder = os.path.join(args.log_root, save_folder)
model_folder = os.path.join(args.model_root, save_folder)
makedirs(log_folder)
makedirs(model_folder)
setattr(args, 'log_folder', log_folder)
setattr(args, 'model_folder', model_folder)
logger = create_logger(log_folder, args.todo, 'info')
print_args(args, logger)
# model = WideResNet(depth=34, num_classes=10, widen_factor=10, dropRate=0.0)
model = models.resnet50(pretrained=args.pretrain)
num_classes=8
# model.classifier = nn.Linear(model.classifier.in_features, num_classes)
model.fc = nn.Linear(model.fc.in_features, num_classes)
attack = FastGradientSignUntargeted(model,
args.epsilon,
args.alpha,
min_val=0,
max_val=1,
max_iters=args.k,
_type=args.perturbation_type)
if torch.cuda.is_available():
model.cuda()
# model = nn.DataParallel(model).cuda()
trainer = Trainer(args, logger, attack)
if args.todo == 'train':
transform_train = tv.transforms.Compose([
tv.transforms.Resize(256),
tv.transforms.ToTensor(),
tv.transforms.Lambda(lambda x: F.pad(x.unsqueeze(0),
(4*6,4*6,4*6,4*6), mode='constant', value=0).squeeze()),
tv.transforms.ToPILImage(),
tv.transforms.RandomHorizontalFlip(),
tv.transforms.ColorJitter(brightness=0.3, contrast=0.3,
saturation=0.3, hue=0.3),
# tv.transforms.RandomRotation(25),
tv.transforms.RandomAffine(25, translate=(0.2, 0.2),
scale=(0.8,1.2), shear=10),
tv.transforms.RandomCrop(256),
tv.transforms.ToTensor(),
AddGaussianNoise(0.5, args.epsilon)
])
tr_dataset = patd.PatchDataset(path_to_images=args.data_root,
fold='train',
sample=args.subsample,
transform=transform_train)
tr_loader = DataLoader(tr_dataset, batch_size=args.batch_size, shuffle=True, num_workers=24)
# evaluation during training
transform_test = tv.transforms.Compose([
tv.transforms.Resize(256),
# tv.transforms.CenterCrop(224),
tv.transforms.ToTensor(),
# tv.transforms.Normalize(mean, std)
])
te_dataset = patd.PatchDataset(path_to_images=args.data_root,
fold='valid',
transform=transform_test)
te_loader = DataLoader(te_dataset, batch_size=args.batch_size, shuffle=False, num_workers=24)
trainer.train(model, tr_loader, te_loader, args.adv_train)
elif args.todo == 'test':
te_dataset = patd.PatchDataset(path_to_images=args.data_root,
fold='test',
transform=tv.transforms.Compose([
tv.transforms.Resize(256),
tv.transforms.ToTensor(),
]))
te_loader = DataLoader(te_dataset, batch_size=1, shuffle=False, num_workers=1)
checkpoint = torch.load(args.load_checkpoint)
model.load_state_dict(checkpoint)
std_acc, adv_acc = trainer.test(model, te_loader, adv_test=True, use_pseudo_label=False, if_AUC=True)
print("std acc: %.4f, adv_acc: %.4f" % (std_acc * 100, adv_acc * 100))
else:
raise NotImplementedError
adv_list = []
pred_list = []
with torch.no_grad():
model = WideResNet(depth=34, num_classes=10, widen_factor=10, dropRate=0.0)
load_model(model, args.load_checkpoint)
if torch.cuda.is_available():
model.cuda()
attack = FastGradientSignUntargeted(model,
max_epsilon,
args.alpha,
min_val=0,
max_val=1,
max_iters=args.k,
_type=perturbation_type)
adv_data = attack.perturb(data, label, 'mean', False)
output = model(adv_data, _eval=True)
pred = torch.max(output, dim=1)[1]
adv_list.append(adv_data.cpu().numpy().squeeze() * 255.0) # (N, 28, 28)
pred_list.append(pred.cpu().numpy())
data = data.cpu().numpy().squeeze() # (N, 28, 28)
data *= 255.0
label = label.cpu().numpy()
def main(args):
save_folder = '%s_%s' % (args.dataset, args.affix)
log_folder = os.path.join(args.log_root, save_folder)
model_folder = os.path.join(args.model_root, save_folder)
makedirs(log_folder)
makedirs(model_folder)
setattr(args, 'log_folder', log_folder)
setattr(args, 'model_folder', model_folder)
logger = create_logger(log_folder, args.todo, 'info')
print_args(args, logger)
# Using a WideResNet model
model = WideResNet(depth=34, num_classes=10, widen_factor=1, dropRate=0.0)
flop, param = get_model_infos(model, (1, 3, 32, 32))
logger.info('Model Info: FLOP = {:.2f} M, Params = {:.2f} MB'.format(
flop, param))
# Configuring the train attack mode
if args.adv_train_mode == 'FGSM':
train_attack = FastGradientSignUntargeted(model,
args.epsilon,
args.alpha,
min_val=0,
max_val=1,
max_iters=args.k,
_type=args.perturbation_type,
logger=logger)
elif args.adv_train_mode == 'CW':
mean = [0]
std = [1]
inputs_box = (min((0 - m) / s for m, s in zip(mean, std)),
max((1 - m) / s for m, s in zip(mean, std)))
train_attack = carlini_wagner_L2.L2Adversary(targeted=False,
confidence=0.0,
search_steps=10,
optimizer_lr=5e-4,
logger=logger)
# Configuring the test attack mode
if args.adv_test_mode == 'FGSM':
test_attack = FastGradientSignUntargeted(model,
args.epsilon,
args.alpha,
min_val=0,
max_val=1,
max_iters=args.k,
_type=args.perturbation_type,
logger=logger)
elif args.adv_test_mode == 'CW':
mean = [0]
std = [1]
inputs_box = (min((0 - m) / s for m, s in zip(mean, std)),
max((1 - m) / s for m, s in zip(mean, std)))
test_attack = carlini_wagner_L2.L2Adversary(targeted=False,
confidence=0.0,
search_steps=10,
optimizer_lr=5e-4,
logger=logger)
if torch.cuda.is_available():
model.cuda()
trainer = Trainer(args, logger, train_attack, test_attack)
if args.todo == 'train':
transform_train = tv.transforms.Compose([
tv.transforms.ToTensor(),
tv.transforms.Lambda(lambda x: F.pad(
x.unsqueeze(0),
(4, 4, 4, 4), mode='constant', value=0).squeeze()),
tv.transforms.ToPILImage(),
tv.transforms.RandomCrop(32),
tv.transforms.RandomHorizontalFlip(),
tv.transforms.ToTensor(),
])
tr_dataset = tv.datasets.CIFAR10(args.data_root,
train=True,
transform=transform_train,
download=True)
tr_loader = DataLoader(tr_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=4)
# evaluation during training
te_dataset = tv.datasets.CIFAR10(args.data_root,
train=False,
transform=tv.transforms.ToTensor(),
download=True)
te_loader = DataLoader(te_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=4)
trainer.train(model, tr_loader, te_loader, args.adv_train)
elif args.todo == 'test':
pass
else:
raise NotImplementedError
def main():
""" main function
"""
### header
parser = argparse.ArgumentParser()
# path
parser.add_argument('--root-path', default=CFG.root_path, help="root path")
parser.add_argument('--log-path', default=CFG.log_path, help="log path")
parser.add_argument('--model-path',
default=CFG.model_path,
help="model path")
parser.add_argument('--pretrained-path', help='pretrained path')
# image
parser.add_argument('--transform-version',
default=0,
type=int,
help="image transform version ex) 0, 1, 2 ...")
parser.add_argument('--image-size',
default=64,
type=int,
help="image size(64)")
# model
parser.add_argument('--model-name',
default=CFG.model_name,
help=f"model name({CFG.model_name})")
parser.add_argument('--backbone-name',
default=CFG.backbone_name,
help=f"backbone name({CFG.backbone_name})")
# learning
parser.add_argument('--batch-size',
default=CFG.batch_size,
type=int,
help=f"batch size({CFG.batch_size})")
parser.add_argument('--learning-rate',
default=CFG.learning_rate,
type=float,
help=f"learning rate({CFG.learning_rate})")
parser.add_argument('--num-epochs',
default=CFG.num_epochs,
type=int,
help=f"number of epochs({CFG.num_epochs})")
# etc
parser.add_argument("--seed",
default=CFG.seed,
type=int,
help=f"seed({CFG.seed})")
parser.add_argument("--workers",
default=CFG.workers,
type=int,
help=f"number of workers({CFG.workers})")
parser.add_argument("--debug", action="store_true", help="debug mode")
parser.add_argument("--val-fold",
default=CFG.val_fold,
choices=[list(range(0, CFG.n_folds))],
help=f"fold number for validation({CFG.val_fold})")
args = parser.parse_args()
# path
CFG.root_path = args.root_path
CFG.model_path = args.model_path
CFG.log_path = args.log_path
CFG.pretrained_path = args.pretrained_path
# image
CFG.transform_version = args.transform_version
CFG.image_size = args.image_size
# model
CFG.model_name = args.model_name
CFG.backbone_name = args.backbone_name
# learning
CFG.batch_size = args.batch_size
CFG.learning_rate = args.learning_rate
CFG.num_epochs = args.num_epochs
# etc
CFG.seed = args.seed
CFG.workers = args.workers
CFG.debug = args.debug
# get device
CFG.device = get_device()
# get version
_, version, _ = sys.argv[0].split('/')
CFG.version = version
# update log path
if not CFG.debug:
CFG.log_path = os.path.join(CFG.log_path, CFG.version)
os.makedirs(CFG.log_path, exist_ok=True)
CFG.log_path = os.path.join(
CFG.log_path, f'exp_{get_exp_id(CFG.log_path, prefix="exp_")}')
os.makedirs(CFG.log_path, exist_ok=True)
else:
CFG.log_path = os.path.join(CFG.log_path, "debug")
os.makedirs(CFG.log_path, exist_ok=True)
CFG.log_path = os.path.join(CFG.log_path, "debug")
os.makedirs(CFG.log_path, exist_ok=True)
# update model path
if not CFG.debug:
CFG.model_path = os.path.join(CFG.model_path, version)
os.makedirs(CFG.model_path, exist_ok=True)
CFG.model_path = os.path.join(
CFG.model_path, f'exp_{get_exp_id(CFG.model_path, prefix="exp_")}')
os.makedirs(CFG.model_path, exist_ok=True)
else:
CFG.model_path = os.path.join(CFG.model_path, "debug")
os.makedirs(CFG.model_path, exist_ok=True)
CFG.model_path = os.path.join(CFG.model_path, "debug")
os.makedirs(CFG.model_path, exist_ok=True)
pprint({k: v for k, v in dict(CFG.__dict__).items() if '__' not in k})
json.dump({k: v
for k, v in dict(CFG.__dict__).items() if '__' not in k},
open(os.path.join(CFG.log_path, 'CFG.json'), "w"))
print()
### Seed all
seed_everything(CFG.seed)
### Data Related
# load raw data
print("Load Raw Data")
train_df, test_df, ss_df = load_data(CFG)
# preprocess data
print("Preprocess Data")
train_df = preprocess_data(CFG, train_df)
# split data
print("Split Data")
train_df, valid_df = split_data(CFG, train_df)
# get transform
print("Get Transforms")
train_transforms, test_transforms = get_transform(CFG)
# get dataset
print("Get Dataset")
trn_data = DFDDataset(CFG, train_df, train_transforms)
val_data = DFDDataset(CFG, valid_df, test_transforms)
### Model related
# get learner
learner = Learner(CFG)
learner.name = f"model.fold_{CFG.val_fold}"
if CFG.pretrained_path:
print("Load Pretrained Model")
print(f"... Pretrained Info - {CFG.pretrained_path}")
learner.load(CFG.pretrained_path, f"model_state_dict")
model = learner.best_model.to(CFG.deivce)
else:
print(f"Load Model")
model = get_model(CFG)
model = model.to(CFG.device)
if torch.cuda.device_count() > 1:
model = nn.DataParallel(model)
# get attacker
attack = FastGradientSignUntargeted(CFG,
model,
0.0157,
0.00784,
min_val=0,
max_val=1,
max_iters=10,
_type="linf")
# get optimizer
optimizer = optim.Adam(model.parameters(), lr=CFG.learning_rate)
# get scheduler
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer,
mode='min',
patience=2,
verbose=False,
factor=0.5)
### train related
# train model
learner.train(trn_data, val_data, model, optimizer, scheduler, attack)
def main(args):
save_folder = '%s_%s' % (args.dataset, args.affix)
log_folder = os.path.join(args.log_root, save_folder)
model_folder = os.path.join(args.model_root, save_folder)
makedirs(log_folder)
makedirs(model_folder)
setattr(args, 'log_folder', log_folder)
setattr(args, 'model_folder', model_folder)
logger = create_logger(log_folder, args.todo, 'info')
print_args(args, logger)
#return
# model = WideResNet(depth=34, num_classes=10, widen_factor=10, dropRate=0.0)
print('create model')
model = ResNet34()
print('done...')
# file_name = os.path.join(args.model_folder, 'checkpoint.pth')
# torch.save(model.state_dict(), file_name)
# load_model(model, file_name, args)
print('create attack')
attack = FastGradientSignUntargeted(model,
args.epsilon,
args.alpha,
min_val=0,
max_val=1,
max_iters=args.k,
_type=args.perturbation_type)
print('done...')
print('create trainer')
trainer = Trainer(args, logger, attack)
print('done...')
if args.todo == 'train':
print('start train')
if torch.cuda.is_available():
model.cuda()
transform_train = tv.transforms.Compose([
tv.transforms.ToTensor(),
tv.transforms.Lambda(lambda x: F.pad(x.unsqueeze(0), (4, 4, 4, 4), mode='constant', value=0).squeeze()),
tv.transforms.ToPILImage(),
tv.transforms.RandomCrop(32),
tv.transforms.RandomHorizontalFlip(),
tv.transforms.ToTensor(),
])
print('get dataset')
print('dataset_root: ', args.data_root)
tr_dataset = tv.datasets.CIFAR10(args.data_root, train=True, transform=transform_train, download=True)
tr_loader = DataLoader(tr_dataset, batch_size=args.batch_size, shuffle=True, num_workers=4)
# evaluation during training
te_dataset = tv.datasets.CIFAR10(args.data_root, train=False, transform=tv.transforms.ToTensor(), download=True)
te_loader = DataLoader(te_dataset, batch_size=args.batch_size, shuffle=False, num_workers=4)
trainer.train(model, tr_loader, te_loader, args.adv_train)
elif args.todo == 'valid':
load_model(model, args.load_checkpoint, args)
if torch.cuda.is_available():
model.cuda()
te_dataset = tv.datasets.CIFAR10(args.data_root, train=False, transform=tv.transforms.ToTensor(), download=True)
te_loader = DataLoader(te_dataset, batch_size=args.batch_size, shuffle=False, num_workers=4)
test_acc, adv_acc = trainer.test(model, te_loader, adv_test=False)
print('Test accuracy is %.3f' % test_acc)
else:
raise NotImplementedError
def fit(train_loader,
val_loader,
model,
loss_fn,
optimizer,
scheduler,
n_epochs,
cuda,
log_interval,
model_save_dir,
metrics=[],
start_epoch=0,
criterion=None,
domain_cls=False,
unsup_da=False,
adv_train=False,
adv_epsilon=None,
adv_alph=None,
adv_iter=None,
eval_train_dataset=None,
eval_test_dataset=None):
best_val_loss = None
start_time = time.time()
# for epoch in range(0, start_epoch):
# scheduler.step()
writer = SummaryWriter('runs/{}'.format(model_save_dir))
for epoch in range(start_epoch, n_epochs):
# Train stage
if adv_train:
attack = FastGradientSignUntargeted(model,
adv_epsilon,
adv_alph,
min_val=-2.5,
max_val=3.0,
max_iters=adv_iter,
_type='linf') # l2
else:
attack = None
if not domain_cls:
train_loss, metrics = train_epoch(train_loader, model, attack,
loss_fn, optimizer, cuda, writer,
epoch, log_interval, metrics,
adv_train)
message = 'Epoch: {}/{}. Train set: Average loss: {:.4f} Elapsed time: {}s'\
.format(epoch + 1, n_epochs, train_loss, int(time.time() - start_time))
else:
train_loss, train_loss_sim, train_loss_domain_cls, metrics = train_domain_classifier_epoch(train_loader,\
model, attack, loss_fn, criterion, optimizer, cuda, writer, epoch, log_interval, metrics, unsup_da)
message = 'Epoch: {}/{}. Train set: Average loss: {:.4f} loss-sim: {:.4f} loss-domain-cls: {:.4f} Elapsed time: {}s'\
.format(epoch + 1, n_epochs, train_loss, train_loss_sim, train_loss_domain_cls, int(time.time() - start_time))
for metric in metrics:
message += '\t{}: {:.4f}'.format(metric.name(), metric.value())
summ_step = (epoch + 1) * len(train_loader) - 1
# Test stage
if not domain_cls:
val_loss, val_loss_sim, val_loss_mix_div, metrics = test_epoch(
val_loader, model, loss_fn, cuda, metrics)
val_loss /= len(val_loader)
val_loss_sim /= len(val_loader)
val_loss_mix_div /= len(val_loader)
writer.add_scalars('Loss/total', {'validation': val_loss},
summ_step)
writer.add_scalars('Loss/similarity', {'validation': val_loss_sim},
summ_step)
writer.add_scalars('Loss/miture-divergence',
{'validation': val_loss_mix_div}, summ_step)
message += '\nEpoch: {}/{}. Validation set: Average loss: {:.4f} loss-sim: {:.4f} loss-mixture-div: {:.4f}'\
.format(epoch + 1, n_epochs, val_loss, val_loss_sim, val_loss_mix_div)
else:
val_loss, val_loss_sim, val_loss_domain_cls, metrics = test_domain_classifier_epoch(
val_loader, model, loss_fn, criterion, cuda, metrics, unsup_da)
val_loss /= len(val_loader)
val_loss_sim /= len(val_loader)
val_loss_domain_cls /= len(val_loader)
writer.add_scalars('Loss/total', {'validation': val_loss},
summ_step)
writer.add_scalars('Loss/similarity', {'validation': val_loss_sim},
summ_step)
writer.add_scalars('Loss/domain_clf',
{'validation': val_loss_domain_cls}, summ_step)
message += '\nEpoch: {}/{}. Validation set: Average loss: {:.4f} loss-sim: {:.4f} loss-domain-cls: {:.4f}'\
.format(epoch + 1, n_epochs, val_loss, val_loss_sim, val_loss_domain_cls)
if best_val_loss is None or best_val_loss > val_loss:
best_val_loss = val_loss
torch.save(
model.module.state_dict(),
os.path.join(model_save_dir, '{}.pth'.format(str(epoch).zfill(5))))
for metric in metrics:
message += '\t{}: {}'.format(metric.name(), metric.value())
writer.add_scalars('Metric/{}'.format(metric.name()),
{'validation': metric.value()}, summ_step)
scheduler.step(val_loss)
print(message)
def main(args):
save_folder = '%s_%s' % (args.dataset, args.affix)
log_folder = os.path.join(args.log_root, save_folder)
model_folder = os.path.join(args.model_root, save_folder)
makedirs(log_folder)
makedirs(model_folder)
setattr(args, 'log_folder', log_folder)
setattr(args, 'model_folder', model_folder)
logger = create_logger(log_folder, args.todo, 'info')
print_args(args, logger)
model = Classifier(10)
attack = FastGradientSignUntargeted(model,
args.epsilon,
args.alpha,
min_val=-1,
max_val=1,
max_iters=args.k,
_type=args.perturbation_type)
trainer = Trainer(args, logger, attack)
if args.todo == 'train':
if torch.cuda.is_available():
model.cuda()
tr_dataset = tv.datasets.MNIST(args.data_root,
train=True,
transform=tv.transforms.Compose([
tv.transforms.Resize(args.img_size),
tv.transforms.ToTensor(),
tv.transforms.Normalize([0.5],
[0.5])
]),
download=True)
tr_loader = DataLoader(tr_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=4)
# evaluation during training
te_dataset = tv.datasets.MNIST(args.data_root,
train=False,
transform=tv.transforms.Compose([
tv.transforms.Resize(args.img_size),
tv.transforms.ToTensor(),
tv.transforms.Normalize([0.5],
[0.5])
]),
download=True)
te_loader = DataLoader(te_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=4)
trainer.train(model, tr_loader, te_loader, args.adv_train)
elif args.todo == 'valid':
load_model(model, args.load_checkpoint, args)
if torch.cuda.is_available():
model.cuda()
te_dataset = tv.datasets.MNIST(args.data_root,
train=False,
transform=tv.transforms.Compose([
tv.transforms.Resize(args.img_size),
tv.transforms.ToTensor(),
tv.transforms.Normalize([0.5],
[0.5])
]),
download=True)
te_loader = DataLoader(te_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=4)
test_acc, adv_acc = trainer.test(model, te_loader, adv_test=False)
print('Test accuracy is %.3f' % test_acc)
else:
raise NotImplementedError