# 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'))