def evaluate_ofa_random_sample(path, data_loader, batch_size=100, device='cuda:0', ensemble=False): net_acc = [] for i, id in enumerate(net_id): acc = "" for j in range(2, len(id)): if id[j] == '.': acc = id[j - 2] + id[j - 1] + id[j] + id[j + 1] net_acc.append(acc) id = np.argsort(np.array(net_acc)) new_net_id = copy.deepcopy(net_id) for i, sortid in enumerate(id): new_net_id[i] = net_id[sortid] print('new_net_id', new_net_id) n = len(net_id) best_acc = 0 acc_list = [] space = [] best_team = [] for k in range(20): nets = [] team = [] i = random.randint(0, n - 1) j = (i + random.randint(1, n - 1)) % n print('i:{} j:{}'.format(i, j)) team.append(j) team.append(i) net, image_size = ofa_specialized(net_id=new_net_id[j], pretrained=True) nets.append(net) net, image_size = ofa_specialized(net_id=new_net_id[i], pretrained=True) nets.append(net) acc = ensemble_validate(nets, path, image_size, data_loader, batch_size, device) print('net i:{} netj:{} acc:{}'.format(new_net_id[i], new_net_id[j], acc)) acc_list.append(acc) if acc > best_acc: best_acc = acc best_team = team avg_acc = np.mean(acc_list) std_acc = np.std(acc_list, ddof=1) var_acc = np.var(acc_list) print("avg{} var{} std{}".format(avg_acc, std_acc, var_acc)) print('best_random_team best_acc{}'.format(best_team, best_acc)) space.append(best_acc) print('space:{}'.format(space)) return new_net_id[best_team[0]], new_net_id[best_team[1]]
def __init__(self, data_dir=None): self.__net_id = "flops@[email protected]_finetune@75" __cwd = os.getcwd() if data_dir is not None: _p = os.path.join(data_dir, self.__class__.__module__, self.get_version()) _p = os.path.abspath(_p) os.makedirs(_p, exist_ok=True) os.chdir(_p) try: # NOTE: hard-coded download dir of ofa_specialized is under CWD. self.__model, self.__image_size = ofa_specialized(self.__net_id, pretrained=True) finally: os.chdir(__cwd) mobi = self.__model assert isinstance(mobi, MobileNetV3) def _forward__fv(x): """Patch of bound method MobileNetV3.forward to extract the feature vectors without classification layer""" x = mobi.first_conv(x) for block in mobi.blocks: x = block(x) x = mobi.final_expand_layer(x) x = mobi.global_avg_pool(x) # global average pooling x = mobi.feature_mix_layer(x) x = x.view(x.size(0), -1) # x = mobi.classifier(x) return x self.__model.forward = _forward__fv self.__model.eval() self.__dim = self.__model.feature_mix_layer.conv.out_channels # to normalize according to model trained on imagenet self.__mean = [0.485, 0.456, 0.406] self.__std = [0.229, 0.224, 0.225]
def evaluate_ofa_best_acc_team(path, data_loader, batch_size=100, device='cuda:0', ensemble=False): net_acc = [] for i, id in enumerate(net_id): acc = "" for j in range(2, len(id)): if id[j] == '.': acc = id[j - 2] + id[j - 1] + id[j] + id[j + 1] net_acc.append(acc) id = np.argsort(np.array(net_acc)) new_net_id = copy.deepcopy(net_id) for i, sortid in enumerate(id): new_net_id[i] = net_id[sortid] print('new_net_id', new_net_id) n = len(net_id) best_acc = 0 space = [] best_team = [] i = n - 1 for j in range(18, n): nets = [] team = [] team.append(j) team.append(i) net, image_size = ofa_specialized(net_id=new_net_id[j], pretrained=True) nets.append(net) net, image_size = ofa_specialized(net_id=new_net_id[i], pretrained=True) nets.append(net) acc = ensemble_validate(nets, path, image_size, data_loader, batch_size, device) print('net i:{} netj:{} acc:{}'.format(new_net_id[i], new_net_id[j], acc)) if acc > best_acc: best_acc = acc best_team = team print('space {} best_acc{}'.format(i + 1, best_acc)) space.append(best_acc) print('space:{}'.format(space)) return new_net_id[best_team[0]], new_net_id[best_team[1]]
def main(args): save_folder = args.save_folder model_folder = os.path.join(args.model_root, save_folder) makedirs(model_folder) setattr(args, 'model_folder', model_folder) logger = create_logger(model_folder, 'train', 'info') print_args(args, logger) # seed torch.manual_seed(args.seed) if torch.cuda.is_available(): torch.cuda.manual_seed(args.seed) torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False # ResNet if "ResNet" in args.model: depth_ = args.model.split('-')[1] # when it is dst, resnet block is special if args.prune_method != 'dst': p_type = None res_dict = { '18': resnet18(pretrained=args.pretrained, progress=True, prune_type=p_type), '34': resnet34(pretrained=args.pretrained, progress=True, prune_type=p_type), '50': resnet50(pretrained=args.pretrained, progress=True, prune_type=p_type), '101': resnet101(pretrained=args.pretrained, progress=True, prune_type=p_type) } net = res_dict[depth_] #elif 'efficientnet' in args.model: # net = EfficientNet.from_pretrained(args.model) elif args.model == 'efficientnet_b0': print('efficientnet-b0 load...') net = efficientnet_b0(pretrained=args.pretrained) # MobileNet elif args.model == "mobilenetv3-large-1.0": print('mobilenetv3-large-1.0') net = mobilenetv3_large_100(pretrained=args.pretrained) elif args.model == 'once-mobilenetv3-large-1.0': print('once-mobilenetv3-large-1.0') net, image_size = ofa_specialized( 'note8_lat@[email protected]_finetune@25', pretrained=args.pretrained) elif args.model == 'mobilenetv2-120d': print('mobilenetv2-120d load...') net = mobilenetv2_120d(pretrained=args.pretrained) # conv1 trainable if args.conv1_not_train: print('conv1 weight not train') if args.model == "mobilenetv3-large-1.0": for param in net.conv_stem.parameters(): param.requires_grad = False elif "ResNet" in args.model: for param in net.conv1.parameters(): param.requires_grad = False else: assert (False, 'not ready') # custom pretrain path if args.pretrain_path: print('load custom pretrain weight...') net.load_state_dict(torch.load(args.pretrain_path)) net2 = copy.deepcopy(net) # for save removed_models device = torch.device("cuda" if torch.cuda.is_available() else "cpu") net = nn.DataParallel(net) net.to(device) # KD if args.KD: print('knowledge distillation model load!') teacher_net, image_size = ofa_specialized( 'flops@[email protected]_finetune@75', pretrained=True) # 79.6% teacher_net = nn.DataParallel(teacher_net) teacher_net.to(device) # set trainer if args.KD: trainer = Trainer_KD(args, logger) else: trainer = Trainer(args, logger) # loss loss = nn.CrossEntropyLoss() # dataloader if args.model != 'once-mobilenetv3-large-1': image_size = 224 if args.dataset == 'imagenet': train_loader = torch.utils.data.DataLoader( datasets.ImageNet( '/data/imagenet/', split='train', download=False, transform=transforms.Compose([ transforms.RandomSizedCrop(image_size), transforms.RandomHorizontalFlip(), #ImageNetPolicy(), transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) ])), batch_size=args.batch_size, shuffle=True, num_workers=args.num_worker, pin_memory=True) # test_loader = torch.utils.data.DataLoader(datasets.ImageNet( '/data/imagenet/', split='val', download=False, transform=transforms.Compose([ transforms.Resize(int(image_size / 0.875)), transforms.CenterCrop(image_size), transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) ])), batch_size=args.batch_size, shuffle=False, num_workers=args.num_worker, pin_memory=True) # optimizer & scheduler optimizer = torch.optim.SGD(net.parameters(), lr=args.learning_rate, momentum=args.momentum, weight_decay=args.weight_decay) if args.scheduler == 'multistep': scheduler = torch.optim.lr_scheduler.MultiStepLR( optimizer, milestones=eval(args.multi_step_epoch), gamma=args.multi_step_gamma) elif args.scheduler == 'plateau': scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( optimizer=optimizer, mode='max', patience=3, verbose=True, factor=0.3, threshold=1e-4, min_lr=1e-6) # pruning if args.prune_method == 'global': if args.prune_type == 'group_filter': tmps = [] for n, conv in enumerate(net.modules()): if isinstance(conv, nn.Conv2d): if conv.weight.shape[1] <= 3: continue tmp_pruned = conv.weight.data.clone() original_size = tmp_pruned.size() # (out, ch, h, w) tmp_pruned = tmp_pruned.view(original_size[0], -1) # (out, inp) #append_size = 4 - tmp_pruned.shape[1] % 4 #tmp_pruned = torch.cat((tmp_pruned, tmp_pruned[:, 0:append_size]), 1) tmp_pruned = tmp_pruned.view(tmp_pruned.shape[0], -1, args.block_size) # out, -1, 4 tmp_pruned = tmp_pruned.abs().mean(2, keepdim=True).expand( tmp_pruned.shape) # out, -1, 4 tmp = tmp_pruned.flatten() tmps.append(tmp) tmps = torch.cat(tmps) num = tmps.shape[0] * (1 - args.sparsity) #sparsity 0.2 top_k = torch.topk(tmps, int(num), sorted=True) threshold = top_k.values[-1] for n, conv in enumerate(net.modules()): if isinstance(conv, nn.Conv2d): if conv.weight.shape[1] <= 3: continue tmp_pruned = conv.weight.data.clone() original_size = tmp_pruned.size() tmp_pruned = tmp_pruned.view(original_size[0], -1) #append_size = 4 - tmp_pruned.shape[1] % 4 #tmp_pruned = torch.cat((tmp_pruned, tmp_pruned[:, 0:append_size]), 1) tmp_pruned = tmp_pruned.view(tmp_pruned.shape[0], -1, args.block_size) # out, -1, 4 tmp_pruned = tmp_pruned.abs().mean(2, keepdim=True).expand( tmp_pruned.shape) # out,-1, 4 tmp_pruned = tmp_pruned.ge(threshold) tmp_pruned = tmp_pruned.view(original_size[0], -1) # out, inp #tmp_pruned = tmp_pruned[:, 0: conv.weight.data[0].nelement()] tmp_pruned = tmp_pruned.contiguous().view( original_size) # out, ch, h, w prune.custom_from_mask(conv, name='weight', mask=tmp_pruned) elif args.prune_type == 'group_channel': tmps = [] for n, conv in enumerate(net.modules()): if isinstance(conv, nn.Conv2d): if conv.weight.shape[1] <= 3: continue tmp_pruned = conv.weight.data.clone() original_size = tmp_pruned.size() # (out, ch, h, w) tmp_pruned = tmp_pruned.view(original_size[0], -1) # (out, inp) #append_size = 4 - tmp_pruned.shape[1] % 4 #tmp_pruned = torch.cat((tmp_pruned, tmp_pruned[:, 0:append_size]), 1) tmp_pruned = tmp_pruned.view( -1, args.block_size, tmp_pruned.shape[1]) # out, -1, 4 tmp_pruned = tmp_pruned.abs().mean(1, keepdim=True).expand( tmp_pruned.shape) # out, -1, 4 tmp = tmp_pruned.flatten() tmps.append(tmp) tmps = torch.cat(tmps) num = tmps.shape[0] * (1 - args.sparsity) #sparsity 0.2 top_k = torch.topk(tmps, int(num), sorted=True) threshold = top_k.values[-1] for n, conv in enumerate(net.modules()): if isinstance(conv, nn.Conv2d): if conv.weight.shape[1] <= 3: continue tmp_pruned = conv.weight.data.clone() original_size = tmp_pruned.size() tmp_pruned = tmp_pruned.view(original_size[0], -1) #append_size = 4 - tmp_pruned.shape[1] % 4 #tmp_pruned = torch.cat((tmp_pruned, tmp_pruned[:, 0:append_size]), 1) tmp_pruned = tmp_pruned.view( -1, args.block_size, tmp_pruned.shape[1]) # out, -1, 4 tmp_pruned = tmp_pruned.abs().mean(1, keepdim=True).expand( tmp_pruned.shape) # out,-1, 4 tmp_pruned = tmp_pruned.ge(threshold) #tmp_pruned = tmp_pruned.view(original_size[0], -1) # out, inp #tmp_pruned = tmp_pruned[:, 0: conv.weight.data[0].nelement()] tmp_pruned = tmp_pruned.contiguous().view( original_size) # out, ch, h, w prune.custom_from_mask(conv, name='weight', mask=tmp_pruned) elif args.prune_type == 'filter': tmps = [] for n, conv in enumerate(net.modules()): if isinstance(conv, nn.Conv2d): if conv.weight.shape[1] <= 3: continue tmp_pruned = conv.weight.data.clone() original_size = tmp_pruned.size() tmp_pruned = tmp_pruned.view(original_size[0], -1) tmp_pruned = tmp_pruned.abs().mean(1, keepdim=True).expand( tmp_pruned.shape) tmp = tmp_pruned.flatten() tmps.append(tmp) tmps = torch.cat(tmps) num = tmps.shape[0] * (1 - args.sparsity) #sparsity 0.5 top_k = torch.topk(tmps, int(num), sorted=True) threshold = top_k.values[-1] for n, conv in enumerate(net.modules()): if isinstance(conv, nn.Conv2d): if conv.weight.shape[1] <= 3: continue tmp_pruned = conv.weight.data.clone() original_size = tmp_pruned.size() tmp_pruned = tmp_pruned.view(original_size[0], -1) tmp_pruned = tmp_pruned.abs().mean(1, keepdim=True).expand( tmp_pruned.shape) tmp = tmp_pruned.flatten() tmp_pruned = tmp_pruned.ge(threshold) tmp_pruned = tmp_pruned.view(original_size[0], -1) tmp_pruned = tmp_pruned[:, 0:conv.weight.data[0].nelement()] tmp_pruned = tmp_pruned.contiguous().view(original_size) prune.custom_from_mask(conv, name='weight', mask=tmp_pruned) elif args.prune_type == 'channel': tmps = [] for n, conv in enumerate(net.modules()): if isinstance(conv, nn.Conv2d): if conv.weight.shape[1] <= 3: continue tmp_pruned = conv.weight.data.clone() original_size = tmp_pruned.size() tmp_pruned = tmp_pruned.view(original_size[0], -1) tmp_pruned = tmp_pruned.abs().mean(0, keepdim=True).expand( tmp_pruned.shape) tmp = tmp_pruned.flatten() tmps.append(tmp) tmps = torch.cat(tmps) num = tmps.shape[0] * (1 - args.sparsity) #sparsity 0.5 top_k = torch.topk(tmps, int(num), sorted=True) threshold = top_k.values[-1] for n, conv in enumerate(net.modules()): if isinstance(conv, nn.Conv2d): if conv.weight.shape[1] <= 3: continue tmp_pruned = conv.weight.data.clone() original_size = tmp_pruned.size() tmp_pruned = tmp_pruned.view(original_size[0], -1) tmp_pruned = tmp_pruned.abs().mean(0, keepdim=True).expand( tmp_pruned.shape) tmp = tmp_pruned.flatten() tmp_pruned = tmp_pruned.ge(threshold) tmp_pruned = tmp_pruned.view(original_size[0], -1) tmp_pruned = tmp_pruned[:, 0:conv.weight.data[0].nelement()] tmp_pruned = tmp_pruned.contiguous().view(original_size) prune.custom_from_mask(conv, name='weight', mask=tmp_pruned) print( f'model pruned!!(sparsity : {args.sparsity : .2f}, prune_method : {args.prune_method}, prune_type : {args.prune_type}-level pruning' ) elif args.prune_method == 'uniform': assert False, 'uniform code is not ready' elif args.prune_method == 'dst': print( f'model pruned!!(prune_method : {args.prune_method}, prune_type : {args.prune_type}-level pruning' ) elif args.prune_method == None: print('Not pruned model training started!') # Training if args.KD: trainer.train(net, teacher_net, loss, device, train_loader, test_loader, optimizer=optimizer, scheduler=scheduler) else: trainer.train(net, loss, device, train_loader, test_loader, optimizer=optimizer, scheduler=scheduler) # save removed models filename = os.path.join(args.model_folder, 'pruned_models.pth') temp = torch.load(filename) temp_dict = OrderedDict() for i in temp: if ('orig' in i): value = temp[i] * temp[i.split('_orig')[0] + '_mask'] temp_dict[i.split('module.')[1].split('_orig')[0]] = value elif 'mask' not in i: temp_dict[i.split('module.')[1]] = temp[i] net2.load_state_dict(temp_dict) save_model(net2, os.path.join(args.model_folder, 'removed_models.pth')) print('saved removed models')
metavar='NET', default='pixel1_lat@[email protected]_finetune@75', choices=specialized_network_list, help='OFA specialized networks: ' + ' | '.join(specialized_network_list) + ' (default: pixel1_lat@[email protected]_finetune@75)') 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 net, image_size = ofa_specialized(net_id=args.net, pretrained=True) args.batch_size = args.batch_size * max(len(device_list), 1) data_loader = torch.utils.data.DataLoader( datasets.ImageFolder( osp.join( args.path, 'val'), transforms.Compose( [ transforms.Resize(int(math.ceil(image_size / 0.875))), transforms.CenterCrop(image_size), transforms.ToTensor(), transforms.Normalize( mean=[ 0.485,
def evaluate_ofa_specialized(path, data_loader, batch_size=100, device='cuda:0'): def select_platform_name(): valid_platform_name = [ 'pixel1', 'pixel2', 'note10', 'note8', 's7edge', 'lg-g8', '1080ti', 'v100', 'tx2', 'cpu', 'flops' ] print( "Please select a hardware platform from ('pixel1', 'pixel2', 'note10', 'note8', 's7edge', 'lg-g8', '1080ti', 'v100', 'tx2', 'cpu', 'flops')!\n" ) while True: platform_name = input() platform_name = platform_name.lower() if platform_name in valid_platform_name: return platform_name print( "Platform name is invalid! Please select in ('pixel1', 'pixel2', 'note10', 'note8', 's7edge', 'lg-g8', '1080ti', 'v100', 'tx2', 'cpu', 'flops')!\n" ) def select_netid(platform_name): platform_efficiency_map = { 'pixel1': { 143: 'pixel1_lat@[email protected]_finetune@75', 132: 'pixel1_lat@[email protected]_finetune@75', 79: 'pixel1_lat@[email protected]_finetune@75', 58: 'pixel1_lat@[email protected]_finetune@75', 40: 'pixel1_lat@[email protected]_finetune@25', 28: 'pixel1_lat@[email protected]_finetune@25', 20: 'pixel1_lat@[email protected]_finetune@25', }, 'pixel2': { 62: 'pixel2_lat@[email protected]_finetune@25', 50: 'pixel2_lat@[email protected]_finetune@25', 35: 'pixel2_lat@[email protected]_finetune@25', 25: 'pixel2_lat@[email protected]_finetune@25', }, 'note10': { 64: 'note10_lat@[email protected]_finetune@75', 50: 'note10_lat@[email protected]_finetune@75', 41: 'note10_lat@[email protected]_finetune@75', 30: 'note10_lat@[email protected]_finetune@75', 22: 'note10_lat@[email protected]_finetune@25', 16: 'note10_lat@[email protected]_finetune@25', 11: 'note10_lat@[email protected]_finetune@25', 8: 'note10_lat@[email protected]_finetune@25', }, 'note8': { 65: 'note8_lat@[email protected]_finetune@25', 49: 'note8_lat@[email protected]_finetune@25', 31: 'note8_lat@[email protected]_finetune@25', 22: 'note8_lat@[email protected]_finetune@25', }, 's7edge': { 88: 's7edge_lat@[email protected]_finetune@25', 58: 's7edge_lat@[email protected]_finetune@25', 41: 's7edge_lat@[email protected]_finetune@25', 29: 's7edge_lat@[email protected]_finetune@25', }, 'lg-g8': { 24: 'LG-G8_lat@[email protected]_finetune@25', 16: 'LG-G8_lat@[email protected]_finetune@25', 11: 'LG-G8_lat@[email protected]_finetune@25', 8: 'LG-G8_lat@[email protected]_finetune@25', }, '1080ti': { 27: '1080ti_gpu64@[email protected]_finetune@25', 22: '1080ti_gpu64@[email protected]_finetune@25', 15: '1080ti_gpu64@[email protected]_finetune@25', 12: '1080ti_gpu64@[email protected]_finetune@25', }, 'v100': { 11: 'v100_gpu64@[email protected]_finetune@25', 9: 'v100_gpu64@[email protected]_finetune@25', 6: 'v100_gpu64@[email protected]_finetune@25', 5: 'v100_gpu64@[email protected]_finetune@25', }, 'tx2': { 96: 'tx2_gpu16@[email protected]_finetune@25', 80: 'tx2_gpu16@[email protected]_finetune@25', 47: 'tx2_gpu16@[email protected]_finetune@25', 35: 'tx2_gpu16@[email protected]_finetune@25', }, 'cpu': { 17: 'cpu_lat@[email protected]_finetune@25', 15: 'cpu_lat@[email protected]_finetune@25', 11: 'cpu_lat@[email protected]_finetune@25', 10: 'cpu_lat@[email protected]_finetune@25', }, 'flops': { 595: 'flops@[email protected]_finetune@75', 482: 'flops@[email protected]_finetune@75', 389: 'flops@[email protected]_finetune@75', } } sub_efficiency_map = platform_efficiency_map[platform_name] if not platform_name == 'flops': print( "Now, please specify a latency constraint for model specialization among", sorted(list(sub_efficiency_map.keys())), 'ms. (Please just input the number.) \n') else: print( "Now, please specify a FLOPs constraint for model specialization among", sorted(list(sub_efficiency_map.keys())), 'MFLOPs. (Please just input the number.) \n') while True: efficiency_constraint = input() if not efficiency_constraint.isdigit(): print('Sorry, please input an integer! \n') continue efficiency_constraint = int(efficiency_constraint) if not efficiency_constraint in sub_efficiency_map.keys(): print('Sorry, please choose a value from: ', sorted(list(sub_efficiency_map.keys())), '.\n') continue return sub_efficiency_map[efficiency_constraint] # platform_name = select_platform_name() # net_id = select_netid(platform_name) net_id = 'flops@[email protected]_finetune@75' print(net_id) net, image_size, net_config = ofa_specialized(net_id=net_id, pretrained=True) # validate(net, path, image_size, data_loader, batch_size, device) return net_id, net_config