[Bug] Imcompatible with TimmClassifier

[Weiyun1025]

Describe the bug

if the architecture or teacher of SingleTeacherDistill is a instance of TimmClassifier, the ckpt can not be load correctly.

To Reproduce

from mmengine.config import Config
from mmrazor.registry import MODELS
from mmengine.runner import load_checkpoint, save_checkpoint

model = MODELS.build(Config.fromfile('./debug_config.py')['model'])

save_checkpoint(model.state_dict(), './temp.pth')
load_checkpoint(model, './temp.pth', strict=True)

Config

student = dict(
    _scope_='mmcls',
    type='TimmClassifier',
    model_name='deit_tiny_patch16_224',
)

teacher = dict(
    _scope_='mmcls',
    type='TimmClassifier',
    model_name='cait_s24_224',
)

model = dict(
    _scope_='mmrazor',
    type='SingleTeacherDistill',
    architecture=student,
    teacher=teacher,
    distiller=dict(
        type='ConfigurableDistiller',
    ),
)

[wilxy]

Thank you for your issue.

This bug is casued by some problems of the TimmClassifier class in mmcls. I have submitted a PR to mmcls to address this bug. You can replace the original file in mmcls with the timm.py file in my PR to solve this bug.

Please let me know if there still exits any problems.

[Weiyun1025]

Thank you for your issue.

This bug is casued by some problems of the TimmClassifier class in mmcls. I have submitted a PR to mmcls to address this bug. You can replace the original file in mmcls with the timm.py file in my PR to solve this bug.

Please let me know if there still exits any problems.

Thank you for your response. That successfully solve my problem.