[Fix] Fix wrong default CLASSES settings in `init_model`.

pull/1037/head
mzr1996 2022-09-20 17:46:52 +08:00
parent e9e2d48cb2
commit b35ee5d139
1 changed files with 2 additions and 2 deletions

View File

@ -40,11 +40,11 @@ def init_model(config, checkpoint=None, device='cuda:0', options=None):
if 'CLASSES' in checkpoint.get('meta', {}):
model.CLASSES = checkpoint['meta']['CLASSES']
else:
from mmcls.datasets import ImageNet
from mmcls.datasets.categories import IMAGENET_CATEGORIES
warnings.simplefilter('once')
warnings.warn('Class names are not saved in the checkpoint\'s '
'meta data, use imagenet by default.')
model.CLASSES = ImageNet.CLASSES
model.CLASSES = IMAGENET_CATEGORIES
model.cfg = config # save the config in the model for convenience
model.to(device)
model.eval()