From ecabf988820d93e8e6b092af6051ec093d965a19 Mon Sep 17 00:00:00 2001 From: Ma Zerun Date: Tue, 27 Jul 2021 05:59:05 -0400 Subject: [PATCH] [Fix] Use local ATTENTION registery to avoid conflict with other repo. (#375) --- mmcls/models/builder.py | 3 +++ mmcls/models/utils/attention.py | 2 +- 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/mmcls/models/builder.py b/mmcls/models/builder.py index 7b61742df7a..c1caf0c7971 100644 --- a/mmcls/models/builder.py +++ b/mmcls/models/builder.py @@ -1,4 +1,5 @@ from mmcv.cnn import MODELS as MMCV_MODELS +from mmcv.cnn.bricks.registry import ATTENTION as MMCV_ATTENTION from mmcv.utils import Registry MODELS = Registry('models', parent=MMCV_MODELS) @@ -9,6 +10,8 @@ LOSSES = MODELS CLASSIFIERS = MODELS +ATTENTION = Registry('attention', parent=MMCV_ATTENTION) + def build_backbone(cfg): """Build backbone.""" diff --git a/mmcls/models/utils/attention.py b/mmcls/models/utils/attention.py index a4e918b8d12..f75a69c09c1 100644 --- a/mmcls/models/utils/attention.py +++ b/mmcls/models/utils/attention.py @@ -1,11 +1,11 @@ import torch import torch.nn as nn import torch.nn.functional as F -from mmcv.cnn.bricks.registry import ATTENTION from mmcv.cnn.bricks.transformer import build_dropout from mmcv.cnn.utils.weight_init import trunc_normal_ from mmcv.runner.base_module import BaseModule +from ..builder import ATTENTION from .helpers import to_2tuple