import os import torch import argparse import torch.nn as nn import numpy as np from tqdm.auto import tqdm from torch.utils.data import DataLoader import torch.nn as nn import torch import random from torch import nn from torch.optim import * from torch.optim.lr_scheduler import * from torch.utils.data import DataLoader from torchprofile import profile_macs from torchvision.datasets import * from torchvision.transforms import * from proard.model_zoo import DYN_net from proard.nas.accuracy_predictor import AccuracyPredictor,ResNetArchEncoder,RobustnessPredictor,MobileNetArchEncoder,Accuracy_Robustness_Predictor from proard.nas.efficiency_predictor import ResNet50FLOPsModel,Mbv3FLOPsModel,ProxylessNASFLOPsModel from proard.nas.search_algorithm import EvolutionFinder,DynIndividual_mbv,DynIndividual_res,DynRandomSampler,DynProblem_mbv,DynProblem_res,DynSampling,individual_to_arch_res,individual_to_arch_mbv from utils.profile import trainable_param_num from pymoo.core.individual import Individual from pymoo.core.mutation import Mutation from pymoo.core.population import Population from pymoo.core.problem import Problem from pymoo.core.sampling import Sampling from pymoo.core.variable import Choice from pymoo.operators.crossover.ux import UniformCrossover from pymoo.operators.mutation.pm import PolynomialMutation from pymoo.operators.mutation.rm import ChoiceRandomMutation from pymoo.operators.selection.rnd import RandomSelection from pymoo.operators.selection.tournament import TournamentSelection from pymoo.algorithms.moo.nsga2 import NSGA2 from pymoo.algorithms.moo.sms import SMSEMOA from pymoo.algorithms.moo.spea2 import SPEA2 from pymoo.optimize import minimize from pymoo.termination import get_termination from pymoo.termination.default import DefaultMultiObjectiveTermination from pymoo.core.callback import Callback from pymoo.util.display.column import Column from pymoo.util.display.output import Output from proard.classification.run_manager import ClassificationRunConfig, RunManager parser = argparse.ArgumentParser() parser.add_argument( "-p", "--path", help="The path of cifar10", type=str, default="/dataset/cifar10" ) parser.add_argument("-g", "--gpu", help="The gpu(s) to use", type=str, default="all") parser.add_argument( "-b", "--batch-size", help="The batch on every device for validation", type=int, default=100, ) parser.add_argument("-j", "--workers", help="Number of workers", type=int, default=20) parser.add_argument( "-n", "--net", metavar="DYNNET", default="ResNet50", choices=[ "ResNet50", "MBV3", "ProxylessNASNet", ], help="dynamic networks", ) parser.add_argument( "--dataset", type=str, default="cifar10" ,choices=["cifar10", "cifar100", "imagenet"] ) parser.add_argument( "--attack", type=str, default="linf-pgd" ,choices=['fgsm', 'linf-pgd', 'fgm', 'l2-pgd', 'linf-df', 'l2-df', 'linf-apgd', 'l2-apgd','squar_attack','autoattack','apgd_ce'] ) parser.add_argument("--train_criterion", type=str, default="trades",choices=["trades","sat","mart","hat"]) parser.add_argument( "--robust_mode", type=bool, default=True ) args = parser.parse_args() if args.gpu == "all": device_list = range(torch.cuda.device_count()) args.gpu = ",".join(str(_) for _ in device_list) else: device_list = [int(_) for _ in args.gpu.split(",")] os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu args.batch_size = args.batch_size * max(len(device_list), 1) run_config = ClassificationRunConfig(attack_type=args.attack, dataset= args.dataset, test_batch_size=args.batch_size, n_worker=args.workers,robust_mode=args.robust_mode) dyn_network = DYN_net(args.net,args.robust_mode,args.dataset,args.train_criterion, pretrained=True,run_config=run_config) if args.net == "ResNet50": efficiency_predictor = ResNet50FLOPsModel(dyn_network) arch = ResNetArchEncoder(image_size_list=[32],depth_list=[0,1,2],expand_list=[0.2,0.25,0.35],width_mult_list=[0.65,0.8,1.0]) accuracy_robustness_predictor = Accuracy_Robustness_Predictor(arch) accuracy_robustness_predictor.load_state_dict(torch.load("./acc_rob_data_{}_{}_{}/src/model_acc_rob.pth".format(args.dataset,args.net,args.train_criterion))) elif args.net == "MBV3": efficiency_predictor = Mbv3FLOPsModel(dyn_network) arch = MobileNetArchEncoder(image_size_list=[32],depth_list=[2,3,4],expand_list=[3,4,6],ks_list=[3,5,7]) accuracy_robustness_predictor = Accuracy_Robustness_Predictor(arch) accuracy_robustness_predictor.load_state_dict(torch.load("./acc_rob_data_{}_{}_{}/src/model_acc_rob.pth".format(args.dataset,args.net,args.train_criterion))) elif args.net == "ProxylessNASNet": efficiency_predictor = ProxylessNASFLOPsModel(dyn_network) arch = MobileNetArchEncoder(image_size_list=[32],depth_list=[2,3,4],expand_list=[3,4,6],width_mult_list=[3,5,7]) accuracy_robustness_predictor = Accuracy_Robustness_Predictor(arch) accuracy_robustness_predictor.load_state_dict(torch.load("./acc_rob_data_{}_{}_{}/src/model_acc_rob.pth".format(args.dataset,args.net,args.train_criterion))) ##### Test ################################################# dyn_sampler = DynRandomSampler(arch, efficiency_predictor) # arch1, eff1 = dyn_sampler.random_sample() # arch2, eff2 = dyn_sampler.random_sample() # print(accuracy_predictor.predict_acc([arch1, arch2])) # print(arch1,eff1) ################################################## """ Hyperparameters - P: size of the population in each generation (number of individuals) - N: number of generations to run the algorithm - mutate_prob: probability of gene mutation in the evolutionary search """ P = 100 N = 100 mutation_prob = 0.5 # variables options if args.net == 'ResNet50': search_space = { 'e': [0.2, 0.25, 0.35], 'd': [0, 1, 2], 'w': [0 ,1 ,2], 'image_size': [32] } else: search_space = { 'ks': [3, 5, 7], 'e': [3, 4, 6], 'd': [2, 3, 4], 'image_size': [32] } #---------------------------- # units num_blocks = arch.max_n_blocks num_stages = arch.n_stage Flops_constraints = 1600 if args.net == "ResNet50": problem = DynProblem_res(efficiency_predictor, accuracy_robustness_predictor, num_blocks, num_stages, search_space,Flops_constraints) else: problem = DynProblem_mbv(efficiency_predictor, accuracy_robustness_predictor, num_blocks, num_stages, search_space,Flops_constraints) mutation_rc = ChoiceRandomMutation(prob=1.0, prob_var=0.1) crossover_ux = UniformCrossover(prob=1.0) # selection_tournament = TournamentSelection( # func_comp=accuracy_predictor.predict_acc, # pressure=2 # ) termination_default = DefaultMultiObjectiveTermination( xtol=1e-8, cvtol=1e-6, ftol=0.0025, period=30, n_max_gen=1000, n_max_evals=100000 ) termination_gen = get_termination("n_gen", N) np.random.seed(42) random.seed(42) if args.net=="ResNet50": init_pop = Population(individuals=[DynIndividual_res(dyn_sampler.random_sample(), accuracy_robustness_predictor) for _ in range(P)]) else: init_pop = Population(individuals=[DynIndividual_mbv(dyn_sampler.random_sample(), accuracy_robustness_predictor) for _ in range(P)]) algorithm = NSGA2( pop_size=P, sampling=DynSampling(), # selection=selection_tournament, crossover=crossover_ux, mutation=mutation_rc, # mutation=mutation_pm, # survival=RankAndCrowdingSurvival(), # output=MultiObjectiveOutput(), # **kwargs ) res_nsga2 = minimize( problem, algorithm, termination=termination_gen, seed=1, #verbose=True, verbose=False, save_history=True, ) # print(100-res_nsga2.history[99].pop.get('F')[:,0],100-res_nsga2.history[99].pop.get('F')[:,1]) # a = individual_to_arch_res(res_nsga2.pop.get('X'),num_blocks)[0] # # print(a) # # a['d'][3] = int(a['d'][3]) # a['d'][4] = int(a['d'][4]) # dyn_network.set_active_subnet(**a) # subnet = dyn_network.get_active_subnet(preserve_weight=True) # run_manager = RunManager(".tmp/eval_subnet", subnet, run_config, init=False) # run_config.data_provider.assign_active_img_size(32) # run_manager.reset_running_statistics(net=subnet) # print("Test random subnet:") # # print(subnet.module_str) # loss, (top1, top5,robust1,robust5) = run_manager.validate(net=subnet,is_test=True) # print("Results: loss=%.5f,\t top1=%.1f,\t top5=%.1f,\t robust1=%.1f,\t robust5=%.1f" % (loss, top1, top5,robust1,robust5)) np.savetxt("./results/acc_gen0.csv", 100-res_nsga2.history[0].pop.get('F')[:,0], delimiter=",") np.savetxt("./results/acc_gen99.csv", 100-res_nsga2.history[99].pop.get('F')[:,0], delimiter=",") np.savetxt("./results/rob_gen0.csv", 100-res_nsga2.history[0].pop.get('F')[:,1], delimiter=",") np.savetxt("./results/rob_gen99.csv", 100-res_nsga2.history[99].pop.get('F')[:,1], delimiter=",") np.savetxt("./results/flops_gen99.csv", res_nsga2.history[99].pop.get('G'), delimiter=",") # np.savetxt("./results/robs.csv", np.array(robs), delimiter=",") from matplotlib import pyplot as plt from matplotlib.ticker import FormatStrFormatter from matplotlib.ticker import AutoMinorLocator, MultipleLocator # NSGA-II population progression x_min, x_max, y_min, y_max = 80, 93, 47, 56 ax_limits = [x_min, x_max, y_min, y_max] #------------------------------------------------- # plot fig, ax = plt.subplots(dpi=600) gen0 = 0 gen1 = 99 print(100-res_nsga2.history[gen1].pop.get('F')[:,0], 100 - res_nsga2.history[gen1].pop.get('F')[:,1]) # gen2 = 99 # print(res_nsga2.history[gen0].pop.get('F')[:,0],res_nsga2.history[gen0].pop.get('F')[:,1] ) ax.plot(100-res_nsga2.history[gen0].pop.get('F')[:,0], 100 - res_nsga2.history[gen0].pop.get('F')[:,1] , 'o', label=f'Population at generation #{gen0+1}', color='red', alpha=0.5) ax.plot(100-res_nsga2.history[gen1].pop.get('F')[:,0], 100 - res_nsga2.history[gen1].pop.get('F')[:,1] , 'o', label=f'Population at generation #{gen1+1}', color='green', alpha=0.5) # ax.plot(res_nsga2.history[gen2].pop.get('F')[:,0], 100 - res_nsga2.history[gen2].pop.get('F')[:,1], 'o', label=f'Population at generation #{gen2+1}', color='orange', alpha=0.5) # ax.plot(res_nsga2.history[gen3].pop.get('F')[:,0], 100 - res_nsga2.history[gen3].pop.get('F')[:,1], 'o', label=f'Population at generation #{gen3+1}', color='blue', alpha=0.5) #------------------------------------------------- # text ax.grid(True, linestyle=':') ax.set_xlabel('Accuracy (%)') ax.set_ylabel('Robustness (%)') ax.set_title('NSGA-II solutions progression For Fixed number of FLOPs'), ax.legend() #------------------------------------------------- # x-axis ax.xaxis.set_major_locator(MultipleLocator(1)) ax.xaxis.set_minor_locator(MultipleLocator(1)) # y-axis ax.yaxis.set_major_locator(MultipleLocator(1)) ax.yaxis.set_minor_locator(MultipleLocator(1)) # ax.yaxis.set_major_formatter(FormatStrFormatter('%.1f')) ax.set(xlim=(ax_limits[0], ax_limits[1]), ylim=(ax_limits[2], ax_limits[3])) #------------------------------------------------- plt.savefig('nsga2_pop_progression_debug.png') fig.set_dpi(100) # plt.close(fig) # plt.show() # finder = EvolutionFinder(efficiency_predictor,accuracy_predictor,Robustness_predictor) # valid_constraint_range = 800 # best_valids, best_info = finder.run_evolution_search(constraint=valid_constraint_range,verbose=True) # print(efficiency_predictor.get_efficiency(best_info[2])) # dyn_network.set_active_subnet(best_info[2]['d'],best_info[2]['e'],best_info[2]['w']) # subnet = dyn_network.get_active_subnet(preserve_weight=True) # run_config = CifarRunConfig_robust(test_batch_size=args.batch_size, n_worker=args.workers) # run_manager = RunManager_robust(".tmp/eval_subnet", subnet, run_config, init=False) # run_config.data_provider.assign_active_img_size(32) # run_manager.reset_running_statistics(net=subnet) # loss, (top1, top5,robust1,robust5) = run_manager.validate(net=subnet) # print("Results: loss=%.5f,\t top1=%.1f,\t top5=%.1f,\t robust1=%.1f,\t robust5=%.1f" % (loss, top1, top5,robust1,robust5)) # print("number of parameter={}M".format(trainable_param_num(subnet)))