def test_complexity(self, cfg_file, cx_expected):
"""Test complexity of a single model with the specified config."""
cfg_init = cfg.clone()
cfg.merge_from_file(cfg_file)
cx = net.complexity(builders.get_model())
cfg.merge_from_other_cfg(cfg_init)
self.assertEqual(cx_expected, cx)
def get_model_data(name, timings, errors):
"""Get model data for a single model."""
# Load model config
reset_cfg()
cfg.merge_from_file(model_zoo.get_config_file(name))
config_url, _, model_id, _, weight_url_full = model_zoo.get_model_info(
name)
# Get model complexity
cx = net.complexity(builders.get_model())
# Inference time is measured in ms with a reference batch_size and num_gpus
batch_size, num_gpus = 64, 1
reference = batch_size / cfg.TEST.BATCH_SIZE * cfg.NUM_GPUS / num_gpus
infer_time = timings[name]["test_fw_time"] * reference * 1000
# Training time is measured in hours for 100 epochs over the ImageNet train set
iterations = 1281167 / cfg.TRAIN.BATCH_SIZE * 100
train_time = timings[name]["train_fw_bw_time"] * iterations / 3600
# Gather all data about the model
return {
"config_url": "configs/" + config_url,
"flops": round(cx["flops"] / 1e9, 1),
"params": round(cx["params"] / 1e6, 1),
"acts": round(cx["acts"] / 1e6, 1),
"batch_size": cfg.TRAIN.BATCH_SIZE,
"infer_time": round(infer_time),
"train_time": round(train_time, 1),
"error": round(errors[name]["top1_err"], 1),
"model_id": model_id,
"weight_url": weight_url_full,
}
def test_timing(key):
"""Measure the timing of a single model."""
reset_cfg()
cfg.merge_from_file(model_zoo.get_config_file(key))
cfg.PREC_TIME.WARMUP_ITER, cfg.PREC_TIME.NUM_ITER = 5, 50
cfg.OUT_DIR, cfg.LOG_DEST = tempfile.mkdtemp(), "file"
dist.multi_proc_run(num_proc=cfg.NUM_GPUS, fun=trainer.time_model)
log_file = os.path.join(cfg.OUT_DIR, "stdout.log")
data = logging.sort_log_data(logging.load_log_data(log_file))["iter_times"]
shutil.rmtree(cfg.OUT_DIR)
return data
def test_error(key):
"""Measure the error of a single model."""
reset_cfg()
cfg.merge_from_file(model_zoo.get_config_file(key))
cfg.TEST.WEIGHTS = model_zoo.get_weights_file(key)
cfg.OUT_DIR, cfg.LOG_DEST = tempfile.mkdtemp(), "file"
dist.multi_proc_run(num_proc=cfg.NUM_GPUS, fun=trainer.test_model)
log_file = os.path.join(cfg.OUT_DIR, "stdout.log")
data = logging.sort_log_data(logging.load_log_data(log_file))["test_epoch"]
data = {"top1_err": data["top1_err"][-1], "top5_err": data["top5_err"][-1]}
shutil.rmtree(cfg.OUT_DIR)
return data
def dump_complexity():
"""Measure the complexity of every model in the configs/ directory."""
complexity = {"date-created": str(datetime.datetime.now())}
cfg_files = [os.path.join(r, f) for r, _, fs in os.walk("configs/") for f in fs]
cfg_files = sorted(f for f in cfg_files if ".yaml" in f)
for cfg_file in cfg_files:
cfg_init = cfg.clone()
cfg.merge_from_file(cfg_file)
complexity[cfg_file] = net.complexity(builders.get_model())
cfg.merge_from_other_cfg(cfg_init)
with open(_COMPLEXITY_FILE, "w") as file:
json.dump(complexity, file, sort_keys=True, indent=4)
def build_model(name, pretrained=False, cfg_list=()):
"""Constructs a predefined model (note: loads global config as well)."""
# Load the config
reset_cfg()
config_file = get_config_file(name)
cfg.merge_from_file(config_file)
cfg.merge_from_list(cfg_list)
# Construct model
model = builders.build_model()
# Load pretrained weights
if pretrained:
weights_file = get_weights_file(name)
cp.load_checkpoint(weights_file, model)
return model
def main():
# Parse cmd line args
args = parse_args()
# Load config options
cfg.merge_from_file(args.cfg_file)
cfg.merge_from_list(args.opts)
assert_cfg()
cfg.freeze()
# Perform evaluation
if cfg.NUM_GPUS > 1:
mpu.multi_proc_run(num_proc=cfg.NUM_GPUS, fun=single_proc_test)
else:
single_proc_test()
def main():
# Parse cmd line args
args = parse_args()
# Load config options
cfg.merge_from_file(args.cfg_file)
cfg.merge_from_list(args.opts)
assert_and_infer_cfg()
cfg.freeze()
# Ensure that the output dir exists
os.makedirs(cfg.OUT_DIR, exist_ok=True)
# Save the config
dump_cfg()
# Perform training
if cfg.NUM_GPUS > 1:
mpu.multi_proc_run(num_proc=cfg.NUM_GPUS, fun=single_proc_train)
else:
single_proc_train()
# Compute the predictions
preds = model(inputs)
# Compute the errors
top1_err, top5_err = mu.topk_errors(preds, labels, [1, 5])
# Combine the errors across the GPUs
# if cfg.NUM_GPUS > 1:
# top1_err = du.scaled_all_reduce([top1_err])
# #as above returns a list
# top1_err = top1_err[0]
# Copy the errors from GPU to CPU (sync point)
top1_err = top1_err.item()
# Multiply by Number of GPU's as top1_err is scaled by 1/Num_GPUs
misclassifications += top1_err * inputs.size(0) * cfg.NUM_GPUS
totalSamples += inputs.size(0) * cfg.NUM_GPUS
test_meter.iter_toc()
# Update and log stats
test_meter.update_stats(top1_err=top1_err,
mb_size=inputs.size(0) * cfg.NUM_GPUS)
test_meter.log_iter_stats(cur_epoch, cur_iter)
test_meter.iter_tic()
# Log epoch stats
test_meter.log_epoch_stats(cur_epoch)
test_meter.reset()
return misclassifications / totalSamples
if __name__ == "__main__":
cfg.merge_from_file(argparser().parse_args().cfg_file)
main(cfg)
def test_complexity(key):
"""Measure the complexity of a single model."""
reset_cfg()
cfg_file = os.path.join(_PYCLS_DIR, key)
cfg.merge_from_file(cfg_file)
return net.complexity(builders.get_model())
def main():
# load pretrained model
checkpoint = torch.load(args.checkpoint_path)
try:
model_arch = checkpoint['model_name']
patch_size = checkpoint['patch_size']
prime_size = checkpoint['patch_size']
flops = checkpoint['flops']
model_flops = checkpoint['model_flops']
policy_flops = checkpoint['policy_flops']
fc_flops = checkpoint['fc_flops']
anytime_classification = checkpoint['anytime_classification']
budgeted_batch_classification = checkpoint[
'budgeted_batch_classification']
dynamic_threshold = checkpoint['dynamic_threshold']
maximum_length = len(checkpoint['flops'])
except:
print(
'Error: \n'
'Please provide essential information'
'for customized models (as we have done '
'in pre-trained models)!\n'
'At least the following information should be Given: \n'
'--model_name: name of the backbone CNNs (e.g., resnet50, densenet121)\n'
'--patch_size: size of image patches (i.e., H\' or W\' in the paper)\n'
'--flops: a list containing the Multiply-Adds corresponding to each '
'length of the input sequence during inference')
model_configuration = model_configurations[model_arch]
if args.eval_mode > 0:
# create model
if 'resnet' in model_arch:
model = resnet.resnet50(pretrained=False)
model_prime = resnet.resnet50(pretrained=False)
elif 'densenet' in model_arch:
model = eval('densenet.' + model_arch)(pretrained=False)
model_prime = eval('densenet.' + model_arch)(pretrained=False)
elif 'efficientnet' in model_arch:
model = create_model(model_arch,
pretrained=False,
num_classes=1000,
drop_rate=0.3,
drop_connect_rate=0.2)
model_prime = create_model(model_arch,
pretrained=False,
num_classes=1000,
drop_rate=0.3,
drop_connect_rate=0.2)
elif 'mobilenetv3' in model_arch:
model = create_model(model_arch,
pretrained=False,
num_classes=1000,
drop_rate=0.2,
drop_connect_rate=0.2)
model_prime = create_model(model_arch,
pretrained=False,
num_classes=1000,
drop_rate=0.2,
drop_connect_rate=0.2)
elif 'regnet' in model_arch:
import pycls.core.model_builder as model_builder
from pycls.core.config import cfg
cfg.merge_from_file(model_configuration['cfg_file'])
cfg.freeze()
model = model_builder.build_model()
model_prime = model_builder.build_model()
traindir = args.data_url + 'train/'
valdir = args.data_url + 'val/'
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_set = datasets.ImageFolder(
traindir,
transforms.Compose([
transforms.RandomResizedCrop(
model_configuration['image_size'],
interpolation=model_configuration['dataset_interpolation']
),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(), normalize
]))
train_set_index = torch.randperm(len(train_set))
train_loader = torch.utils.data.DataLoader(
train_set,
batch_size=256,
num_workers=32,
pin_memory=False,
sampler=torch.utils.data.sampler.SubsetRandomSampler(
train_set_index[-200000:]))
val_loader = torch.utils.data.DataLoader(datasets.ImageFolder(
valdir,
transforms.Compose([
transforms.Resize(
int(model_configuration['image_size'] /
model_configuration['crop_pct']),
interpolation=model_configuration['dataset_interpolation']
),
transforms.CenterCrop(model_configuration['image_size']),
transforms.ToTensor(), normalize
])),
batch_size=256,
shuffle=False,
num_workers=16,
pin_memory=False)
state_dim = model_configuration['feature_map_channels'] * math.ceil(
patch_size / 32) * math.ceil(patch_size / 32)
memory = Memory()
policy = ActorCritic(model_configuration['feature_map_channels'],
state_dim,
model_configuration['policy_hidden_dim'],
model_configuration['policy_conv'])
fc = Full_layer(model_configuration['feature_num'],
model_configuration['fc_hidden_dim'],
model_configuration['fc_rnn'])
model = nn.DataParallel(model.cuda())
model_prime = nn.DataParallel(model_prime.cuda())
policy = policy.cuda()
fc = fc.cuda()
model.load_state_dict(checkpoint['model_state_dict'])
model_prime.load_state_dict(checkpoint['model_prime_state_dict'])
fc.load_state_dict(checkpoint['fc'])
policy.load_state_dict(checkpoint['policy'])
budgeted_batch_flops_list = []
budgeted_batch_acc_list = []
print('generate logits on test samples...')
test_logits, test_targets, anytime_classification = generate_logits(
model_prime, model, fc, memory, policy, val_loader, maximum_length,
prime_size, patch_size, model_arch)
if args.eval_mode == 2:
print('generate logits on training samples...')
dynamic_threshold = torch.zeros([39, maximum_length])
train_logits, train_targets, _ = generate_logits(
model_prime, model, fc, memory, policy, train_loader,
maximum_length, prime_size, patch_size, model_arch)
for p in range(1, 40):
print('inference: {}/40'.format(p))
_p = torch.FloatTensor(1).fill_(p * 1.0 / 20)
probs = torch.exp(torch.log(_p) * torch.range(1, maximum_length))
probs /= probs.sum()
if args.eval_mode == 2:
dynamic_threshold[p - 1] = dynamic_find_threshold(
train_logits, train_targets, probs)
acc_step, flops_step = dynamic_evaluate(test_logits, test_targets,
flops,
dynamic_threshold[p - 1])
budgeted_batch_acc_list.append(acc_step)
budgeted_batch_flops_list.append(flops_step)
budgeted_batch_classification = [
budgeted_batch_flops_list, budgeted_batch_acc_list
]
print('model_arch :', model_arch)
print('patch_size :', patch_size)
print('flops :', flops)
print('model_flops :', model_flops)
print('policy_flops :', policy_flops)
print('fc_flops :', fc_flops)
print('anytime_classification :', anytime_classification)
print('budgeted_batch_classification :', budgeted_batch_classification)
def main():
if not os.path.isdir(args.work_dirs):
mkdir_p(args.work_dirs)
record_path = args.work_dirs + '/GF-' + str(args.model_arch) \
+ '_patch-size-' + str(args.patch_size) \
+ '_T' + str(args.T) \
+ '_train-stage' + str(args.train_stage)
if not os.path.isdir(record_path):
mkdir_p(record_path)
record_file = record_path + '/record.txt'
# *create model* #
model_configuration = model_configurations[args.model_arch]
if 'resnet' in args.model_arch:
model_arch = 'resnet'
model = resnet.resnet50(pretrained=False)
model_prime = resnet.resnet50(pretrained=False)
elif 'densenet' in args.model_arch:
model_arch = 'densenet'
model = eval('densenet.' + args.model_arch)(pretrained=False)
model_prime = eval('densenet.' + args.model_arch)(pretrained=False)
elif 'efficientnet' in args.model_arch:
model_arch = 'efficientnet'
model = create_model(args.model_arch, pretrained=False, num_classes=1000,
drop_rate=0.3, drop_connect_rate=0.2)
model_prime = create_model(args.model_arch, pretrained=False, num_classes=1000,
drop_rate=0.3, drop_connect_rate=0.2)
elif 'mobilenetv3' in args.model_arch:
model_arch = 'mobilenetv3'
model = create_model(args.model_arch, pretrained=False, num_classes=1000,
drop_rate=0.2, drop_connect_rate=0.2)
model_prime = create_model(args.model_arch, pretrained=False, num_classes=1000,
drop_rate=0.2, drop_connect_rate=0.2)
elif 'regnet' in args.model_arch:
model_arch = 'regnet'
import pycls.core.model_builder as model_builder
from pycls.core.config import cfg
cfg.merge_from_file(model_configuration['cfg_file'])
cfg.freeze()
model = model_builder.build_model()
model_prime = model_builder.build_model()
fc = Full_layer(model_configuration['feature_num'],
model_configuration['fc_hidden_dim'],
model_configuration['fc_rnn'])
if args.train_stage == 1:
model.load_state_dict(torch.load(args.model_path))
model_prime.load_state_dict(torch.load(args.model_prime_path))
else:
checkpoint = torch.load(args.checkpoint_path)
model.load_state_dict(checkpoint['model_state_dict'])
model_prime.load_state_dict(checkpoint['model_prime_state_dict'])
fc.load_state_dict(checkpoint['fc'])
train_configuration = train_configurations[model_arch]
if args.train_stage != 2:
if train_configuration['train_model_prime']:
optimizer = torch.optim.SGD([{'params': model.parameters()},
{'params': model_prime.parameters()},
{'params': fc.parameters()}],
lr=0, # specify in adjust_learning_rate()
momentum=train_configuration['momentum'],
nesterov=train_configuration['Nesterov'],
weight_decay=train_configuration['weight_decay'])
else:
optimizer = torch.optim.SGD([{'params': model.parameters()},
{'params': fc.parameters()}],
lr=0, # specify in adjust_learning_rate()
momentum=train_configuration['momentum'],
nesterov=train_configuration['Nesterov'],
weight_decay=train_configuration['weight_decay'])
training_epoch_num = train_configuration['epoch_num']
else:
optimizer = None
training_epoch_num = 15
criterion = nn.CrossEntropyLoss().cuda()
model = nn.DataParallel(model.cuda())
model_prime = nn.DataParallel(model_prime.cuda())
fc = fc.cuda()
traindir = args.data_url + 'train/'
valdir = args.data_url + 'val/'
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
train_set = datasets.ImageFolder(traindir, transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
]))
train_set_index = torch.randperm(len(train_set))
train_loader = torch.utils.data.DataLoader(train_set, batch_size=256, num_workers=32, pin_memory=False,
sampler=torch.utils.data.sampler.SubsetRandomSampler(
train_set_index[:]))
val_loader = torch.utils.data.DataLoader(
datasets.ImageFolder(valdir, transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize, ])),
batch_size=train_configuration['batch_size'], shuffle=False, num_workers=32, pin_memory=False)
if args.train_stage != 1:
state_dim = model_configuration['feature_map_channels'] * math.ceil(args.patch_size / 32) * math.ceil(args.patch_size / 32)
ppo = PPO(model_configuration['feature_map_channels'], state_dim,
model_configuration['policy_hidden_dim'], model_configuration['policy_conv'])
if args.train_stage == 3:
ppo.policy.load_state_dict(checkpoint['policy'])
ppo.policy_old.load_state_dict(checkpoint['policy'])
else:
ppo = None
memory = Memory()
if args.resume:
resume_ckp = torch.load(args.resume)
start_epoch = resume_ckp['epoch']
print('resume from epoch: {}'.format(start_epoch))
model.module.load_state_dict(resume_ckp['model_state_dict'])
model_prime.module.load_state_dict(resume_ckp['model_prime_state_dict'])
fc.load_state_dict(resume_ckp['fc'])
if optimizer:
optimizer.load_state_dict(resume_ckp['optimizer'])
if ppo:
ppo.policy.load_state_dict(resume_ckp['policy'])
ppo.policy_old.load_state_dict(resume_ckp['policy'])
ppo.optimizer.load_state_dict(resume_ckp['ppo_optimizer'])
best_acc = resume_ckp['best_acc']
else:
start_epoch = 0
best_acc = 0
for epoch in range(start_epoch, training_epoch_num):
if args.train_stage != 2:
print('Training Stage: {}, lr:'.format(args.train_stage))
adjust_learning_rate(optimizer, train_configuration,
epoch, training_epoch_num, args)
else:
print('Training Stage: {}, train ppo only'.format(args.train_stage))
train(model_prime, model, fc, memory, ppo, optimizer, train_loader, criterion,
args.print_freq, epoch, train_configuration['batch_size'], record_file, train_configuration, args)
acc = validate(model_prime, model, fc, memory, ppo, optimizer, val_loader, criterion,
args.print_freq, epoch, train_configuration['batch_size'], record_file, train_configuration, args)
if acc > best_acc:
best_acc = acc
is_best = True
else:
is_best = False
save_checkpoint({
'epoch': epoch + 1,
'model_state_dict': model.module.state_dict(),
'model_prime_state_dict': model_prime.module.state_dict(),
'fc': fc.state_dict(),
'acc': acc,
'best_acc': best_acc,
'optimizer': optimizer.state_dict() if optimizer else None,
'ppo_optimizer': ppo.optimizer.state_dict() if ppo else None,
'policy': ppo.policy.state_dict() if ppo else None,
}, is_best, checkpoint=record_path)
