def get_classifier(mode, n_classes=10):
if mode == 'resnet18':
classifier = ResNet18(num_classes=n_classes)
elif mode == 'resnet34':
classifier = ResNet34(num_classes=n_classes)
elif mode == 'resnet50':
classifier = ResNet50(num_classes=n_classes)
elif mode == 'resnet18_imagenet':
classifier = resnet18(num_classes=n_classes)
elif mode == 'resnet50_imagenet':
classifier = resnet50(num_classes=n_classes)
else:
raise NotImplementedError()
return classifier
def get_classifier(mode, n_classes=10):
if mode == 'resnet18':
classifier = ResNet18(num_classes=n_classes)
elif mode == 'resnet34':
classifier = ResNet34(num_classes=n_classes)
elif mode == 'resnet50':
classifier = ResNet50(num_classes=n_classes)
elif mode == 'resnet18_imagenet':
classifier = resnet18(num_classes=n_classes)
elif mode == 'resnet50_imagenet':
classifier = resnet50(num_classes=n_classes)
elif mode == 'live':
classifier = FeatherNet(input_size=128, se=True, avgdown=True)
# classifier = ResNet18(num_classes=n_classes)
# classifier = LiveModel()
else:
raise NotImplementedError()
return classifier
.format(top1=top1, top5=top5, error1=100 - top1.avg), log)
return top1.avg, losses.avg
print("=> creating model '{}'".format(args.arch))
if args.dataset == 'imagenet':
if args.arch == 'resnet101':
net = resnet_imagenet.resnet101()
elif args.arch == 'resnet50':
net = resnet_imagenet.resnet50()
elif args.arch == 'resnet34':
net = resnet_imagenet.resnet34()
elif args.arch == 'resnet18':
net = resnet_imagenet.resnet18()
else:
if args.arch == 'resnet110':
net = models.resnet110(num_classes=10)
elif args.arch == 'resnet56':
net = models.resnet56(num_classes=10)
elif args.arch == 'resnet32':
net = models.resnet32(num_classes=10)
elif args.arch == 'resnet20':
net = models.resnet20(num_classes=10)
if args.dataset == 'imagenet':
if args.arch == 'resnet101':
state_dict = torch.load(
os.path.join(args.pretrain_path, 'resnet101-5d3b4d8f.pth'))
train_loader = torch.utils.data.DataLoader(train_data, batch_size=args.batch_size, shuffle=True,
num_workers=args.workers, pin_memory=True)
test_loader = torch.utils.data.DataLoader(test_data, batch_size=args.batch_size, shuffle=False,
num_workers=args.workers, pin_memory=True)
print("=> creating model '{}'".format(args.arch))
if args.dataset=='imagenet':
if args.arch=='resnet34':
net=resnet_imagenet.resnet34(args=args)
elif args.arch=='resnet18':
net=resnet_imagenet.resnet18(args=args)
else:
if args.arch=='dyresnet20':
net=resnet_cifar.resnet20(num_classes=num_classes,args=args)
elif args.arch=='dyresnet32':
net=resnet_cifar.resnet32(num_classes=num_classes,args=args)
elif args.arch=='dyresnet56':
net=resnet_cifar.resnet56(num_classes=num_classes,args=args)
if args.dataset=='imagenet':
if args.arch=='resnet34':
def load_model(opt):
if opt.from_modelzoo:
if opt.pretrained:
print("=> using pre-trained model '{}'".format(opt.arch))
model = models.__dict__[opt.model_def](pretrained=True)
else:
print("=> creating model '{}'".format(opt.arch))
model = models.__dict__[opt.model_def]()
return model
else:
if opt.pretrained_file != '':
model = torch.load(opt.pretrained_filedir)
else:
if opt.model_def == 'alexnet':
model = alexnet.Net()
if opt.cuda:
model = model.cuda()
elif opt.model_def == 'mobilenet':
model = mobilenet.Net(nClasses=opt.nclasses,
width_mult=opt.widthmult,
gtp=opt.grouptype,
gsz=opt.sp,
expsz=opt.exp)
if opt.cuda:
model = model.cuda()
elif opt.model_def == 'alexnetexpander':
model = alexnetexpander.Net()
if opt.cuda:
model = model.cuda()
elif opt.model_def == 'vgg16cifar':
model = vggcifar.vgg16()
if opt.cuda:
model = model.cuda()
elif opt.model_def == 'vgg16cifar_bn':
model = vggcifar.vgg16_bn()
if opt.cuda:
model = model.cuda()
elif opt.model_def == 'vgg16cifarexpander':
model = vggcifarexpander.vgg16()
if opt.cuda:
model = model.cuda()
elif opt.model_def == 'vgg16cifar_bnexpander':
model = vggcifarexpander.vgg16_bn()
if opt.cuda:
model = model.cuda()
elif opt.model_def == 'densenet_cifar':
model = densenet_cifar.DenseNet3(opt.layers,
opt.nclasses,
opt.growth,
reduction=opt.reduce,
bottleneck=opt.bottleneck,
dropRate=opt.droprate)
if opt.cuda:
model = model.cuda()
elif opt.model_def == 'densenetgrouped_cifar':
model = densenetgrouped_cifar.DenseNet3(
opt.layers,
opt.nclasses,
opt.growth,
reduction=opt.reduce,
bottleneck=opt.bottleneck,
dropRate=opt.droprate)
if opt.cuda:
model = model.cuda()
elif opt.model_def == 'densenetexpander_cifar':
model = densenetexpander_cifar.DenseNet3(
opt.layers,
opt.nclasses,
opt.growth,
reduction=opt.reduce,
bottleneck=opt.bottleneck,
dropRate=opt.droprate,
expandSize=opt.expandSize)
if opt.cuda:
model = model.cuda()
elif opt.model_def == 'densenet121':
model = densenet.densenet121()
if opt.cuda:
model = model.cuda()
elif opt.model_def == 'densenet169':
model = densenet.densenet169()
if opt.cuda:
model = model.cuda()
elif opt.model_def == 'densenet161':
model = densenet.densenet161()
if opt.cuda:
model = model.cuda()
elif opt.model_def == 'densenet201':
model = densenet.densenet201()
if opt.cuda:
model = model.cuda()
elif opt.model_def == 'densenetexpander121':
model = densenetexpander.densenet121(expandSize=opt.expandSize)
if opt.cuda:
model = model.cuda()
elif opt.model_def == 'densenetexpander169':
model = densenetexpander.densenet169(expandSize=opt.expandSize)
if opt.cuda:
model = model.cuda()
elif opt.model_def == 'densenetexpander161':
model = densenetexpander.densenet161(expandSize=opt.expandSize)
if opt.cuda:
model = model.cuda()
elif opt.model_def == 'densenetexpander201':
model = densenetexpander.densenet201(expandSize=opt.expandSize)
if opt.cuda:
model = model.cuda()
elif opt.model_def == 'resnet34':
model = resnet.resnet34()
if opt.cuda:
model = model.cuda()
elif opt.model_def == 'resnet50':
model = resnet.resnet50()
if opt.cuda:
model = model.cuda()
elif opt.model_def == 'resnet101':
model = resnet.resnet101()
if opt.cuda:
model = model.cuda()
elif opt.model_def == 'resnet152':
model = resnet.resnet152()
if opt.cuda:
model = model.cuda()
elif opt.model_def == 'resnetexpander34':
model = resnetexpander.resnet34(opt.expandSize)
if opt.cuda:
model = model.cuda()
elif opt.model_def == 'resnetexpander50':
model = resnetexpander.resnet50(opt.expandSize)
if opt.cuda:
model = model.cuda()
elif opt.model_def == 'resnetexpander101':
model = resnetexpander.resnet101(opt.expandSize)
if opt.cuda:
model = model.cuda()
elif opt.model_def == 'resnetexpander152':
model = resnetexpander.resnet152(opt.expandSize)
if opt.cuda:
model = model.cuda()
elif opt.model_def == 'resnet18':
model = resnet.resnet18()
if opt.cuda:
model = model.cuda()
elif opt.model_def == 'resnet50':
model = resnet.resnet50()
if opt.cuda:
model = model.cuda()
return model
num_workers=16,
pin_memory=True)
testloader = D.DataLoader(testset,
batch_size=args.batch_size,
shuffle=False,
num_workers=16,
pin_memory=True)
# global variable recording the best test accuracy
best_test_acc = 0.0
start_epoch = 0
if dataset_type == utils.Data.imagenet:
num_scales = 2
pretrained_agent = resnet_imagenet.resnet18(pretrained=True)
# agent = resnet_imagenet.ResNet(resnet_imagenet.BasicBlock, [2, 2, 2, 2], num_scales)
agent = resnet_imagenet.scalenet18(pretrained=True,
scalelist=[1, 0.5, 0.25])
# pretrained_dict = pretrained_agent.state_dict()
# model_dict = agent.state_dict()
# pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict and not k.startswith('fc')}
# model_dict.update(pretrained_dict)
# agent.load_state_dict(model_dict)
rnet = resnet_imagenet.resnet50(pretrained=True)
else:
# ScaleNet
num_scales = 2
def create_model(args, num_classes):
# determine the correct non-linearity
nl_str = args.nonlin.lower()
if nl_str == 'relu':
nonlin = nn.ReLU
elif nl_str == 'hardtanh':
nonlin = nn.Hardtanh
elif nl_str == 'hardsigmoid':
nonlin = hardsigmoid
assert False
elif nl_str == 'cabs':
nonlin = CAbs
elif nl_str == 'threshrelu':
nonlin = ThresholdReLU
elif nl_str == 'threshrelu2':
nonlin = partial(ThresholdReLU, max_val=2)
elif nl_str == 'step01':
nonlin = partial(Step,
targetprop_rule=args.tp_rule,
make01=True,
scale_by_grad_out=args.tp_grad_scale)
elif nl_str == 'step11':
nonlin = partial(Step,
targetprop_rule=args.tp_rule,
make01=False,
scale_by_grad_out=args.tp_grad_scale)
elif nl_str == 'staircase':
nonlin = partial(Staircase,
targetprop_rule=args.tp_rule,
nsteps=5,
margin=1,
trunc_thresh=2,
scale_by_grad_out=args.tp_grad_scale)
elif nl_str == 'staircase3':
nonlin = partial(Staircase,
targetprop_rule=args.tp_rule,
nsteps=3,
margin=1,
trunc_thresh=2,
scale_by_grad_out=args.tp_grad_scale)
elif nl_str == 'staircase11_3':
nonlin = partial(Staircase,
targetprop_rule=args.tp_rule,
nsteps=3,
margin=1,
trunc_thresh=2,
a=-1,
scale_by_grad_out=args.tp_grad_scale)
elif nl_str == 'staircase100':
nonlin = partial(Staircase,
targetprop_rule=args.tp_rule,
nsteps=100,
margin=1,
trunc_thresh=2,
scale_by_grad_out=args.tp_grad_scale)
elif nl_str == 'staircase_t2':
nonlin = partial(Staircase,
targetprop_rule=args.tp_rule,
nsteps=5,
margin=1,
trunc_thresh=3,
scale_by_grad_out=args.tp_grad_scale)
elif nl_str == 'staircase7':
nonlin = partial(Staircase,
targetprop_rule=args.tp_rule,
nsteps=7,
margin=1,
trunc_thresh=2,
scale_by_grad_out=args.tp_grad_scale)
elif nl_str == 'oldstaircase':
nonlin = partial(OldStaircase,
targetprop_rule=args.tp_rule,
scale_by_grad_out=args.tp_grad_scale)
elif nl_str == 'oldstaircase3':
nonlin = partial(OldStaircase,
targetprop_rule=args.tp_rule,
nsteps=3,
scale_by_grad_out=args.tp_grad_scale)
elif nl_str == 'oldstaircase11':
nonlin = partial(OldStaircase,
targetprop_rule=args.tp_rule,
a=-1,
scale_by_grad_out=args.tp_grad_scale)
elif nl_str == 'oldstaircase11_3':
nonlin = partial(OldStaircase,
targetprop_rule=args.tp_rule,
a=-1,
nsteps=3,
scale_by_grad_out=args.tp_grad_scale)
elif nl_str == 'oldstaircase7':
nonlin = partial(OldStaircase,
targetprop_rule=args.tp_rule,
nsteps=7,
scale_by_grad_out=args.tp_grad_scale)
else:
raise NotImplementedError(
'no other non-linearities currently supported')
if args.ds == 'mnist':
input_shape = (1, 28, 28)
elif args.ds.startswith('cifar'):
input_shape = (3, 32, 32)
elif args.ds == 'svhn':
input_shape = (3, 40, 40)
elif args.ds == 'imagenet':
input_shape = (3, 224, 224)
else:
raise NotImplementedError('no other datasets currently supported')
# create a model with the specified architecture
if args.arch == 'convnet4':
model = ConvNet4(nonlin=nonlin,
use_bn=args.use_bn,
input_shape=input_shape)
assert args.ds == 'mnist' or args.ds == 'cifar10'
elif args.arch == 'alexnet_drf':
assert args.use_bn
model = AlexNetDoReFa(nonlin=nonlin,
no_step_last=False,
use_bn=args.use_bn,
num_classes=num_classes,
data_parallel=len(args.gpus) > 1)
elif args.arch == 'convnet8':
model = ConvNet8(nonlin=nonlin,
use_bn=args.use_bn,
input_shape=input_shape,
no_step_last=False)
elif args.arch == 'resnet18':
assert args.use_bn, 'batchnorm is required for resnet'
if args.ds.startswith('cifar'):
model = ResNet18(nonlin=nonlin)
else:
model = resnet18(False, nonlin=nonlin)
else:
raise NotImplementedError('other models not yet supported')
logging.info("{} model has {} parameters and non-linearity={} ({})".format(
args.arch, sum([p.data.nelement() for p in model.parameters()]),
nl_str, args.tp_rule.name))
if len(args.gpus) > 1 and args.arch != 'alexnet_drf':
model = nn.DataParallel(model)
if args.cuda:
model.cuda()
return model