def get_model(args):
print('=> Building model..')
if args.dataset == 'imagenet':
n_class = 1000
elif args.dataset == 'cifar10':
n_class = 10
else:
raise NotImplementedError
if args.model_type == 'mobilenet':
net = MobileNet(n_class=n_class)
elif args.model_type == 'mobilenetv2':
net = MobileNetV2(n_class=n_class)
elif args.model_type.startswith('resnet'):
net = resnet.__dict__[args.model_type](pretrained=True)
in_features = net.fc.in_features
net.fc = nn.Linear(in_features, n_class)
else:
raise NotImplementedError
if args.ckpt_path is not None:
# the checkpoint can be state_dict exported by amc_search.py or saved by amc_train.py
print('=> Loading checkpoint {} ..'.format(args.ckpt_path))
net.load_state_dict(torch.load(args.ckpt_path, torch.device('cpu')))
if args.mask_path is not None:
SZ = 224 if args.dataset == 'imagenet' else 32
data = torch.randn(2, 3, SZ, SZ)
ms = ModelSpeedup(net, data, args.mask_path, torch.device('cpu'))
ms.speedup_model()
net.to(args.device)
if torch.cuda.is_available() and args.n_gpu > 1:
net = torch.nn.DataParallel(net, list(range(args.n_gpu)))
return net
def get_model(args):
print('=> Building model..')
if args.dataset == 'imagenet':
n_class = 1000
elif args.dataset == 'cifar10':
n_class = 10
else:
raise NotImplementedError
if args.model_type == 'mobilenet':
net = MobileNet(n_class=n_class).cuda()
elif args.model_type == 'mobilenetv2':
net = MobileNetV2(n_class=n_class).cuda()
else:
raise NotImplementedError
if args.ckpt_path is not None:
# the checkpoint can be a saved whole model object exported by amc_search.py, or a state_dict
print('=> Loading checkpoint {} ..'.format(args.ckpt_path))
ckpt = torch.load(args.ckpt_path)
if type(ckpt) == dict:
net.load_state_dict(ckpt['state_dict'])
else:
net = ckpt
net.to(args.device)
if torch.cuda.is_available() and args.n_gpu > 1:
net = torch.nn.DataParallel(net, list(range(args.n_gpu)))
return net
def run_training(config,
n_classes,
train_loader,
valid_loader,
width=1,
mb_version=1):
"""
Whole training procedure with fine-tune after regular training
"""
# defining model
if width > 1:
model = tvm.resnet18(num_classes=n_classes)
else:
if mb_version == 1:
model = MobileNet(n_classes=n_classes, width_mult=width)
else:
model = MobileNetV2(n_classes=n_classes, width_mult=width)
model = model.to(config['device'])
# print out number of parameters
num_params = 0
for p in model.parameters():
num_params += np.prod(p.size())
print(f"width={width}, num_params {num_params}")
# defining loss criterion, optimizer and learning rate scheduler
criterion = t.nn.CrossEntropyLoss()
opt = t.optim.Adam(model.parameters(), config['lr'])
sched = t.optim.lr_scheduler.MultiStepLR(opt, [3, 6])
# training process with Adam
tr_loss, tr_accuracy, valid_loss, valid_accuracy = train(
config, model, train_loader, valid_loader, criterion, opt, sched)
# training process with SGDR
opt = t.optim.SGD(model.parameters(), config['lr'] / 10, momentum=0.9)
sched = SGDR(opt, 3, 1.2)
tr_loss_finetune, tr_accuracy_finetune, valid_loss_finetune, valid_accuracy_finetune = train(
config, model, train_loader, valid_loader, criterion, opt, sched)
return [
tr_loss + tr_loss_finetune, tr_accuracy + tr_accuracy_finetune,
valid_loss + valid_loss_finetune,
valid_accuracy + valid_accuracy_finetune
]
