mirror of
https://github.com/huggingface/pytorch-image-models.git
synced 2025-06-03 15:01:08 +08:00
f902bcd54c
* select_conv2d -> create_conv2d * added create_attn to create attention module from string/bool/module * factor padding helpers into own file, use in both conv2d_same and avg_pool2d_same * add some more test eca resnet variants * minor tweaks, naming, comments, consistency
31 lines
893 B
Python
31 lines
893 B
Python
""" Select AttentionFactory Method
|
|
|
|
Hacked together by Ross Wightman
|
|
"""
|
|
import torch
|
|
from .se import SEModule
|
|
from .eca import EcaModule, CecaModule
|
|
|
|
|
|
def create_attn(attn_type, channels, **kwargs):
|
|
module_cls = None
|
|
if attn_type is not None:
|
|
if isinstance(attn_type, str):
|
|
attn_type = attn_type.lower()
|
|
if attn_type == 'se':
|
|
module_cls = SEModule
|
|
elif attn_type == 'eca':
|
|
module_cls = EcaModule
|
|
elif attn_type == 'eca':
|
|
module_cls = CecaModule
|
|
else:
|
|
assert False, "Invalid attn module (%s)" % attn_type
|
|
elif isinstance(attn_type, bool):
|
|
if attn_type:
|
|
module_cls = SEModule
|
|
else:
|
|
module_cls = attn_type
|
|
if module_cls is not None:
|
|
return module_cls(channels, **kwargs)
|
|
return None
|