def get_net_model(net='alexnet',
pretrained_dataset='imagenet',
dropout=False,
pretrained=True):
if net == 'alexnet':
model = myalexnet(pretrained=(pretrained_dataset == 'imagenet')
and pretrained,
dropout=dropout)
teacher_model = alexnet(pretrained=(pretrained_dataset == 'imagenet'))
elif net == 'mobilenet-imagenet':
model = MobileNet(num_classes=1001, dropout=dropout)
if pretrained and pretrained_dataset == 'imagenet':
model.load_state_dict(torch.load(imagenet_pretrained_mbnet_path))
teacher_model = MobileNet(num_classes=1001)
if os.path.isfile(imagenet_pretrained_mbnet_path):
teacher_model.load_state_dict(
torch.load(imagenet_pretrained_mbnet_path))
else:
warnings.warn('failed to import teacher model!')
elif net == 'erfnet-cityscapes':
model = erfnet(pretrained=(pretrained_dataset == 'cityscapes')
and pretrained,
num_classes=20,
dropout=dropout)
teacher_model = erfnet(pretrained=(pretrained_dataset == 'cityscapes'),
num_classes=20)
else:
raise NotImplementedError
for p in teacher_model.parameters():
p.requires_grad = False
teacher_model.eval()
return model, teacher_model
if args.preprocess == 'rescale':
model = EnergyEstimateNet(n_nodes=[len(Alexnet_width_ub) - 1, 1],
preprocessor=EnergyEstimateWidthRescale(
scales=(Alexnet_width_ub)))
else:
raise NotImplementedError
elif args.net == 'mobilenet':
if args.preprocess == 'rescale':
model = EnergyEstimateNet(n_nodes=[len(Mobilenet_width_ub) - 1, 1],
preprocessor=EnergyEstimateWidthRescale(
scales=(Mobilenet_width_ub)))
else:
raise NotImplementedError
elif args.net == 'erfnet':
if args.preprocess == 'rescale':
width_ub = [3] + erfnet().get_cpwub() + [20]
model = EnergyEstimateNet(
n_nodes=[len(width_ub) - 1, 1],
preprocessor=EnergyEstimateWidthRescale(scales=(width_ub)))
else:
raise NotImplementedError
else:
raise NotImplementedError
if not args.pinv:
model.cuda()
# args.wd /= tr_features.shape[0]
optimizer = torch.optim.Adam(model.regressor.parameters(),
lr=1e-3,
weight_decay=args.wd)
print(args.__dict__)
if args.net == 'alexnet':
net_class = CustomizedAlexnet
width_ub = [64, 192, 384, 256, 256, 4096, 4096]
h, w = 224, 224
elif args.net == 'mobilenet':
net_class = CustomizedMobileNet
width_ub = [
32, 64, 128, 128, 256, 256, 512, 512, 512, 512, 512, 512, 1024,
1024
]
h, w = 224, 224
elif args.net == 'erfnet':
net_class = customized_erfnet
width_ub = erfnet().get_cpwub()
h, w = 512, 1024
else:
raise NotImplementedError('not supported network architecture')
test_num = args.test_num
nclasses = args.num_classes
fake_img = torch.randn([1, 3, h, w], dtype=torch.float32)
random.seed(1)
use_cuda = not args.cpu
if use_cuda:
fake_img = fake_img.cuda(device=args.gpuid)
# test the upper-bound and lower-bound
model = net_class(width=[3] + [1] * len(width_ub) + [nclasses])
model.eval()
if use_cuda: