# Azure's Mongo db Configuration
uri = "mongodb://*****:*****@smo"\
"othing-db.documents.azure.com:10255/?ssl=true&replicaSet=globaldb"
client = pymongo.MongoClient(uri)
database = client['certification'] # Make sure that this database exists on Azure
results_db = database[args.dataset] # Make sure that this collection exists in the db
#######################################################################################
if __name__ == "__main__":
# load the base classifier
if args.base_classifier in IMAGENET_CLASSIFIERS:
assert args.dataset == 'imagenet'
# loading pretrained imagenet architectures
base_classifier = get_architecture(args.base_classifier ,args.dataset, pytorch_pretrained=True)
else:
checkpoint = torch.load(args.base_classifier)
base_classifier = get_architecture(checkpoint['arch'], args.dataset)
base_classifier.load_state_dict(checkpoint['state_dict'])
if args.denoiser != '':
checkpoint = torch.load(args.denoiser)
if "off-the-shelf-denoiser" in args.denoiser:
denoiser = get_architecture('orig_dncnn', args.dataset)
denoiser.load_state_dict(checkpoint)
else:
denoiser = get_architecture(checkpoint['arch'] ,args.dataset)
denoiser.load_state_dict(checkpoint['state_dict'])
base_classifier = torch.nn.Sequential(denoiser, base_classifier)
def main():
if not os.path.exists(args.outdir):
os.mkdir(args.outdir)
device = torch.device("cuda")
torch.cuda.set_device(args.gpu)
logfilename = os.path.join(args.outdir, args.logname)
log(logfilename, "Hyperparameter List")
log(logfilename, "Epochs: {:}".format(args.epochs))
log(logfilename, "Learning Rate: {:}".format(args.lr))
log(logfilename, "Alpha: {:}".format(args.alpha))
log(logfilename, "Keep ratio: {:}".format(args.keep_ratio))
test_acc_list = []
for _ in range(args.round):
train_dataset = get_dataset(args.dataset, 'train')
test_dataset = get_dataset(args.dataset, 'test')
pin_memory = (args.dataset == "imagenet")
train_loader = DataLoader(train_dataset,
shuffle=True,
batch_size=args.batch,
num_workers=args.workers,
pin_memory=pin_memory)
test_loader = DataLoader(test_dataset,
shuffle=False,
batch_size=args.batch,
num_workers=args.workers,
pin_memory=pin_memory)
# Loading the base_classifier
base_classifier = get_architecture(args.arch, args.dataset, device)
checkpoint = torch.load(args.savedir)
base_classifier.load_state_dict(checkpoint['state_dict'])
base_classifier.eval()
print("Loaded the base_classifier")
original_acc = model_inference(base_classifier,
test_loader,
device,
display=True)
log(logfilename,
"Original Model Test Accuracy: {:.5}".format(original_acc))
print("Original Model Test Accuracy, ", original_acc)
# Creating a fresh copy of network not affecting the original network.
net = copy.deepcopy(base_classifier)
net = net.to(device)
# Generating the mask 'm'
for layer in net.modules():
if isinstance(layer, nn.Linear) or isinstance(layer, nn.Conv2d):
layer.weight_mask = nn.Parameter(torch.ones_like(layer.weight))
layer.weight.requires_grad = True
layer.weight_mask.requires_grad = True
# This is the monkey-patch overriding layer.forward to custom function.
# layer.forward will pass nn.Linear with weights: 'w' and 'm' elementwised
if isinstance(layer, nn.Linear):
layer.forward = types.MethodType(mask_forward_linear, layer)
if isinstance(layer, nn.Conv2d):
layer.forward = types.MethodType(mask_forward_conv2d, layer)
criterion = nn.NLLLoss().to(
device) # I added Log Softmax layer to all architecture.
optimizer = SGD(
net.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=0) # weight_decay = 0 for training the mask.
sparsity, total = 0, 0
for layer in net.modules():
if isinstance(layer, nn.Linear) or isinstance(layer, nn.Conv2d):
boolean_list = layer.weight_mask.data > args.threshold
sparsity += (boolean_list == 1).sum()
total += layer.weight.numel()
# Training the mask with the training set.
# You can set the maximum number of loop in case the sparsity on auxiliary parameter
# do not go below target sparsity.
for epoch in range(300):
if epoch % 5 == 0:
print("Current epochs: ", epoch)
print("Sparsity: {:}".format(sparsity))
train_loss = mask_train(train_loader,
net,
criterion,
optimizer,
epoch,
device,
alpha=args.alpha,
display=False)
acc = model_inference(net, test_loader, device, display=False)
log(logfilename,
"Epoch {:}, Mask Update Test Acc: {:.5}".format(epoch, acc))
sparsity, total = 0, 0
for layer in net.modules():
if isinstance(layer, nn.Linear) or isinstance(
layer, nn.Conv2d):
boolean_list = layer.weight_mask.data > args.threshold
sparsity += (boolean_list == 1).sum()
total += layer.weight.numel()
if sparsity <= total * args.keep_ratio:
print("Current epochs breaking loop at {:}".format(epoch))
break
mask_update_acc = model_inference(net,
test_loader,
device,
display=True)
log(logfilename,
"Mask Update Test Accuracy: {:.5}".format(mask_update_acc))
# This line allows to calculate the threshold to satisfy the keep_ratio.
c_abs = []
for layer in net.modules():
if isinstance(layer, nn.Linear) or isinstance(layer, nn.Conv2d):
c_abs.append(torch.abs(layer.weight_mask))
all_scores = torch.cat([torch.flatten(x) for x in c_abs])
num_params_to_keep = int(len(all_scores) * args.keep_ratio)
threshold, _ = torch.topk(all_scores, num_params_to_keep, sorted=True)
threshold = threshold[-1]
keep_masks = []
for c in c_abs:
keep_masks.append((c >= threshold).float())
print(
"Number of ones.",
torch.sum(torch.cat([torch.flatten(x == 1) for x in keep_masks])))
# Updating the weight with elementwise product of update c.
for layer in net.modules():
if isinstance(layer, nn.Linear) or isinstance(layer, nn.Conv2d):
# We update the weight by elementwise multiplication between
# weight 'w' and mask 'm'.
layer.weight.data = layer.weight.data * layer.weight_mask.data
layer.zeros = nn.Parameter(torch.zeros_like(
layer.weight)) # Dummy parameter.
layer.ones = nn.Parameter(torch.ones_like(
layer.weight)) # Dummy parameter.
layer.weight_mask.data = torch.where(
torch.abs(layer.weight_mask) <= threshold, layer.zeros,
layer.ones
) # Updated weight_mask becomes the mask with element
# 0 and 1 again.
# Temporarily disabling the backprop for both 'w' and 'm'.
layer.weight.requires_grad = False
layer.weight_mask.requires_grad = False
if isinstance(layer, nn.Linear):
layer.forward = types.MethodType(mask_forward_linear, layer)
if isinstance(layer, nn.Conv2d):
layer.forward = types.MethodType(mask_forward_conv2d, layer)
weight_update_acc = model_inference(net,
test_loader,
device,
display=True)
log(logfilename,
"Weight Update Test Accuracy: {:.5}".format(weight_update_acc))
# Calculating the sparsity of the network.
remain = 0
total = 0
for layer in net.modules():
if isinstance(layer, nn.Linear) or isinstance(layer, nn.Conv2d):
total += torch.norm(torch.ones_like(layer.weight),
p=1) # Counting total num parameter
remain += torch.norm(layer.weight_mask.data,
p=1) # Counting ones in the mask.
# Disabling backprop except weight 'w' for the finetuning.
layer.zeros.requires_grad = False
layer.ones.requires_grad = False
layer.weight_mask.requires_grad = False
layer.weight.requires_grad = True
if isinstance(layer, nn.Linear):
layer.forward = types.MethodType(mask_forward_linear, layer)
if isinstance(layer, nn.Conv2d):
layer.forward = types.MethodType(mask_forward_conv2d, layer)
log(logfilename, "Sparsity: {:.3}".format(remain / total))
print("Sparsity: ", remain / total)
# --------------------------------
# We need to transfer the weight we learned from "net" to "base_classifier".
for (layer1, layer2) in zip(base_classifier.modules(), net.modules()):
if isinstance(layer1, (nn.Linear, nn.Conv2d)) or isinstance(
layer2, (nn.Linear, nn.Conv2d)):
layer1.weight.data = layer2.weight.data
if layer1.bias != None:
layer1.bias.data = layer2.bias.data
layer1.bias.requires_grad = True
layer1.weight.requires_grad = True
# Applying the mask to the base_classifier.
apply_prune_mask(base_classifier, keep_masks)
# --------------------------------
optimizer = SGD(base_classifier.parameters(),
lr=1e-3,
momentum=args.momentum,
weight_decay=args.weight_decay)
loss = nn.NLLLoss()
scheduler = StepLR(optimizer,
step_size=args.lr_step_size,
gamma=args.gamma)
test_acc = []
# Finetuning via ignite
trainer = create_supervised_trainer(base_classifier, optimizer,
nn.NLLLoss(), device)
evaluator = create_supervised_evaluator(base_classifier, {
'accuracy': Accuracy(),
'nll': Loss(loss)
}, device)
pbar = ProgressBar()
pbar.attach(trainer)
@trainer.on(Events.ITERATION_COMPLETED)
def log_training_loss(engine):
iter_in_epoch = (engine.state.iteration -
1) % len(train_loader) + 1
if engine.state.iteration % args.print_freq == 0:
pbar.log_message("Epoch[{}] Iteration[{}/{}] Loss: {:.2f}"
"".format(engine.state.epoch, iter_in_epoch,
len(train_loader),
engine.state.output))
@trainer.on(Events.EPOCH_COMPLETED)
def log_epoch(engine):
scheduler.step()
evaluator.run(test_loader)
metrics = evaluator.state.metrics
avg_accuracy = metrics['accuracy']
avg_nll = metrics['nll']
pbar.log_message(
"Validation Results - Epoch: {} Avg accuracy: {:.5f} Avg loss: {:.3f}"
.format(engine.state.epoch, avg_accuracy, avg_nll))
log(
logfilename,
"Validation - Epoch: {} Avg accuracy: {:.5f} Avg loss: {:.3f}"
.format(engine.state.epoch, avg_accuracy, avg_nll))
test_acc.append(avg_accuracy)
if avg_accuracy >= max(test_acc):
print("Saving the model at Epoch {:}".format(
engine.state.epoch))
torch.save(
{
'arch': args.arch,
'state_dict': base_classifier.state_dict(),
'optimizer': optimizer.state_dict(),
}, os.path.join(args.outdir, 'checkpoint.pth.tar'))
if engine.state.epoch == args.epochs:
test_acc_list.append(max(test_acc))
log(logfilename,
"Finetuned Test Accuracy: {:.5f}".format(max(test_acc)))
print("Finetuned Test Accuracy: ", max(test_acc))
trainer.run(train_loader, args.epochs)
log(logfilename, "This is the test accuracy list for args.round.")
log(logfilename, str(test_acc_list))
parser.add_argument("--N0", type=int, default=100)
parser.add_argument("--split",
choices=["train", "test"],
default="test",
help="train or test set")
parser.add_argument("--verbstep",
type=int,
default=100,
help="print for how many subslices")
args = parser.parse_args()
if __name__ == '__main__':
# load the base classifier
checkpoint = torch.load(args.base_classifier)
base_classifier = get_architecture(checkpoint["arch"], args.dataset)
base_classifier.load_state_dict(checkpoint['state_dict'])
# iterate through the dataset
dataset = get_dataset(args.dataset, args.split)
# create the smooothed classifier g
smoothed_classifier = StrictRotationSmooth(base_classifier,
get_num_classes(args.dataset),
args.noise_sd)
tot, tot_good = 0, 0
for i in range(len(dataset)):
# only certify every args.skip examples, and stop after args.max examples
def main():
if args.gpu:
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
if not os.path.exists(args.outdir):
os.makedirs(args.outdir)
# Copy code to output directory
copy_code(args.outdir)
train_dataset = get_dataset(args.dataset, 'train')
test_dataset = get_dataset(args.dataset, 'test')
pin_memory = (args.dataset == "imagenet")
train_loader = DataLoader(train_dataset,
shuffle=True,
batch_size=args.batch,
num_workers=args.workers,
pin_memory=pin_memory)
test_loader = DataLoader(test_dataset,
shuffle=False,
batch_size=args.batch,
num_workers=args.workers,
pin_memory=pin_memory)
model = get_architecture(args.arch, args.dataset)
criterion = CrossEntropyLoss().cuda()
optimizer = SGD(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
scheduler = StepLR(optimizer,
step_size=args.lr_step_size,
gamma=args.gamma)
starting_epoch = 0
logfilename = os.path.join(args.outdir, 'log.txt')
## Resume from checkpoint if exists and if resume flag is True
model_path = os.path.join(args.outdir, 'checkpoint.pth.tar')
if args.resume and os.path.isfile(model_path):
print("=> loading checkpoint '{}'".format(model_path))
checkpoint = torch.load(model_path,
map_location=lambda storage, loc: storage)
starting_epoch = checkpoint['epoch']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})".format(
model_path, checkpoint['epoch']))
else:
if args.resume:
print("=> no checkpoint found at '{}'".format(args.outdir))
init_logfile(
logfilename,
"epoch\ttime\tlr\ttrainloss\ttestloss\ttrainAcc\ttestAcc")
for epoch in range(starting_epoch, args.epochs):
before = time.time()
train_loss, train_acc = train(train_loader, model, criterion,
optimizer, epoch, args.noise_sd)
test_loss, test_acc = test(test_loader, model, criterion,
args.noise_sd)
after = time.time()
log(
logfilename, "{}\t{:.3}\t{:.3}\t{:.3}\t{:.3}\t{:.3}\t{:.3}".format(
epoch, after - before,
scheduler.get_lr()[0], train_loss, test_loss, train_acc,
test_acc))
scheduler.step(epoch)
torch.save(
{
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
}, os.path.join(args.outdir, 'checkpoint.pth.tar'))
def main():
if not os.path.exists(args.outdir):
os.mkdir(args.outdir)
device = torch.device("cuda")
torch.cuda.set_device(args.gpu)
logfilename = os.path.join(args.outdir, args.logname)
log(logfilename, "Hyperparameter List")
log(logfilename, "Epochs: {:}".format(args.epochs))
log(logfilename, "Learning Rate: {:}".format(args.lr))
log(logfilename, "Alpha: {:}".format(args.alpha))
log(logfilename, "Keep ratio: {:}".format(args.keep_ratio))
log(logfilename, "Warmup Epochs: {:}".format(args.epochs_warmup))
test_acc_list = []
for _ in range(args.round):
train_dataset = get_dataset(args.dataset, 'train')
test_dataset = get_dataset(args.dataset, 'test')
pin_memory = (args.dataset == "imagenet")
train_loader = DataLoader(train_dataset,
shuffle=True,
batch_size=args.batch,
num_workers=args.workers,
pin_memory=pin_memory)
test_loader = DataLoader(test_dataset,
shuffle=False,
batch_size=args.batch,
num_workers=args.workers,
pin_memory=pin_memory)
# Loading the base_classifier
base_classifier = get_architecture(args.arch, args.dataset, device)
print("Loaded the base_classifier")
criterion = nn.NLLLoss().to(device)
optimizer = SGD(base_classifier.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# Warmup training for the rewinding.
for epoch in range(args.epochs_warmup):
print("Warmup Training Epochs: {:}".format(epoch))
train_loss, train_top1, train_top5 = utils.train(train_loader,
base_classifier,
criterion,
optimizer,
epoch,
device,
print_freq=100,
display=False)
original_acc = model_inference(base_classifier,
test_loader,
device,
display=True)
log(logfilename,
"Warmup Model Test Accuracy: {:.5}".format(original_acc))
print("Warmup Model Test Accuracy, ", original_acc)
# Creating a fresh copy of network not affecting the original network.
# Goal is to find the supermask.
net = copy.deepcopy(base_classifier)
net = net.to(device)
# Generating the mask 'm'
for layer in net.modules():
if isinstance(layer, nn.Linear) or isinstance(layer, nn.Conv2d):
layer.weight_mask = nn.Parameter(torch.ones_like(layer.weight))
layer.weight.requires_grad = True
layer.weight_mask.requires_grad = True
# This is the monkey-patch overriding layer.forward to custom function.
# layer.forward will pass nn.Linear with weights: 'w' and 'm' elementwised
if isinstance(layer, nn.Linear):
layer.forward = types.MethodType(mask_forward_linear, layer)
if isinstance(layer, nn.Conv2d):
layer.forward = types.MethodType(mask_forward_conv2d, layer)
criterion = nn.NLLLoss().to(device)
optimizer = SGD(net.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=0)
# weight_decay = 0 for training the mask.
sparsity, total = 0, 0
for layer in net.modules():
if isinstance(layer, nn.Linear) or isinstance(layer, nn.Conv2d):
boolean_list = layer.weight_mask.data > args.threshold
sparsity += (boolean_list == 1).sum()
total += layer.weight.numel()
# Training the mask with the training set.
for epoch in range(300):
if epoch % 5 == 0:
print("Current epochs: ", epoch)
print("Sparsity: {:}".format(sparsity))
before = time.time()
train_loss = mask_train(train_loader,
net,
criterion,
optimizer,
epoch,
device,
alpha=args.alpha,
display=False)
acc = model_inference(net, test_loader, device, display=False)
log(logfilename,
"Epoch {:}, Mask Update Test Acc: {:.5}".format(epoch, acc))
sparsity = 0
total = 0
for layer in net.modules():
if isinstance(layer, nn.Linear) or isinstance(
layer, nn.Conv2d):
boolean_list = layer.weight_mask.data > 1e-2
sparsity += (boolean_list == 1).sum()
total += layer.weight.numel()
if sparsity <= total * args.keep_ratio:
print("Current epochs breaking loop at {:}".format(epoch))
break
# This line allows to calculate the threshold to satisfy the keep_ratio.
c_abs = []
for layer in net.modules():
if isinstance(layer, nn.Linear) or isinstance(layer, nn.Conv2d):
c_abs.append(torch.abs(layer.weight_mask))
all_scores = torch.cat([torch.flatten(x) for x in c_abs])
num_params_to_keep = int(len(all_scores) * args.keep_ratio)
threshold, _ = torch.topk(all_scores, num_params_to_keep, sorted=True)
threshold = threshold[-1]
keep_masks = []
for c in c_abs:
keep_masks.append((c >= threshold).float())
print(
"Number of ones.",
torch.sum(torch.cat([torch.flatten(x == 1) for x in keep_masks])))
# Applying the mask to the original network.
apply_prune_mask(base_classifier, keep_masks)
mask_update_acc = model_inference(base_classifier,
test_loader,
device,
display=True)
log(logfilename,
"Untrained Network Test Accuracy: {:.5}".format(mask_update_acc))
print("Untrained Network Test Accuracy: {:.5}".format(mask_update_acc))
optimizer = SGD(base_classifier.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
loss = nn.NLLLoss()
scheduler = MultiStepLR(
optimizer,
milestones=[
int(args.epochs * 0.5) - args.epochs_warmup,
int(args.epochs * 0.75) - args.epochs_warmup
],
last_epoch=-1)
test_acc = [] # Collecting the test accuracy
# Finetuning via ignite
trainer = create_supervised_trainer(base_classifier, optimizer,
nn.NLLLoss(), device)
evaluator = create_supervised_evaluator(base_classifier, {
'accuracy': Accuracy(),
'nll': Loss(loss)
}, device)
pbar = ProgressBar()
pbar.attach(trainer)
@trainer.on(Events.ITERATION_COMPLETED)
def log_training_loss(engine):
iter_in_epoch = (engine.state.iteration -
1) % len(train_loader) + 1
if engine.state.iteration % args.print_freq == 0:
pbar.log_message("Epoch[{}] Iteration[{}/{}] Loss: {:.2f}"
"".format(engine.state.epoch, iter_in_epoch,
len(train_loader),
engine.state.output))
@trainer.on(Events.EPOCH_COMPLETED)
def log_epoch(engine):
scheduler.step()
evaluator.run(test_loader)
metrics = evaluator.state.metrics
avg_accuracy = metrics['accuracy']
avg_nll = metrics['nll']
pbar.log_message(
"Validation Results - Epoch: {} Avg accuracy: {:.3f} Avg loss: {:.3f}"
.format(engine.state.epoch + args.epochs_warmup, avg_accuracy,
avg_nll))
log(
logfilename,
"Validation - Epoch: {} Avg accuracy: {:.3f} Avg loss: {:.3f}"
.format(engine.state.epoch + args.epochs_warmup, avg_accuracy,
avg_nll))
test_acc.append(avg_accuracy)
if avg_accuracy >= max(test_acc):
print(
"Saving the model at Epoch {:}".format(engine.state.epoch +
args.epochs_warmup))
torch.save(
{
'arch': args.arch,
'state_dict': base_classifier.state_dict(),
'optimizer': optimizer.state_dict(),
}, os.path.join(args.outdir, 'checkpoint.pth.tar'))
if engine.state.epoch + args.epochs_warmup == args.epochs:
test_acc_list.append(max(test_acc))
log(logfilename,
"Finetuned Test Accuracy: {:.5f}".format(max(test_acc)))
print("Finetuned Test Accuracy: ", max(test_acc))
trainer.run(train_loader, args.epochs - args.epochs_warmup)
log(logfilename, "This is the test accuracy list for args.round.")
log(logfilename, str(test_acc_list))
def load_config(cfg_id=1):
EXPR_CFG_FILE = "cfg/experiment_%d.cfg" % cfg_id
config = configparser.ConfigParser()
config.read(EXPR_CFG_FILE)
experimentID = config["experiment"]["ID"]
options = config["visualize"]
clean_data_root = options["clean_data_root"]
poison_root = options["poison_root"]
denoiser_path = options["denoiser_path"]
arch = options["arch"]
noise_sd = float(options["noise_sd"])
patch_size = int(options["patch_size"])
rand_loc = options.getboolean("rand_loc")
trigger_id = int(options["trigger_id"])
num_poison = int(options["num_poison"])
num_classes = int(options["num_classes"])
batch_size = int(options["batch_size"])
options = config["poison_generation"]
target_wnid = options["target_wnid"]
source_wnid_list = options["source_wnid_list"].format(experimentID)
num_source = int(options["num_source"])
#############################################################################################################
feature_extract = True
# Models to choose from [resnet, alexnet, vgg, squeezenet, densenet, inception]
model_name = arch
# Initialize the model for this run
model_ft, input_size = initialize_model(model_name,
num_classes,
feature_extract,
use_pretrained=True)
# Transforms
data_transforms = transforms.Compose([
transforms.Resize((input_size, input_size)),
transforms.ToTensor(),
])
dataset_clean = LabeledDataset(
clean_data_root + "/train",
"data/{}/finetune_filelist.txt".format(experimentID), data_transforms)
dataset_test = LabeledDataset(
clean_data_root + "/val",
"data/{}/test_filelist.txt".format(experimentID), data_transforms)
dataset_patched = LabeledDataset(
clean_data_root + "/val",
"data/{}/patched_filelist.txt".format(experimentID), data_transforms)
dataloaders_dict = {}
dataloaders_dict['test'] = torch.utils.data.DataLoader(
dataset_test, batch_size=batch_size, shuffle=False, num_workers=4)
dataloaders_dict['patched'] = torch.utils.data.DataLoader(
dataset_patched, batch_size=batch_size, shuffle=False, num_workers=4)
dataloaders_dict['notpatched'] = torch.utils.data.DataLoader(
dataset_patched, batch_size=batch_size, shuffle=False, num_workers=4)
################################################################################################################
trans_trigger = transforms.Compose([
transforms.Resize((patch_size, patch_size)),
transforms.ToTensor(),
])
trigger = Image.open(
'data/triggers/trigger_{}.png'.format(trigger_id)).convert('RGB')
trigger = trans_trigger(trigger).unsqueeze(0).cuda()
checkpointDir = "finetuned_models/" + experimentID + "/" + str(arch) + "/rand_loc_" + str(rand_loc) + \
"/patch_size_" + str(patch_size) + "/num_poison_" + str(num_poison) + "/trigger_" + str(trigger_id)
normalize = NormalizeByChannelMeanStd(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
model = nn.Sequential(normalize, model_ft)
checkpoint = torch.load(os.path.join(checkpointDir, "poisoned_model.pt"),
map_location='cuda:0')
model.load_state_dict(checkpoint['state_dict'])
classifier = model
classifier.eval()
#######################################################################################
# Load denoiser
checkpoint = torch.load(os.path.join(denoiser_path, "checkpoint.pth.tar"))
denoiser = get_architecture(checkpoint['arch'], 'imagenet')
denoiser.load_state_dict(checkpoint['state_dict'])
denoised_classifier = torch.nn.Sequential(denoiser.module, model)
denoised_classifier = torch.nn.DataParallel(denoised_classifier).cuda()
#######################################################################################
denoised_classifier = denoised_classifier.cuda()
return dataloaders_dict, classifier, denoised_classifier, trigger
def main():
if args.gpu:
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
if not os.path.exists(args.outdir):
os.mkdir(args.outdir)
train_dataset = get_dataset(args.dataset, 'train')
test_dataset = get_dataset(args.dataset, 'test')
pin_memory = (args.dataset == "imagenet")
train_loader = DataLoader(train_dataset,
shuffle=True,
batch_size=args.batch,
num_workers=args.workers,
pin_memory=pin_memory)
test_loader = DataLoader(test_dataset,
shuffle=False,
batch_size=args.batch,
num_workers=args.workers,
pin_memory=pin_memory)
if args.base_classifier == 'ori':
model = get_architecture(args.arch, args.dataset)
else:
checkpoint = torch.load(args.base_classifier)
model = get_architecture(checkpoint["arch"], args.dataset)
model.load_state_dict(checkpoint['state_dict'])
logfilename = os.path.join(args.outdir, 'log.txt')
init_logfile(logfilename,
"epoch\ttime\tlr\ttrain loss\ttrain acc\ttestloss\ttest acc")
criterion = CrossEntropyLoss().cuda()
optimizer = SGD(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
scheduler = StepLR(optimizer,
step_size=args.lr_step_size,
gamma=args.gamma)
for epoch in range(args.epochs):
scheduler.step(epoch)
before = time.time()
train_loss, train_acc = train(train_loader, model, criterion,
optimizer, epoch)
test_loss, test_acc = test(test_loader, model, criterion)
after = time.time()
log(
logfilename, "{}\t{:.3}\t{:.3}\t{:.3}\t{:.3}\t{:.3}\t{:.3}".format(
epoch, str(datetime.timedelta(seconds=(after - before))),
scheduler.get_lr()[0], train_loss, train_acc, test_loss,
test_acc))
torch.save(
{
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
}, os.path.join(args.outdir, 'checkpoint.pth.tar'))
def main(args):
if args.gpu:
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
# load the base classifier
checkpoint = torch.load(args.base_classifier)
base_classifier = get_architecture(checkpoint["arch"], args.dataset)
print('arch:', checkpoint['arch'])
if checkpoint["arch"] == 'resnet50' and args.dataset == "imagenet":
try:
base_classifier.load_state_dict(checkpoint['state_dict'])
except Exception as e:
print('direct load failed, try alternative')
try:
base_classifier = torchvision.models.resnet50(
pretrained=False).cuda()
base_classifier.load_state_dict(checkpoint['state_dict'])
# fix
# normalize_layer = get_normalize_layer('imagenet').cuda()
# base_classifier = torch.nn.Sequential(normalize_layer, base_classifier)
except Exception as e:
print('alternative failed again, try alternative 2')
base_classifier = torchvision.models.resnet50(
pretrained=False).cuda()
# base_classifier.load_state_dict(checkpoint['state_dict'])
normalize_layer = get_normalize_layer('imagenet').cuda()
base_classifier = torch.nn.Sequential(normalize_layer,
base_classifier)
base_classifier.load_state_dict(checkpoint['state_dict'])
else:
base_classifier.load_state_dict(checkpoint['state_dict'])
# iterate through the dataset
dataset = get_dataset(args.dataset, args.split)
# generate transformer
transformer = gen_inference_transformer(args, dataset[0][0])
smoothed_classifier = SemanticSmooth(base_classifier,
get_num_classes(args.dataset),
transformer)
# generate image-level transform and params
tinst1, tfunc1, tinst2, tfunc2, param1l, param1r, param2l, param2r, candidates = gen_transform_and_params(
args, dataset[0][0])
# init random number generator
# m1 = Uniform(param1l, param1r)
# if param2l is not None:
# m2 = Uniform(param2l, param2r)
# m1 = Uniform(param1l, param1r)
m1 = Beta(0.5, 0.5)
if param2l is not None:
m2 = Beta(0.5, 0.5)
# m2 = Uniform(param2l, param2r)
# init metrics
tot = tot_benign = tot_robust = 0
# [main] attack section
for i in range(len(dataset)):
# only certify every args.skip examples
if i % args.skip != 0:
continue
print('working on #', i)
(x, y) = dataset[i]
pred = predict(smoothed_classifier, base_classifier, args, x)
if pred != y:
pass
else:
tot_benign += 1
robust = True
for j in range(0, args.tries):
param_sample1 = (m1.sample() * (param1r - param1l) +
param1l).item()
if param2l is not None:
param_sample2 = (m2.sample() * (param2r - param2l) +
param2l).item()
else:
param_sample2 = None
xp = process(x, tfunc1, tinst1, tfunc2, tinst2, param_sample1,
param_sample2)
pre_loss = getloss(smoothed_classifier, base_classifier, args,
xp, y)
if param_sample2 is None:
print(
f"{i} > {j}/{args.tries} begin para1={param_sample1:4.2f} loss={pre_loss}",
flush=True)
else:
print(
f"{i} > {j}/{args.tries} begin para1={param_sample1:4.2f} para2={param_sample2:4.2f} loss={pre_loss}",
flush=True)
# first work on param1
eps = (param1r - param1l) / args.stepdiv
xp_l = process(x, tfunc1, tinst1, tfunc2, tinst2,
min(max(param_sample1 - eps, param1l), param1r),
param_sample2)
xp_r = process(x, tfunc1, tinst1, tfunc2, tinst2,
min(max(param_sample1 + eps, param1l), param1r),
param_sample2)
loss_l = getloss(smoothed_classifier, base_classifier, args,
xp_l, y)
loss_r = getloss(smoothed_classifier, base_classifier, args,
xp_r, y)
coef_1 = 1 if loss_r > loss_l else -1
now_loss = max(loss_l, loss_r)
now_param1 = param_sample1
if now_loss > pre_loss:
while True:
incre = min(max(now_param1 + coef_1 * eps, param1l),
param1r)
new_xp = process(x, tfunc1, tinst1, tfunc2, tinst2,
incre, param_sample2)
new_loss = getloss(smoothed_classifier,
base_classifier, args, new_xp, y)
# print(f"{i} > {j}/{args.tries} iter para1={now_param1 + coef_1 * eps} loss={new_loss}", flush=True)
if new_loss < now_loss or (
not param1l < incre < param1r):
break
now_param1 = incre
now_loss = new_loss
tmp_l = now_param1 - coef_1 * eps
tmp_r = now_param1 + coef_1 * eps
tmp_l = min(max(tmp_l, param1l), param1r)
tmp_r = min(max(tmp_r, param1l), param1r)
# tri-section search
while tmp_r - tmp_l > eps / args.stepdiv:
tmp_m1 = (2.0 * tmp_l + tmp_r) / 3.0
tmp_m2 = (tmp_l + 2.0 * tmp_r) / 3.0
xp_m1 = process(x, tfunc1, tinst1, tfunc2, tinst2, tmp_m1,
param_sample2)
xp_m2 = process(x, tfunc1, tinst1, tfunc2, tinst2, tmp_m2,
param_sample2)
loss_m1 = getloss(smoothed_classifier, base_classifier,
args, xp_m1, y)
loss_m2 = getloss(smoothed_classifier, base_classifier,
args, xp_m2, y)
# print(f"{i} > {j}/{args.tries} search para1={tmp_m1} loss={loss_m1}", flush=True)
# print(f"{i} > {j}/{args.tries} search para1={tmp_m2} loss={loss_m2}", flush=True)
if loss_m1 > loss_m2:
tmp_r = tmp_m2
else:
tmp_l = tmp_m1
targ_param1 = (tmp_l + tmp_r) / 2.0
# now work on param2
if tfunc2 is not None:
eps = (param2r - param2l) / args.stepdiv
xp = process(x, tfunc1, tinst1, tfunc2, tinst2,
targ_param1, param_sample2)
pre_loss2 = getloss(smoothed_classifier, base_classifier,
args, xp, y)
xp_l = process(
x, tfunc1, tinst1, tfunc2, tinst2, targ_param1,
min(max(param_sample2 - eps, param2l), param2r))
xp_r = process(
x, tfunc1, tinst1, tfunc2, tinst2, targ_param1,
min(max(param_sample2 + eps, param2l), param2r))
loss_l = getloss(smoothed_classifier, base_classifier,
args, xp_l, y)
loss_r = getloss(smoothed_classifier, base_classifier,
args, xp_r, y)
coef_2 = 1 if loss_r > loss_l else -1
now_loss = max(loss_l, loss_r)
now_param2 = param_sample2
if now_loss > pre_loss2:
while True:
incre = min(
max(now_param2 + coef_2 * eps, param2l),
param2r)
new_xp = process(x, tfunc1, tinst1, tfunc2, tinst2,
targ_param1, incre)
new_loss = getloss(smoothed_classifier,
base_classifier, args, new_xp,
y)
if new_loss < now_loss or (
not param2l < incre < param2r):
break
now_param2 = incre
now_loss = new_loss
tmp_l = now_param2 - coef_2 * eps
tmp_r = now_param2 + coef_2 * eps
tmp_l = min(max(tmp_l, param2l), param2r)
tmp_r = min(max(tmp_r, param2l), param2r)
# tri-section search
while tmp_r - tmp_l > eps / args.stepdiv:
tmp_m1 = (2.0 * tmp_l + tmp_r) / 3.0
tmp_m2 = (tmp_l + 2.0 * tmp_r) / 3.0
xp_m1 = process(x, tfunc1, tinst1, tfunc2, tinst2,
targ_param1, tmp_m1)
xp_m2 = process(x, tfunc1, tinst1, tfunc2, tinst2,
targ_param1, tmp_m2)
loss_m1 = getloss(smoothed_classifier, base_classifier,
args, xp_m1, y)
loss_m2 = getloss(smoothed_classifier, base_classifier,
args, xp_m2, y)
if loss_m1 > loss_m2:
tmp_r = tmp_m2
else:
tmp_l = tmp_m1
targ_param2 = (tmp_l + tmp_r) / 2.0
xp = tfunc1(tinst1, x, targ_param1)
if param2l is not None:
xp = tfunc2(tinst2, xp, targ_param2)
xp = xp.type_as(x)
fin_loss = getloss(smoothed_classifier, base_classifier, args,
xp, y)
if param_sample2 is None:
print(
f"{i} > {j}/{args.tries} end para1={targ_param1:4.2f} loss={fin_loss}",
flush=True)
else:
print(
f"{i} > {j}/{args.tries} end para1={targ_param1:4.2f} para2={targ_param2:4.2f} loss={fin_loss}",
flush=True)
if args.transtype in [
'rotation-brightness-l2', 'scaling-brightness-l2'
]:
# compute the gradient by soft label and empirical mean
smoothed_classifier.base_classifier.eval()
grad = torch.zeros(
(args.N0, xp.shape[0], xp.shape[1], xp.shape[2]))
n = 0
while n < args.N0:
now_batch = min(args.batch, args.N0 - n)
if args.nosmooth is True:
batch_noised = xp.repeat((1, 1, 1, 1))
else:
batch = xp.repeat((now_batch, 1, 1, 1))
batch_noised = smoothed_classifier.transformer.process(
batch).cuda()
batch_noised = Variable(
batch_noised.data,
requires_grad=True).contiguous()
opt = torch.optim.Adam([batch_noised], lr=1e-3)
opt.zero_grad()
loss = torch.nn.CrossEntropyLoss()(
smoothed_classifier.base_classifier(batch_noised),
torch.tensor(
[y],
dtype=torch.long).expand(now_batch).cuda())
loss.backward()
grad[n:n + now_batch, :, :, :] = batch_noised.grad.data
n += now_batch
grad = torch.mean(grad, dim=0)
unit_grad = F.normalize(grad,
p=2,
dim=list(range(grad.dim())))
delta = unit_grad * args.l2_r
# print(xp)
xp = xp + delta
# print(xp + delta)
# print(delta)
# print(torch.norm(delta.reshape(-1)))
if args.nosmooth is True:
base_classifier.eval()
pred = base_classifier(xp.cuda().unsqueeze(0)).argmax(1)[0]
else:
pred = smoothed_classifier.predict(xp, args.N, args.p,
args.batch)
if (pred != y):
robust = False
break
print(f"{i} > {j}/{args.tries}", flush=True)
tot_robust += int(robust)
tot += 1
print(
f'#{i} clean acc={tot_benign / tot} robust acc={tot_robust / tot}')
if args.outfile is None:
param_str = ''
if args.transtype != 'translation':
param_str = f'{param1r}'
if param2r is not None:
param_str += f'_{param2r}'
else:
param_str = f'{args.displacement}'
args.outfile = args.transtype + '/' + args.dataset + '/' + param_str + '/' + 'result.txt'
out_full_path = os.path.join(args.outfolder, args.outfile)
print('output result to ' + out_full_path)
if not os.path.exists(os.path.dirname(out_full_path)):
os.makedirs(os.path.dirname(out_full_path))
f = open(out_full_path, 'a')
f.write(
f'clean {tot_benign / tot},{tot_benign} robust={tot_robust / tot},{tot_robust} tot={tot}\n'
)
f.close()
print('done')
def main(args):
if args.gpu:
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
# load the base classifier
checkpoint = torch.load(args.base_classifier)
base_classifier = get_architecture(checkpoint["arch"], args.dataset)
print('arch:', checkpoint['arch'])
if checkpoint["arch"] == 'resnet50' and args.dataset == "imagenet":
try:
base_classifier.load_state_dict(checkpoint['state_dict'])
except Exception as e:
print('direct load failed, try alternative')
try:
base_classifier = torchvision.models.resnet50(
pretrained=False).cuda()
base_classifier.load_state_dict(checkpoint['state_dict'])
# fix
# normalize_layer = get_normalize_layer('imagenet').cuda()
# base_classifier = torch.nn.Sequential(normalize_layer, base_classifier)
except Exception as e:
print('alternative failed again, try alternative 2')
base_classifier = torchvision.models.resnet50(
pretrained=False).cuda()
# base_classifier.load_state_dict(checkpoint['state_dict'])
normalize_layer = get_normalize_layer('imagenet').cuda()
base_classifier = torch.nn.Sequential(normalize_layer,
base_classifier)
base_classifier.load_state_dict(checkpoint['state_dict'])
else:
base_classifier.load_state_dict(checkpoint['state_dict'])
# iterate through the dataset
dataset = get_dataset(args.dataset, args.split)
# generate transformer
transformer = gen_inference_transformer(args, dataset[0][0])
smoothed_classifier = SemanticSmooth(base_classifier,
get_num_classes(args.dataset),
transformer)
# generate image-level transform and params
tinst1, tfunc1, tinst2, tfunc2, param1l, param1r, param2l, param2r, candidates = gen_transform_and_params(
args, dataset[0][0])
# init random number generator
m1 = Uniform(param1l, param1r)
if param2l is not None:
m2 = Uniform(param2l, param2r)
# init metrics
tot = tot_benign = tot_robust = 0
# [main] attack section
for i in range(len(dataset)):
# only certify every args.skip examples
if i % args.skip != 0:
continue
print('working on #', i)
(x, y) = dataset[i]
# clean_x = x.cuda().unsqueeze(0)
pred = smoothed_classifier.predict(x, args.N0, args.p, args.batch)
if pred != y:
pass
else:
tot_benign += 1
robust = True
for j in range(0, args.tries):
xp = None
if args.transtype == 'translation':
param_sample1 = candidates[int(m1.sample().item())]
xp = tfunc1(tinst1, x, param_sample1[0].item(),
param_sample1[1].item())
else:
param_sample1 = m1.sample().item()
if param2l is not None:
param_sample2 = m2.sample().item()
xp = tfunc1(tinst1, x, param_sample1)
if param2l is not None:
xp = tfunc2(tinst2, xp, param_sample2)
# xp = xp.contiguous().cuda()
# xp_old = xp
# if args.l2 is not None and args.l2 > EPS:
# xp = fgsm(model, xp, torch.tensor([y], dtype=torch.long).expand(now_batch).cuda(), args.l2)
# print(torch.norm((xp_old - xp).reshape(xp.size()[0], -1), dim=1))
xp = xp.type_as(x)
if args.transtype in [
'rotation-brightness-l2', 'scaling-brightness-l2'
]:
# compute the gradient by soft label and empirical mean
smoothed_classifier.base_classifier.eval()
grad = torch.zeros(
(args.N0, xp.shape[0], xp.shape[1], xp.shape[2]))
n = 0
while n < args.N0:
now_batch = min(args.batch, args.N0 - n)
batch = xp.repeat((now_batch, 1, 1, 1))
batch_noised = smoothed_classifier.transformer.process(
batch).cuda()
batch_noised = Variable(
batch_noised.data,
requires_grad=True).contiguous()
opt = torch.optim.Adam([batch_noised], lr=1e-3)
opt.zero_grad()
loss = torch.nn.CrossEntropyLoss()(
smoothed_classifier.base_classifier(batch_noised),
torch.tensor(
[y],
dtype=torch.long).expand(now_batch).cuda())
loss.backward()
grad[n:n + now_batch, :, :, :] = batch_noised.grad.data
n += now_batch
grad = torch.mean(grad, dim=0)
unit_grad = F.normalize(grad,
p=2,
dim=list(range(grad.dim())))
delta = unit_grad * args.l2_r
# print(xp)
xp = xp + delta
# print(xp + delta)
# print(delta)
# print(torch.norm(delta.reshape(-1)))
pred = smoothed_classifier.predict(xp, args.N0, args.p,
args.batch)
if (pred != y):
robust = False
break
print(f"> {j}/{args.tries}", end='\r', flush=True)
tot_robust += int(robust)
tot += 1
print(
f'#{i} clean acc={tot_benign / tot} robust acc={tot_robust / tot}')
if args.outfile is None:
param_str = ''
if args.transtype != 'translation':
param_str = f'{param1r}'
if param2r is not None:
param_str += f'_{param2r}'
else:
param_str = f'{args.displacement}'
args.outfile = args.transtype + '/' + args.dataset + '/' + param_str + '/' + 'result.txt'
out_full_path = os.path.join(args.outfolder, args.outfile)
print('output result to ' + out_full_path)
if not os.path.exists(os.path.dirname(out_full_path)):
os.makedirs(os.path.dirname(out_full_path))
f = open(out_full_path, 'w')
f.write(
f'clean {tot_benign / tot},{tot_benign} robust={tot_robust / tot},{tot_robust} tot={tot}\n'
)
f.close()
print('done')
X_pgd = Variable(torch.clamp(X_pgd, 0, 1), requires_grad=True)
if (step + 1) >= decay_1:
lr = 0.5 * step_size
decay_1 = perturb_steps + 1
if (step + 1) >= decay_2:
lr = 0.25 * step_size
decay_2 = perturb_steps + 1
return X_pgd.data
if __name__ == "__main__":
resnet152 = get_architecture(denoise=False).cuda()
resnet152.eval()
resnet152_denoise = get_architecture(denoise=True).cuda()
resnet152_denoise.eval()
resnet101_denoise = get_architecture(
denoise=True, model_name="Resnet101-DenoiseAll").cuda()
resnet101_denoise.eval()
model_list = [resnet152, resnet152_denoise, resnet101_denoise]
loader = MyCustomDataset(csv_path=args.csv_path, img_path=args.img_path)
attack_loader = torch.utils.data.DataLoader(
dataset=loader,
def main():
if args.gpu:
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
if not os.path.exists(args.outdir):
os.makedirs(args.outdir)
train_dataset = get_dataset(args.dataset, 'train')
test_dataset = get_dataset(args.dataset, 'test')
pin_memory = (args.dataset == "imagenet")
train_loader = DataLoader(train_dataset,
shuffle=True,
batch_size=args.batch,
num_workers=args.workers,
pin_memory=pin_memory)
test_loader = DataLoader(test_dataset,
shuffle=False,
batch_size=args.batch,
num_workers=args.workers,
pin_memory=pin_memory)
model = get_architecture(args.arch, args.dataset)
if args.pretrain is not None:
if args.pretrain == 'torchvision':
# load pretrain model from torchvision
if args.dataset == 'imagenet' and args.arch == 'resnet50':
model = torchvision.models.resnet50(True).cuda()
# fix
normalize_layer = get_normalize_layer('imagenet').cuda()
model = torch.nn.Sequential(normalize_layer, model)
print('loaded from torchvision for imagenet resnet50')
else:
raise Exception(f'Unsupported pretrain arg {args.pretrain}')
else:
# load the base classifier
checkpoint = torch.load(args.pretrain)
model.load_state_dict(checkpoint['state_dict'])
print(f'loaded from {args.pretrain}')
logfilename = os.path.join(args.outdir, 'log.txt')
init_logfile(logfilename,
"epoch\ttime\tlr\ttrain loss\ttrain acc\ttestloss\ttest acc")
writer = SummaryWriter(args.outdir)
canopy = None
for (inputs, targets) in train_loader:
canopy = inputs[0]
break
transformer = gen_transformer(args, canopy)
criterion = CrossEntropyLoss().cuda()
optimizer = SGD(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
scheduler = StepLR(optimizer,
step_size=args.lr_step_size,
gamma=args.gamma)
for epoch in range(args.epochs):
before = time.time()
train_loss, train_acc = train(train_loader, model, criterion,
optimizer, epoch, transformer, writer)
test_loss, test_acc = test(test_loader, model, criterion, epoch,
transformer, writer, args.print_freq)
after = time.time()
scheduler.step(epoch)
log(
logfilename, "{}\t{:.3}\t{:.3}\t{:.3}\t{:.3}\t{:.3}\t{:.3}".format(
epoch, str(datetime.timedelta(seconds=(after - before))),
scheduler.get_lr()[0], train_loss, train_acc, test_loss,
test_acc))
torch.save(
{
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
}, os.path.join(args.outdir, 'checkpoint.pth.tar'))
def get_basemodel(args):
if args.model == 'none': return None
if args.model == 'resnet50' and args.dataset == 'imagenet':
model = models.resnet50(pretrained=True).eval()
normalize = NormalizeLayer(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
else:
path = glob.glob(
os.path.join('./models', args.model, '**', 'checkpoint.pt.best'))
path_tar = glob.glob(
os.path.join('./models', args.model, '**', 'checkpoint.pth.tar'))
if not (len(path) > 0 or
(len(path_tar) > 0 and args.dataset in ['cifar', 'imagenet'])):
print("Could not load model")
sys.exit(1)
if len(path_tar) > 0 and args.dataset in [
'cifar', 'imagenet', 'restricted_imagenet'
]:
sys.path.append('smoothing-adversarial/code')
from architectures import get_architecture
path = path_tar[0]
print('Loading model from', path)
checkpoint = torch.load(path, map_location='cpu')
if args.dataset == 'cifar':
model = get_architecture(checkpoint["arch"], 'cifar10')
else:
model = get_architecture(checkpoint["arch"], args.dataset)
model.load_state_dict(checkpoint['state_dict'])
model = model.to('cpu')
for i, m in enumerate(model):
if isinstance(m, torch.nn.DataParallel):
model[i] = m.module
normalize = None
model = model[1:]
print(model)
else:
path = path[0]
print('Loading model from', path)
if args.dataset in ['imagenet', 'restricted_imagenet', 'cifar']:
ds_class = datasets.DATASETS[args.dataset]
ds = ds_class("./ds/imagenet"
if args.dataset != 'cifar' else './ds/cifar')
model, _ = model_utils.make_and_restore_model(
arch=('resnet18'
if args.dataset == 'cifar' else 'resnet50'),
dataset=ds,
resume_path=path,
parallel=False)
normalize = model.normalizer
model = model.model
elif args.dataset in ['mnist', 'fashionmnist', 'GTSRB']:
if 'mnist' in args.dataset:
num_classes = 10
color_channels = 1
mean = torch.tensor([0.1307])
std = torch.tensor([0.3081])
if 'convnet' in path:
print('convenet')
model = MNISTConvNet()
else:
model = resnet18(num_classes=num_classes,
color_channels=color_channels)
elif args.dataset == 'GTSRB':
num_classes = 43
color_channels = 3
mean = torch.tensor([0.3337, 0.3064, 0.3171])
std = torch.tensor([0.2672, 0.2564, 0.2629])
model = resnet18(num_classes=num_classes,
color_channels=color_channels)
model = RobustnessModelInnerWrapper(model)
d = argparse.Namespace()
d.mean = mean
d.std = std
model = AttackerModel(model, d)
checkpoint = torch.load(path, pickle_module=dill)
state_dict_path = 'model'
if not ('model' in checkpoint):
state_dict_path = 'state_dict'
sd = checkpoint[state_dict_path]
sd = {k[len('module.'):]: v for k, v in sd.items()}
sd = {(k if 'model.net' in k else k.replace(
'model.', 'model.net.')): v
for k, v in sd.items()}
model.load_state_dict(sd)
normalize = model.normalizer
model = model.model
else:
assert (False)
m = []
if normalize is not None:
m.append(normalize.to(args.device))
if args.radiusDecrease >= 0:
shape = {'rot': 'circ', 'trans': 'rect'}[args.transformation]
size = {
'mnist': (1, 28, 28),
'fashionmnist': (1, 28, 28),
'cifar': (3, 32, 32),
'GTSRB': (3, np.inf, np.inf),
'imagenet': (3, np.inf, np.inf),
'restricted_imagenet': (3, np.inf, np.inf)
}[args.dataset]
if args.resize_post_transform > 0:
size = (size[0], min(size[1], args.resize_post_transform),
min(size[2], args.resize_post_transform))
if args.center_crop_post_transform > 0:
size = (size[0], min(size[1], args.center_crop_post_transform),
min(size[2], args.center_crop_post_transform))
V = VingetteModule(size, shape, args.radiusDecrease)
m.append(V)
m.append(model)
model = torch.nn.Sequential(*m)
if args.use_cuda:
model = torch.nn.DataParallel(model.to(args.device))
model = model.eval().to(args.device)
return model
def main():
torch.manual_seed(0)
if args.gpu:
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
if not os.path.exists(args.outdir):
os.mkdir(args.outdir)
train_loader, test_loader = get_dataset(args.dataset,
'train',
args.batch,
num_workers=args.workers)
# test_dataset = get_dataset(args.dataset, 'valid')
# pin_memory = (args.dataset == "imagenet")
# train_loader = DataLoader(train_dataset, shuffle=True, batch_size=args.batch,
# num_workers=args.workers, pin_memory=pin_memory)
# test_loader = DataLoader(test_dataset, shuffle=False, batch_size=args.batch,
# num_workers=args.workers, pin_memory=pin_memory)
model = get_architecture(args.arch, args.dataset)
logfilename = os.path.join(args.outdir, 'log.txt')
init_logfile(logfilename,
"epoch\ttime\tlr\ttrain loss\ttrain acc\ttestloss\ttest acc")
criterion = CrossEntropyLoss().cuda()
if args.optimizer == 'momentum':
optimizer = SGD(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
elif args.optimizer == 'nesterov':
optimizer = SGD(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay,
nesterov=True)
elif args.optimizer == 'amsgrad':
optimizer = Adam(model.parameters(),
lr=args.lr,
weight_decay=args.weight_decay,
amsgrad=True)
scheduler = StepLR(optimizer,
step_size=args.lr_step_size,
gamma=args.gamma)
best_acc = 0.
for epoch in range(args.epochs):
scheduler.step(epoch)
before = time.time()
train_loss, train_acc = train(train_loader, model, criterion,
optimizer, epoch, args)
test_loss, test_acc = test(test_loader, model, criterion, epoch, args)
after = time.time()
log(
logfilename, "{}\t{:.3}\t{:.3}\t{:.3}\t{:.3}\t{:.3}\t{:.3}".format(
epoch, str(datetime.timedelta(seconds=(after - before))),
scheduler.get_lr()[0], train_loss, train_acc, test_loss,
test_acc))
if args.tune and best_acc < test_acc:
best_acc = test_acc
print('saving best model...')
torch.save(
{
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
}, os.path.join(args.outdir, 'checkpoint.best.tar'))
torch.save(
{
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
}, os.path.join(args.outdir, 'checkpoint.pth.tar'))
def main():
if args.gpu:
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
if not os.path.exists(args.outdir):
os.makedirs(args.outdir)
# Copies files to the outdir to store complete script with each experiment
copy_code(args.outdir)
train_dataset = get_dataset(args.dataset, 'train')
test_dataset = get_dataset(args.dataset, 'test')
pin_memory = (args.dataset == "imagenet")
train_loader = DataLoader(train_dataset,
shuffle=True,
batch_size=args.batch,
num_workers=args.workers,
pin_memory=pin_memory)
test_loader = DataLoader(test_dataset,
shuffle=False,
batch_size=args.batch,
num_workers=args.workers,
pin_memory=pin_memory)
model = get_architecture(args.arch, args.dataset)
if args.attack == 'PGD':
print('Attacker is PGD')
attacker = PGD_L2(steps=args.num_steps,
device='cuda',
max_norm=args.epsilon)
elif args.attack == 'DDN':
print('Attacker is DDN')
attacker = DDN(steps=args.num_steps,
device='cuda',
max_norm=args.epsilon,
init_norm=args.init_norm_DDN,
gamma=args.gamma_DDN)
else:
raise Exception('Unknown attack')
criterion = CrossEntropyLoss().cuda()
optimizer = SGD(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
scheduler = StepLR(optimizer,
step_size=args.lr_step_size,
gamma=args.gamma)
starting_epoch = 0
logfilename = os.path.join(args.outdir, 'log.txt')
# Load latest checkpoint if exists (to handle philly failures)
model_path = os.path.join(args.outdir, 'checkpoint.pth.tar')
if args.resume:
if os.path.isfile(model_path):
print("=> loading checkpoint '{}'".format(model_path))
checkpoint = torch.load(model_path,
map_location=lambda storage, loc: storage)
starting_epoch = checkpoint['epoch']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})".format(
model_path, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(model_path))
if args.adv_training:
init_logfile(
logfilename,
"epoch\ttime\tlr\ttrainloss\ttestloss\ttrainacc\ttestacc\ttestaccNor"
)
else:
init_logfile(
logfilename,
"epoch\ttime\tlr\ttrainloss\ttestloss\ttrainacc\ttestacc")
else:
if args.adv_training:
init_logfile(
logfilename,
"epoch\ttime\tlr\ttrainloss\ttestloss\ttrainacc\ttestacc\ttestaccNor"
)
else:
init_logfile(
logfilename,
"epoch\ttime\tlr\ttrainloss\ttestloss\ttrainacc\ttestacc")
for epoch in range(starting_epoch, args.epochs):
scheduler.step(epoch)
attacker.max_norm = np.min(
[args.epsilon, (epoch + 1) * args.epsilon / args.warmup])
attacker.init_norm = np.min(
[args.epsilon, (epoch + 1) * args.epsilon / args.warmup])
before = time.time()
train_loss, train_acc = train(train_loader, model, criterion,
optimizer, epoch, args.noise_sd,
attacker)
test_loss, test_acc, test_acc_normal = test(test_loader, model,
criterion, args.noise_sd,
attacker)
after = time.time()
if args.adv_training:
log(
logfilename,
"{}\t{:.2f}\t{:.3}\t{:.3}\t{:.3}\t{:.3}\t{:.3}\t{:.3}".format(
epoch, after - before,
scheduler.get_lr()[0], train_loss, test_loss, train_acc,
test_acc, test_acc_normal))
else:
log(
logfilename,
"{}\t{:.2f}\t{:.3}\t{:.3}\t{:.3}\t{:.3}\t{:.3}".format(
epoch, after - before,
scheduler.get_lr()[0], train_loss, test_loss, train_acc,
test_acc))
torch.save(
{
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
}, model_path)
def main():
ckptdir = None if args.ckptdir == 'none' else args.ckptdir
if ckptdir is not None and not os.path.isdir(ckptdir):
os.makedirs(ckptdir)
outdir = None if args.outdir == 'none' else args.outdir
if outdir is not None and not os.path.isdir(outdir):
os.makedirs(outdir)
########## models to train ##########
if args.ens_comp == '3-model-ensemble':
model_list = ['resnet20', 'resnet26', 'resnet32']
elif args.ens_comp == '7-model-ensemble':
model_list = ['lenet', 'alexnet', 'resnet20', 'resnet110', 'densenet', 'vgg16', 'vgg19']
name = model_list if args.name == 'none' else args.name
assert isinstance(name, list), 'name must be a list'
assert len(name) == len(model_list), 'the lengths of name and model_list must be equal'
########## local args ##########
start_epoch = 0
train_epochs = 150
milestones=[60, 90]
########## dataset ##########
trainset, testset, transform_test = get_dataset(args.dataset)
val_set = torch.utils.data.Subset(trainset,
[i for i in range(len(trainset)-args.val_num,len(trainset))])
valloader = torch.utils.data.DataLoader(
val_set, batch_size=args.batch_size, shuffle=True, num_workers=args.workers)
num_classes = 10
device = 'cuda' if torch.cuda.is_available() else 'cpu'
########### load trained models ##############
trained_models = []
for model_id in range(len(model_list)):
checkpoint = torch.load('{}/{}.pth'.format(ckptdir, name[model_id]))
model = get_architecture(model_list[model_id])
model.load_state_dict(checkpoint['net'])
model = Softmaxize(model)
model = model.to(device)
model.eval()
trained_models.append(model)
############### Training Ensemble Paras ###############
ens_model = Ensemble(trained_models, device)
ens_model.set_training_paras()
optimizer = optim.SGD(filter(lambda p: p.requires_grad, ens_model.parameters()),
lr=0.01, momentum=0.9, weight_decay=5e-4)
scheduler = MultiStepLR(optimizer, milestones=milestones, gamma=0.1)
T0 = time.time()
for epoch in range(start_epoch + 1, train_epochs + 1):
print('===train(epoch={}, model=ensemble)==='.format(epoch))
t1 = time.time()
scheduler.step()
ens_model.eval()
ens_train(args.sigma, 1, num_classes, ens_model, valloader, optimizer, device)
t2 = time.time()
print('Elapsed time: {}'.format(t2 - t1))
print('current w:',ens_model.w.data)
print('current alpha:',(ens_model.w.exp()/ens_model.w.exp().sum()).data)
T1 = time.time()
print('Total elapsed time for weight solving: {}'.format(T1 - T0))
# Certify test
print('===test(model=ensemble)===')
t1 = time.time()
ens_model.eval()
certify(ens_model, device, testset, transform_test, num_classes,
'{}/{}'.format(outdir, args.ens_name),
start_img=args.start_img, num_img=args.num_img,
skip=args.skip, sigma=args.sigma)
t2 = time.time()
print('Elapsed time: {}'.format(t2 - t1))
if ckptdir is not None:
# Save checkpoint
print('==> Saving Ens model.pth..')
try:
state = {
'w': ens_model.w.data,
}
torch.save(state, '{}/{}.pth'.format(ckptdir, args.ens_name))
except OSError:
print('OSError while saving model {}.pth'.format(args.ens_name))
print('Ignoring...')
def main():
ckptdir = None if args.ckptdir == 'none' else args.ckptdir
if ckptdir is not None and not os.path.isdir(ckptdir):
os.makedirs(ckptdir)
outdir = None if args.outdir == 'none' else args.outdir
if outdir is not None and not os.path.isdir(outdir):
os.makedirs(outdir)
checkpoint = None if args.resume_ckpt == 'none' else args.resume_ckpt
########## models to train ##########
model = get_architecture(args.arch)
name = args.arch if args.name == 'none' else args.name
########## local args ##########
start_epoch = 0
if args.train_scheme == 'std':
train_epochs = 400
milestones = [150, 300]
if args.train_scheme == 'macer':
train_epochs = 440
milestones = [200, 400]
########## dataset ##########
trainset, testset, transform_test = get_dataset(args.dataset)
train_set = torch.utils.data.Subset(
trainset, [i for i in range(len(trainset) - args.val_num)])
trainloader = torch.utils.data.DataLoader(train_set,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers)
num_classes = 10
device = 'cuda' if torch.cuda.is_available() else 'cpu'
T0 = time.time()
model = model.to(device)
optimizer = optim.SGD(model.parameters(),
lr=args.lr,
momentum=0.9,
weight_decay=5e-4)
scheduler = MultiStepLR(optimizer, milestones=milestones, gamma=0.1)
filename = '{}/{}'.format(outdir, name)
# Resume from checkpoint if required
if checkpoint is not None:
print('==> Resuming from checkpoint..')
print(checkpoint)
checkpoint = torch.load(checkpoint)
model.load_state_dict(checkpoint['net'])
start_epoch = checkpoint['epoch']
scheduler.step(start_epoch)
for epoch in range(start_epoch + 1, train_epochs + 1):
print('===train(epoch={}, model={})==='.format(epoch, name))
t1 = time.time()
model.train()
if args.train_scheme == 'std':
std_train(args.sigma, 1, num_classes, model, trainloader,
optimizer, device)
elif args.train_scheme == 'macer':
macer_train(args.sigma, args.lbd, 16, 16.0, 8.0, num_classes,
model, trainloader, optimizer, device)
else:
raise ValueError('train_scheme must be either std or macer')
scheduler.step()
t2 = time.time()
print('Elapsed time: {}'.format(t2 - t1))
T1 = time.time()
print('Total elapsed time for training: {}'.format(T1 - T0))
# Certify test
print('===test(model={})==='.format(name))
t1 = time.time()
model.eval()
certify(model,
device,
testset,
transform_test,
num_classes,
filename,
start_img=args.start_img,
num_img=args.num_img,
skip=args.skip,
sigma=args.sigma)
t2 = time.time()
print('Elapsed time: {}'.format(t2 - t1))
if ckptdir is not None:
# Save checkpoint
print('==> Saving model {}.pth..'.format(name))
try:
state = {'net': model.state_dict(), 'epoch': epoch}
torch.save(state, '{}/{}.pth'.format(ckptdir, name))
except OSError:
print('OSError while saving {}.pth'.format(name))
print('Ignoring...')
def main():
ckptdir = None if args.ckptdir == 'none' else args.ckptdir
if ckptdir is not None and not os.path.isdir(ckptdir):
os.makedirs(ckptdir)
outdir = None if args.outdir == 'none' else args.outdir
if outdir is not None and not os.path.isdir(outdir):
os.makedirs(outdir)
########## models to train ##########
if args.ens_comp == '3-model-ensemble':
model_list = ['resnet20', 'resnet26', 'resnet32']
elif args.ens_comp == '7-model-ensemble':
model_list = [
'lenet', 'alexnet', 'resnet20', 'resnet110', 'densenet', 'vgg16',
'vgg19'
]
name = model_list if args.name == 'none' else args.name
assert isinstance(name, list), 'name must be a list'
assert len(name) == len(
model_list), 'the lengths of name and model_list must be equal'
########## dataset ##########
trainset, testset, transform_test = get_dataset(args.dataset)
num_classes = 10
device = 'cuda' if torch.cuda.is_available() else 'cpu'
########### load trained models ##############
trained_models = []
for model_id in range(len(model_list)):
checkpoint = torch.load('{}/{}.pth'.format(ckptdir, name[model_id]))
model = get_architecture(model_list[model_id])
model.load_state_dict(checkpoint['net'])
model = Softmaxize(model)
model = model.to(device)
model.eval()
trained_models.append(model)
ens_ckpt = torch.load('{}/{}.pth'.format(ckptdir, args.ens_ckpt))
if args.adp:
ens_model = AEnsemble(trained_models,
device,
w=ens_ckpt['w'],
rt_num=True)
else:
ens_model = Ensemble(trained_models, device, w=ens_ckpt['w'])
ens_model.set_training_paras(weight=False)
# Certify test
print('===test(model=ensemble)===')
t1 = time.time()
ens_model.eval()
certify(ens_model,
device,
testset,
transform_test,
num_classes,
'{}/{}'.format(outdir, args.ens_name),
start_img=args.start_img,
num_img=args.num_img,
skip=args.skip,
sigma=args.sigma,
adp=args.adp)
t2 = time.time()
print('Elapsed time: {}'.format(t2 - t1))
logging.info("Number of clean images: {}".format(len(dataset_clean)))
logging.info("Number of poison images: {}".format(len(dataset_poison)))
normalize = NormalizeByChannelMeanStd(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
model = nn.Sequential(normalize, model_ft)
checkpoint = torch.load(os.path.join(checkpointDir,
"%s_model.pt" % model_type),
map_location='cuda:0')
model.load_state_dict(checkpoint['state_dict'])
#######################################################################################
# Load denoiser
checkpoint = torch.load(os.path.join(denoiser_path, "checkpoint.pth.tar"))
denoiser = get_architecture(checkpoint['arch'], 'imagenet')
denoiser.load_state_dict(checkpoint['state_dict'])
model = torch.nn.Sequential(denoiser.module, model)
#######################################################################################
model = model.cuda()
def visualize(model, dataloader, is_inception=False):
attacker = PGD_L2(steps=num_steps, max_norm=epsilon)
model.eval() # Set model to evaluate mode
running_corrects = 0
def main():
if args.gpu:
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
# args.outdir = os.path.join(os.getenv('PT_DATA_DIR', './'), args.outdir)
if not os.path.exists(args.outdir):
os.makedirs(args.outdir)
# Copy code to output directory
copy_code(args.outdir)
train_dataset = get_dataset(args.dataset, 'train')
test_dataset = get_dataset(args.dataset, 'test')
pin_memory = (args.dataset == "imagenet")
train_loader = DataLoader(train_dataset, shuffle=True, batch_size=args.batch,
num_workers=args.workers, pin_memory=pin_memory)
test_loader = DataLoader(test_dataset, shuffle=False, batch_size=args.batch,
num_workers=args.workers, pin_memory=pin_memory)
## This is used to test the performance of the denoiser attached to a cifar10 classifier
cifar10_test_loader = DataLoader(get_dataset('cifar10', 'test'), shuffle=False, batch_size=args.batch,
num_workers=args.workers, pin_memory=pin_memory)
if args.pretrained_denoiser:
checkpoint = torch.load(args.pretrained_denoiser)
assert checkpoint['arch'] == args.arch
denoiser = get_architecture(checkpoint['arch'], args.dataset)
denoiser.load_state_dict(checkpoint['state_dict'])
else:
denoiser = get_architecture(args.arch, args.dataset)
if args.optimizer == 'Adam':
optimizer = Adam(denoiser.parameters(), lr=args.lr, weight_decay=args.weight_decay)
elif args.optimizer == 'SGD':
optimizer = SGD(denoiser.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)
elif args.optimizer == 'AdamThenSGD':
optimizer = Adam(denoiser.parameters(), lr=args.lr, weight_decay=args.weight_decay)
scheduler = StepLR(optimizer, step_size=args.lr_step_size, gamma=args.gamma)
starting_epoch = 0
logfilename = os.path.join(args.outdir, 'log.txt')
## Resume from checkpoint if exists and if resume flag is True
denoiser_path = os.path.join(args.outdir, 'checkpoint.pth.tar')
if args.resume and os.path.isfile(denoiser_path):
print("=> loading checkpoint '{}'".format(denoiser_path))
checkpoint = torch.load(denoiser_path,
map_location=lambda storage, loc: storage)
assert checkpoint['arch'] == args.arch
starting_epoch = checkpoint['epoch']
denoiser.load_state_dict(checkpoint['state_dict'])
if starting_epoch >= args.start_sgd_epoch and args.optimizer == 'AdamThenSGD ': # Do adam for few steps thaen continue SGD
print("-->[Switching from Adam to SGD.]")
args.lr = args.start_sgd_lr
optimizer = SGD(denoiser.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)
scheduler = StepLR(optimizer, step_size=args.lr_step_size, gamma=args.gamma)
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})"
.format(denoiser_path, checkpoint['epoch']))
else:
if args.resume: print("=> no checkpoint found at '{}'".format(args.outdir))
init_logfile(logfilename, "epoch\ttime\tlr\ttrainloss\ttestloss\ttestAcc")
# the first model is the model we test on, the other 14 models are surrogate models that we use
# (see the paper for more detail)
models = [
"ResNet110", # fixed model
#surrogate models
"WRN", "WRN40", "VGG16", "VGG19", "ResNet18","PreActResNet18","ResNeXt29_2x64d",
"MobileNet","MobileNetV2","SENet18","ShuffleNetV2","EfficientNetB0","GoogLeNet","DenseNet121"
]
path = os.path.join(args.classifiers_path, '{}/noise_0.00/checkpoint.pth.tar')
base_classifiers_paths = [
path.format(model) for model in models
]
base_classifiers = []
if args.classifier_idx == -1:
for base_classifier in base_classifiers_paths:
# load the base classifier
checkpoint = torch.load(base_classifier)
base_classifier = get_architecture(checkpoint["arch"], args.dataset)
base_classifier.load_state_dict(checkpoint['state_dict'])
requires_grad_(base_classifier, False)
base_classifiers.append(base_classifier.eval().cuda())
else:
if args.classifier_idx not in range(len(models)): raise Exception("Unknown model")
print("Model namse: {}".format(base_classifiers_paths[args.classifier_idx]))
checkpoint = torch.load(base_classifiers_paths[args.classifier_idx])
base_classifier = get_architecture(checkpoint["arch"], args.dataset)
base_classifier.load_state_dict(checkpoint['state_dict'])
requires_grad_(base_classifier, False)
base_classifiers.append(base_classifier.eval().cuda())
criterion = CrossEntropyLoss(size_average=None, reduce=None, reduction = 'mean').cuda()
best_acc = 0
for epoch in range(starting_epoch, args.epochs):
before = time.time()
train_loss = train(train_loader, denoiser, criterion, optimizer, epoch, args.noise_sd, base_classifiers[1:])
test_loss, test_acc = test_with_classifier(cifar10_test_loader, denoiser, criterion, args.noise_sd, args.print_freq, base_classifiers[0])
after = time.time()
log(logfilename, "{}\t{:.3}\t{:.3}\t{:.3}\t{:.3}\t{:.3}".format(
epoch, after - before,
args.lr, train_loss, test_loss, test_acc))
scheduler.step(epoch)
args.lr = scheduler.get_lr()[0]
# Switch from Adam to SGD
if epoch == args.start_sgd_epoch and args.optimizer == 'AdamThenSGD ': # Do adam for few steps thaen continue SGD
print("-->[Switching from Adam to SGD.]")
args.lr = args.start_sgd_lr
optimizer = SGD(denoiser.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)
scheduler = StepLR(optimizer, step_size=args.lr_step_size, gamma=args.gamma)
torch.save({
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': denoiser.state_dict(),
'optimizer': optimizer.state_dict(),
}, os.path.join(args.outdir, 'checkpoint.pth.tar'))
if args.objective in ['classification', 'stability'] and test_acc > best_acc:
best_acc = test_acc
else:
continue
torch.save({
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': denoiser.state_dict(),
'optimizer': optimizer.state_dict(),
}, os.path.join(args.outdir, 'best.pth.tar'))
def main():
if not os.path.exists(args.outdir):
os.mkdir(args.outdir)
device = torch.device("cuda")
torch.cuda.set_device(args.gpu)
print(args.weight_decay)
train_dataset = get_dataset(args.dataset, 'train')
test_dataset = get_dataset(args.dataset, 'test')
pin_memory = (args.dataset == "imagenet")
train_loader = DataLoader(train_dataset,
shuffle=True,
batch_size=args.batch,
num_workers=args.workers,
pin_memory=pin_memory)
test_loader = DataLoader(test_dataset,
shuffle=False,
batch_size=args.batch,
num_workers=args.workers,
pin_memory=pin_memory)
model = get_architecture(args.arch, args.dataset, device)
logfilename = os.path.join(args.outdir, 'train_log.txt')
init_logfile(logfilename,
"epoch\ttime\tlr\ttrain loss\ttrain acc\ttestloss\ttest acc")
criterion = nn.NLLLoss().cuda()
optimizer = SGD(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
scheduler = MultiStepLR(optimizer, milestones=[80, 120], last_epoch=-1)
test_acc = []
# Training via ignite
trainer = create_supervised_trainer(model, optimizer, nn.NLLLoss(), device)
evaluator = create_supervised_evaluator(model, {
'accuracy': Accuracy(),
'nll': Loss(criterion)
}, device)
pbar = ProgressBar()
pbar.attach(trainer)
@trainer.on(Events.ITERATION_COMPLETED)
def log_training_loss(engine):
iter_in_epoch = (engine.state.iteration - 1) % len(train_loader) + 1
if engine.state.iteration % args.print_freq == 0:
pbar.log_message("Epoch[{}] Iteration[{}/{}] Loss: {:.2f}"
"".format(engine.state.epoch, iter_in_epoch,
len(train_loader), engine.state.output))
@trainer.on(Events.EPOCH_COMPLETED)
def log_epoch(engine):
scheduler.step()
evaluator.run(test_loader)
metrics = evaluator.state.metrics
avg_accuracy = metrics['accuracy']
avg_nll = metrics['nll']
test_acc.append(avg_accuracy)
if avg_accuracy >= max(test_acc):
print("Saving the model at Epoch {:}".format(engine.state.epoch +
args.epochs_warmup))
torch.save(
{
'arch': args.arch,
'state_dict': base_classifier.state_dict(),
'optimizer': optimizer.state_dict(),
}, os.path.join(args.outdir, 'checkpoint.pth.tar'))
pbar.log_message(
"Validation Results - Epoch: {} Avg accuracy: {:.4f} Avg loss: {:.2f}"
.format(engine.state.epoch, avg_accuracy, avg_nll))
log(
logfilename,
"Validation - Epoch: {} Avg accuracy: {:.4f} Avg loss: {:.3f}".
format(engine.state.epoch, avg_accuracy, avg_nll))
trainer.run(train_loader, args.epochs)
default=100,
help="output frequency")
parser.add_argument('--gpu',
default=None,
type=str,
help='id(s) for CUDA_VISIBLE_DEVICES')
args = parser.parse_args()
if __name__ == '__main__':
if args.gpu:
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
# load the base classifier
checkpoint = torch.load(args.base_classifier)
model = get_architecture(checkpoint["arch"], args.dataset)
if checkpoint["arch"] == 'resnet50' and args.dataset == "imagenet":
try:
model.load_state_dict(checkpoint['state_dict'])
except:
model = torchvision.models.resnet50(pretrained=False).cuda()
# if args.dataset == 'imagenet':
# # directly load from torchvision instead of original architecture API
# model = torchvision.models.resnet50(False).cuda()
# normalize_layer = get_normalize_layer('imagenet').cuda()
# model = torch.nn.Sequential(normalize_layer, model)
# print('loaded from torchvision for imagenet resnet50')
# else:
# model = get_architecture(checkpoint["arch"], args.dataset)
def main():
if args.gpu:
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
if not os.path.exists(args.outdir):
os.mkdir(args.outdir)
if (args.scale_down == 1 or args.dataset == "imagenet"):
train_dataset = get_dataset(args.dataset, 'train')
test_dataset = get_dataset(args.dataset, 'test')
else:
train_dataset = datasets.CIFAR10(
"./dataset_cache",
train=True,
download=True,
transform=transforms.Compose([
transforms.Resize(int(32 / args.scale_down)),
transforms.RandomCrop(int(32 / args.scale_down),
padding=int(4 / args.scale_down)),
transforms.RandomHorizontalFlip(),
transforms.ToTensor()
]))
test_dataset = datasets.CIFAR10("./dataset_cache",
train=False,
download=True,
transform=transforms.Compose([
transforms.Resize(
int(32 / args.scale_down)),
transforms.ToTensor()
]))
pin_memory = (args.dataset == "imagenet")
train_loader = DataLoader(train_dataset,
shuffle=True,
batch_size=args.batch,
num_workers=args.workers,
pin_memory=pin_memory)
test_loader = DataLoader(test_dataset,
shuffle=False,
batch_size=args.batch,
num_workers=args.workers,
pin_memory=pin_memory)
model = get_architecture(args.arch, args.dataset)
#model = torch.nn.DataParallel(model)
logfilename = os.path.join(args.outdir, 'log.txt')
init_logfile(logfilename,
"epoch\ttime\tlr\ttrain loss\ttrain acc\ttestloss\ttest acc")
criterion = CrossEntropyLoss()
optimizer = SGD(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
scheduler = StepLR(optimizer,
step_size=args.lr_step_size,
gamma=args.gamma)
for epoch in range(args.epochs):
before = time.time()
train_loss, train_acc = train(train_loader, model, criterion,
optimizer, epoch, args.noise_sd)
test_loss, test_acc = test(test_loader, model, criterion,
args.noise_sd)
after = time.time()
log(
logfilename, "{}\t{:.3}\t{:.3}\t{:.3}\t{:.3}\t{:.3}\t{:.3}".format(
epoch, str(datetime.timedelta(seconds=(after - before))),
scheduler.get_lr()[0], train_loss, train_acc, test_loss,
test_acc))
scheduler.step(epoch)
torch.save(
{
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
}, os.path.join(args.outdir, 'checkpoint.pth.tar'))
