def model_selector(model_type):
if model_type == 'vgg':
return vgg19(pretrained=False, num_classes=2)
elif model_type == 'resnet':
return resnet18(is_ptrtrained=False)
else:
raise ValueError('')
def load_pretrained_model(pretrained_path, model_type):
checkpoint = torch.load(add_prefix(pretrained_path, 'model_best.pth.tar'))
if model_type == 'vgg':
model = vgg19(pretrained=False, num_classes=2)
print('load vgg successfully.')
elif model_type == 'resnet':
model = resnet18(is_ptrtrained=False)
print('load resnet18 successfully.')
else:
raise ValueError('')
model.load_state_dict(remove_prefix(checkpoint['state_dict']))
return model
def load_pretrained_model(prefix, model_type):
if model_type == 'resnet':
model = resnet18(is_ptrtrained=False)
elif model_type == 'vgg':
model = vgg19(num_classes=2, pretrained=False)
else:
raise ValueError('')
checkpoint = torch.load(add_prefix(prefix, 'model_best.pth.tar'))
print('load pretrained model successfully.')
model.load_state_dict(remove_prefix(checkpoint['state_dict']))
print('best acc=%.4f' % checkpoint['best_accuracy'])
return model
def load_pretrained_model(prefix):
checkpoint = torch.load(add_prefix(prefix, 'model_best.pth.tar'))
model = vgg19(num_classes=2, pretrained=False)
print('load pretrained vgg19 successfully.')
model.load_state_dict(remove_prefix(checkpoint['state_dict']))
return model