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(
