def ofa_specialized(net_id, pretrained=True): url_base = 'https://hanlab.mit.edu/files/OnceForAll/ofa_specialized/' net_config = json.load( open( download_url(url_base + net_id + '/net.config', model_dir='.torch/ofa_specialized/%s/' % net_id))) if net_config['name'] == ProxylessNASNets.__name__: net = ProxylessNASNets.build_from_config(net_config) elif net_config['name'] == MobileNetV3.__name__: net = MobileNetV3.build_from_config(net_config) else: raise ValueError('Not supported network type: %s' % net_config['name']) image_size = json.load( open( download_url(url_base + net_id + '/run.config', model_dir='.torch/ofa_specialized/%s/' % net_id)))['image_size'] if pretrained: init = torch.load(download_url(url_base + net_id + '/init', model_dir='.torch/ofa_specialized/%s/' % net_id), map_location='cpu')['state_dict'] net.load_state_dict(init) return net, image_size
def ofa_net(net_id, pretrained=True): if net_id == 'ofa_proxyless_d234_e346_k357_w1.3': net = OFAProxylessNASNets( dropout_rate=0, width_mult=1.3, ks_list=[3, 5, 7], expand_ratio_list=[3, 4, 6], depth_list=[2, 3, 4], ) elif net_id == 'ofa_mbv3_d234_e346_k357_w1.0': net = OFAMobileNetV3( dropout_rate=0, width_mult=1.0, ks_list=[3, 5, 7], expand_ratio_list=[3, 4, 6], depth_list=[2, 3, 4], ) elif net_id == 'ofa_mbv3_d234_e346_k357_w1.2': net = OFAMobileNetV3( dropout_rate=0, width_mult=1.2, ks_list=[3, 5, 7], expand_ratio_list=[3, 4, 6], depth_list=[2, 3, 4], ) elif net_id == 'ofa_resnet50': net = OFAResNets( dropout_rate=0, depth_list=[0, 1, 2], expand_ratio_list=[0.2, 0.25, 0.35], width_mult_list=[0.65, 0.8, 1.0] ) net_id = 'ofa_resnet50_d=0+1+2_e=0.2+0.25+0.35_w=0.65+0.8+1.0' else: raise ValueError('Not supported: %s' % net_id) if pretrained: url_base = 'https://hanlab.mit.edu/files/OnceForAll/ofa_nets/' init = torch.load( download_url(url_base + net_id, model_dir='.torch/ofa_nets'), map_location='cpu')['state_dict'] net.load_state_dict(init) return net
def __init__(self, local_dir='~/.ofa/latency_tools/', url='https://hanlab.mit.edu/files/proxylessNAS/LatencyTools/mobile_trim.yaml'): if url.startswith('http'): fname = download_url(url, local_dir, overwrite=True) else: fname = url with open(fname, 'r') as fp: self.lut = yaml.load(fp)
def proxylessnas_net(net_id, pretrained=True): net = proxyless_base( net_config='https://hanlab.mit.edu/files/proxylessNAS/%s.config' % net_id, ) if pretrained: net.load_state_dict(torch.load( download_url('https://hanlab.mit.edu/files/proxylessNAS/%s.pth' % net_id), map_location='cpu' )['state_dict']) return net
def ofa_specialized(net_id, pretrained=True): url_base = 'https://hanlab.mit.edu/files/OnceForAll/ofa_specialized/' net_config = json.load(open( download_url(url_base + net_id + '/net.config', model_dir='.torch/ofa_specialized/%s/' % net_id) )) net = get_net_by_name(net_config['name']).build_from_config(net_config) image_size = json.load(open( download_url(url_base + net_id + '/run.config', model_dir='.torch/ofa_specialized/%s/' % net_id) ))['image_size'] if pretrained: init = torch.load( download_url(url_base + net_id + '/init', model_dir='.torch/ofa_specialized/%s/' % net_id), map_location='cpu' )['state_dict'] net.load_state_dict(init) return net, image_size
def __init__(self, pretrained=True, device='cuda:0'): self.device = device self.model = nn.Sequential( nn.Linear(128, 400), nn.ReLU(), nn.Linear(400, 400), nn.ReLU(), nn.Linear(400, 400), nn.ReLU(), nn.Linear(400, 1), ) if pretrained: # load pretrained model fname = download_url( "https://hanlab.mit.edu/files/OnceForAll/tutorial/acc_predictor.pth" ) self.model.load_state_dict( torch.load(fname, map_location=torch.device('cpu'))) self.model = self.model.to(self.device)
def ofa_net(net_id, pretrained=True): from ofa.elastic_nn.modules.dynamic_op import DynamicSeparableConv2d from ofa.elastic_nn.networks import OFAMobileNetV3, OFAProxylessNASNets DynamicSeparableConv2d.KERNEL_TRANSFORM_MODE = 1 if net_id == 'ofa_proxyless_d234_e346_k357_w1.3': net = OFAProxylessNASNets( dropout_rate=0, width_mult_list=1.3, ks_list=[3, 5, 7], expand_ratio_list=[3, 4, 6], depth_list=[2, 3, 4], ) elif net_id == 'ofa_mbv3_d234_e346_k357_w1.0': net = OFAMobileNetV3( dropout_rate=0, width_mult_list=1.0, ks_list=[3, 5, 7], expand_ratio_list=[3, 4, 6], depth_list=[2, 3, 4], ) elif net_id == 'ofa_mbv3_d234_e346_k357_w1.2': net = OFAMobileNetV3( dropout_rate=0, width_mult_list=1.2, ks_list=[3, 5, 7], expand_ratio_list=[3, 4, 6], depth_list=[2, 3, 4], ) else: raise ValueError('Not supported: %s' % net_id) if pretrained: url_base = 'https://hanlab.mit.edu/files/OnceForAll/ofa_nets/' init = torch.load(download_url(url_base + net_id, model_dir='.torch/ofa_nets'), map_location='cpu')['state_dict'] net.load_state_dict(init) return net
def proxyless_base(net_config=None, n_classes=None, bn_param=None, dropout_rate=None, local_path='~/.torch/proxylessnas/'): assert net_config is not None, 'Please input a network config' if 'http' in net_config: net_config_path = download_url(net_config, local_path) else: net_config_path = net_config net_config_json = json.load(open(net_config_path, 'r')) if n_classes is not None: net_config_json['classifier']['out_features'] = n_classes if dropout_rate is not None: net_config_json['classifier']['dropout_rate'] = dropout_rate net = ProxylessNASNets.build_from_config(net_config_json) if bn_param is not None: net.set_bn_param(*bn_param) return net
# we will skip ImageNet evaluation by default** since it will be very slow. # If you are using the GPU, in case you don't have the full dataset, # we will download a subset of ImageNet which contains 2,000 images (~250M) for testing. # If you do have the full ImageNet dataset on your machine, just specify it in `imagenet_data_path` and the downloading script will be skipped. # %% if cuda_available: # path to the ImageNet dataset print("Please input the path to the ImageNet dataset.\n") imagenet_data_path = './imagenet_1k' #input() # if 'imagenet_data_path' is empty, download a subset of ImageNet containing 2000 images (~250M) for test if not os.path.isdir(imagenet_data_path): os.makedirs(imagenet_data_path, exist_ok=True) download_url( 'https://hanlab.mit.edu/files/OnceForAll/ofa_cvpr_tutorial/imagenet_1k.zip', model_dir='data') os.system(' cd data && unzip imagenet_1k 1>/dev/null && cd ..') os.system(' cp -r data/imagenet_1k/* $imagenet_data_path' ) #! pc: copy not done properly os.system(' rm -rf data') print('%s is empty. Download a subset of ImageNet for test.' % imagenet_data_path) print('The ImageNet dataset files are ready.') else: print( 'Since GPU is not found in the environment, we skip all scripts related to ImageNet evaluation.' ) # %% [markdown]
def supporting_elastic_expand(train_func, run_manager, args, validate_func_dict): dynamic_net = run_manager.net if isinstance(dynamic_net, nn.DataParallel): dynamic_net = dynamic_net.module # load stage info stage_info_path = os.path.join(run_manager.path, 'expand.stage') try: stage_info = json.load(open(stage_info_path)) except Exception: stage_info = {'stage': 0} # load pretrained models validate_func_dict['expand_ratio_list'] = sorted( dynamic_net.expand_ratio_list) if args.phase == 1: model_path = download_url( 'https://hanlab.mit.edu/files/OnceForAll/ofa_checkpoints/ofa_D234_E6_K357', model_dir='.torch/ofa_checkpoints/%d' % hvd.rank()) load_models(run_manager, dynamic_net, model_path=model_path) else: model_path = download_url( 'https://hanlab.mit.edu/files/OnceForAll/ofa_checkpoints/ofa_D234_E46_K357', model_dir='.torch/ofa_checkpoints/%d' % hvd.rank()) load_models(run_manager, dynamic_net, model_path=model_path) dynamic_net.re_organize_middle_weights() run_manager.write_log( '%.3f\t%.3f\t%.3f\t%s' % validate(run_manager, **validate_func_dict), 'valid') expand_stage_list = dynamic_net.expand_ratio_list.copy() expand_stage_list.sort(reverse=True) n_stages = len(expand_stage_list) - 1 start_stage = n_stages - 1 for current_stage in range(start_stage, n_stages): run_manager.write_log( '-' * 30 + 'Supporting Elastic Expand Ratio: %s -> %s' % (expand_stage_list[:current_stage + 1], expand_stage_list[:current_stage + 2]) + '-' * 30, 'valid') # add expand list constraints supported_expand = expand_stage_list[:current_stage + 2] if len(set(dynamic_net.ks_list)) == 1 and len( set(dynamic_net.depth_list)) == 1: validate_func_dict['expand_ratio_list'] = supported_expand else: validate_func_dict['expand_ratio_list'] = sorted( {min(supported_expand), max(supported_expand)}) dynamic_net.set_constraint(supported_expand, constraint_type='expand_ratio') # train train_func( run_manager, args, lambda _run_manager, epoch, is_test: validate( _run_manager, epoch, is_test, **validate_func_dict)) # next stage & reset stage_info['stage'] += 1 run_manager.start_epoch = 0 run_manager.best_acc = 0.0 dynamic_net.re_organize_middle_weights( expand_ratio_stage=stage_info['stage']) if isinstance(run_manager, DistributedRunManager): run_manager.broadcast() # save and validate run_manager.save_model(model_name='expand_stage%d.pth.tar' % stage_info['stage']) json.dump(stage_info, open(stage_info_path, 'w'), indent=4) validate_func_dict['expand_ratio_list'] = sorted( dynamic_net.expand_ratio_list) run_manager.write_log( '%.3f\t%.3f\t%.3f\t%s' % validate(run_manager, **validate_func_dict), 'valid')
args.independent_distributed_sampling = False args.kd_ratio = 1.0 args.kd_type = 'ce' if __name__ == '__main__': os.makedirs(args.path, exist_ok=True) # Initialize Horovod hvd.init() # Pin GPU to be used to process local rank (one GPU per process) torch.cuda.set_device(hvd.local_rank()) args.teacher_path = download_url( 'https://hanlab.mit.edu/files/OnceForAll/ofa_checkpoints/ofa_D4_E6_K7', model_dir='.torch/ofa_checkpoints/%d' % hvd.rank() ) num_gpus = hvd.size() torch.manual_seed(args.manual_seed) torch.cuda.manual_seed_all(args.manual_seed) np.random.seed(args.manual_seed) random.seed(args.manual_seed) # image size args.image_size = [int(img_size) for img_size in args.image_size.split(',')] if len(args.image_size) == 1: args.image_size = args.image_size[0] MyRandomResizedCrop.CONTINUOUS = args.continuous_size MyRandomResizedCrop.SYNC_DISTRIBUTED = not args.not_sync_distributed_image_size
args.kd_ratio = -1.0 # not using teacher model args.teacher_model = None else: args.dy_conv_scaling_mode = 1 args.kd_ratio = 1.0 if __name__ == "__main__": os.makedirs(args.path, exist_ok=True) # Initialize Horovod hvd.init() # Pin GPU to be used to process local rank (one GPU per process) torch.cuda.set_device(hvd.local_rank()) args.teacher_path = download_url( "https://hanlab.mit.edu/files/OnceForAll/ofa_checkpoints/ofa_D4_E6_K7", model_dir=".torch/ofa_checkpoints/%d" % hvd.rank(), ) num_gpus = hvd.size() torch.manual_seed(args.manual_seed) torch.cuda.manual_seed_all(args.manual_seed) np.random.seed(args.manual_seed) random.seed(args.manual_seed) # image size args.image_size = [ int(img_size) for img_size in args.image_size.split(",") ] if len(args.image_size) == 1: args.image_size = args.image_size[0]
if args.net == 'proxyless_mobile': net = proxylessnas_mobile(pretrained=False) LiteResidualModule.insert_lite_residual( net, args.lite_residual_downsample, 'bilinear', args.lite_residual_expand, args.lite_residual_ks, 'relu', args.lite_residual_groups, ) # replace bn layers with gn layers replace_bn_with_gn(net, gn_channel_per_group=8) # load pretrained model init_file = download_url( 'https://hanlab.mit.edu/projects/tinyml/tinyTL/files/' 'proxylessnas_mobile+lite_residual@imagenet@ws+gn', model_dir='~/.tinytl/') net.load_state_dict( torch.load(init_file, map_location='cpu')['state_dict']) net.classifier = LinearLayer(net.classifier.in_features, run_config.data_provider.n_classes, dropout_rate=args.dropout) classification_head.append(net.classifier) init_models(classification_head) else: if args.net_path is not None: net_config_path = os.path.join(args.net_path, 'net.config') init_path = os.path.join(args.net_path, 'init') else: base_url = 'https://hanlab.mit.edu/projects/tinyml/tinyTL/files/specialized/%s/' % args.dataset net_config_path = download_url(