net = resnet20_cifar()
pretrain_path = '%s/%s-%s-pretrain.pth' % (save_root, model_name, dataset_name)
else:
raise NotImplementedError
net.load_state_dict(torch.load(pretrain_path), strict=False)
if use_cuda:
net.cuda()
################
# Load Dataset #
################
if dataset_name == 'ImageNet':
try:
train_loader = get_lmdb_imagenet('train', batch_size)
test_loader = get_lmdb_imagenet('test', 100)
except:
train_loader = get_dataloader(dataset_name, 'train', batch_size)
test_loader = get_dataloader(dataset_name, 'test', 100)
else:
train_loader = get_dataloader(dataset_name, 'train', batch_size)
test_loader = get_dataloader(dataset_name, 'test', 100)
####################
# Initial Recorder #
####################
SummaryPath = '%s/runs-FBS/Optimizer-%s-CR-%.3f' \
%(save_root, optimizer_type, final_CR)
if args.exp_spec is not '':
def __init__(self, task_name, task_type = 'prune', optimizer_type = 'adam',
save_root = None, SummaryPath = None, use_cuda = True, **kwargs):
self.task_name = task_name
self.task_type = task_type # prune, soft-quantize
self.model_name, self.dataset_name = task_name.split('-')
self.ratio = 'sample' if self.dataset_name in ['CIFARS'] else -1
#######
# Net #
#######
if task_type == 'prune':
if self.model_name == 'ResNet20':
if self.dataset_name in ['CIFAR10', 'CIFARS']:
self.net = resnet20_cifar()
elif self.dataset_name == 'STL10':
self.net = resnet20_stl()
else:
raise NotImplementedError
elif self.model_name == 'ResNet32':
if self.dataset_name in ['CIFAR10', 'CIFARS']:
self.net = resnet32_cifar()
elif self.dataset_name == 'STL10':
self.net = resnet32_stl()
else:
raise NotImplementedError
elif self.model_name == 'ResNet56':
if self.dataset_name in ['CIFAR10', 'CIFARS']:
self.net = resnet56_cifar()
elif self.dataset_name == 'CIFAR100':
self.net = resnet56_cifar(num_classes=100)
elif self.dataset_name == 'STL10':
self.net = resnet56_stl()
else:
raise NotImplementedError
elif self.model_name == 'ResNet18':
if self.dataset_name == 'ImageNet':
self.net = resnet18()
else:
raise NotImplementedError
elif self.model_name == 'vgg11':
self.net = vgg11() if self.dataset_name == 'CIFAR10' else vgg11_stl10()
else:
print(self.model_name, self.dataset_name)
raise NotImplementedError
elif task_type == 'soft-quantize':
if self.model_name == 'ResNet20':
if self.dataset_name in ['CIFAR10', 'CIFARS']:
self.net = soft_quantized_resnet20_cifar()
elif self.dataset_name in ['STL10']:
self.net = soft_quantized_resnet20_stl()
else:
raise NotImplementedError
else:
raise ('Task type not defined.')
self.meta_opt_flag = True # True for enabling meta leraning
##############
# Meta Prune #
##############
self.mask_dict = dict()
self.meta_grad_dict = dict()
self.meta_hidden_state_dict = dict()
######################
# Meta Soft Quantize #
######################
self.quantized = 0 # Quantized type
self.alpha_dict = dict()
self.alpha_hidden_dict = dict()
self.sq_rate = 0
self.s_rate = 0
self.q_rate = 0
##########
# Record #
##########
self.dataset_type = 'large' if self.dataset_name in ['ImageNet'] else 'small'
self.SummaryPath = SummaryPath
self.save_root = save_root
self.recorder = Recorder(self.SummaryPath, self.dataset_name, self.task_name)
####################
# Load Pre-trained #
####################
self.pretrain_path = '%s/%s-pretrain.pth' %(self.save_root, self.task_name)
self.net.load_state_dict(torch.load(self.pretrain_path))
print('Load pre-trained model from %s' %self.pretrain_path)
if use_cuda:
self.net.cuda()
# Optimizer for this task
if optimizer_type in ['Adam', 'adam']:
self.optimizer = Adam(self.net.parameters(), lr=1e-3)
else:
self.optimizer = SGD(self.net.parameters())
if self.dataset_name == 'ImageNet':
try:
self.train_loader = get_lmdb_imagenet('train', 128)
self.test_loader = get_lmdb_imagenet('test', 100)
except:
self.train_loader = get_dataloader(self.dataset_name, 'train', 128)
self.test_loader = get_dataloader(self.dataset_name, 'test', 100)
else:
self.train_loader = get_dataloader(self.dataset_name, 'train', 128, ratio=self.ratio)
self.test_loader = get_dataloader(self.dataset_name, 'test', 128)
self.iter_train_loader = yielder(self.train_loader)