def main(): global args, iters global file args = parser.parse_args() args.gpu = [int(i) for i in args.gpu.split(',')] torch.cuda.set_device(args.gpu[0] if args.gpu else None) torch.backends.cudnn.benchmark = True L_cls_f = nn.CrossEntropyLoss().cuda() ## Dataset Loading normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) train_loader = torch.utils.data.DataLoader(datasets.CIFAR10( root='./cifar10', train=True, transform=transforms.Compose([ transforms.RandomHorizontalFlip(), transforms.RandomCrop(32, 4), transforms.ToTensor(), normalize ]), download=True), batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True) val_loader = torch.utils.data.DataLoader(datasets.CIFAR10( root='./cifar10', train=False, transform=transforms.Compose([transforms.ToTensor(), normalize])), batch_size=128, shuffle=False, num_workers=args.workers, pin_memory=True) ## Model Initialize and Loading model = resnet.resnet32() args.checkpoint = 'models/model_32.th' model = nn.DataParallel(model, device_ids=args.gpu).cuda() checkpoint = torch.load(args.checkpoint, map_location='cuda:0') model.load_state_dict(checkpoint['state_dict']) model.eval() original_model = copy.deepcopy(model) loss, init_acc = validate(val_loader, model, L_cls_f, '') print('\nOriginal performance. Acc: {:2.2f}%'.format(init_acc)) num_params = get_params(model) num_flops = get_flops(model) ## 1. Initialization process layer_names = get_layer_names(model) expected_flops(copy.deepcopy(model), layer_names[1:], num_params, num_flops, args.ratio) add_1x1_convs(model, layer_names[1:]) print('== 1. Initialization fine-tuning stage. ') model_opt = torch.optim.SGD(model.parameters(), lr=1e-3, momentum=0.9, weight_decay=1e-4) for epochs in range(20): fine_tuning(model, original_model, train_loader, val_loader, L_cls_f, model_opt, False) loss, acc = validate(val_loader, model, L_cls_f, '* ') print("[Init {:02d}] Loss: {:.3f}. Acc: {:2.2f}%. || Param: {:2.2f}% Flop: {:2.2f}%".format(epochs+1, loss, acc, \ get_params(model)/num_params*100, get_flops(model)/num_flops*100)) ## 2. Pruning process, from top to bottom print('\n== 2. Pruning stage. ') for i in range(1, len(layer_names)): index = len(layer_names) - i model = pruning_output_channel(model, original_model, layer_names[index], train_loader, val_loader, L_cls_f) model = pruning_input_channel(model, original_model, layer_names[index], train_loader, val_loader, L_cls_f) loss, acc = validate(val_loader, model, L_cls_f, '* ') print("[Pruning {:02d}]. Loss: {:.3f}. Acc: {:2.2f}%. || Param: {:2.2f}% Flop: {:2.2f}%".format(index, loss, \ acc, get_params(model)/num_params*100, get_flops(model)/num_flops*100)) ## 3. Final Fine-tuning stage print('\n==3. Final fine-tuning stage after pruning.') best_acc = 0 model_opt = torch.optim.SGD(model.parameters(), lr=1e-2, momentum=0.9, weight_decay=1e-4) for epochs in range(args.step_ft): adjust_learning_rate(model_opt, epochs, args.step_ft) fine_tuning(model, original_model, train_loader, val_loader, L_cls_f, model_opt) loss, acc = validate(val_loader, model, L_cls_f, '* ') if acc > best_acc: best_acc = acc print("[Fine-tune {:03d}] Loss: {:.3f}. Acc: {:2.2f}%. || Param: {:2.2f}% Flop: {:2.2f}% Best: {:2.2f}%".format(epochs+1, loss, acc, \ get_params(model)/num_params*100, get_flops(model)/num_flops*100, best_acc)) print("\n[Final] Baseline: {:2.2f}%. After Pruning: {:2.2f}%. || Diff: {:2.2f}% Param: {:2.2f}% Flop: {:2.2f}%".format(\ init_acc, best_acc, init_acc - best_acc, get_params(model)\ /num_params*100, get_flops(model)/num_flops*100))
from model import SequentialImageNetwork, SequentialImageNetworkMod from util import * from datasets import * import re import sys name = sys.argv[1] model_flag = name.split("-")[0] train_flag = name.split("-")[1] if model_flag == "r32p": import resnet model = SequentialImageNetworkMod(resnet.resnet32()).cuda() elif model_flag == "r18": from pytorch_cifar.models import resnet model = SequentialImageNetwork(resnet.ResNet18()).cuda() else: raise NotImplementedError eps = int(re.search(r"[0-9]+$", name).group()) poisoner_flag = name.split("-")[3][:3] clean_label = int(name.split("-")[2][0]) target_label = int(name.split("-")[2][1]) print(f"{model_flag=} {clean_label=} {target_label=} {poisoner_flag=} {eps=}") if len(sys.argv) > 2:
import time import torch import torch.nn as nn import torch.nn.functional as F import torch.nn.init as init import torchvision from torch.autograd import Variable from resnet import resnet32 from collections import OrderedDict import torchvision.datasets as datasets import torchvision.transforms as transforms from custom_loader import val_set # %% device = "cuda" model = resnet32().to(device) # print(model) train_loader = torch.utils.data.DataLoader(val_set, batch_size=128, shuffle=True, num_workers=4, pin_memory=True) # %% saved_model_path = '/root/Adversarial-attacks-DNN-18786/saved_model/resnet32-adv' criterion = nn.CrossEntropyLoss().cuda() optimizer = torch.optim.SGD(model.parameters(), 0.1, momentum=0.9, weight_decay=1e-4)